gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1163

Browse Source 
refactor: update notification handling in activities to include metadata in notifications, enhancing context for error reporting and success messages; update requirements to use version 1.3.0 of the sientia-dataops-library
This commit is contained in:
vitor-aignosi
2025-07-16 13:14:16 -03:00
parent 225964f938
commit b67d23da1f
12 changed files with 449 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
orchestrator/activities/formatters.py
View File

@@ -258,6 +258,8 @@ class Formatters(BaseActivity):
self.logger.info("Creating slot config...")
metadata = input_data.get("metadata", {})
current_slot_config = input_data['current_slot_config']
slot_config = input_data['slot_config']
to_delete = []
@@ -280,21 +282,23 @@ class Formatters(BaseActivity):
return output
def send_success_report(self, message: str, notification_id: str) -> None:
self.notification_handler.build_and_send_notification(
notification_id,
message,
"report_orchestration",
NotificationLevel.INFO
def send_success_report(self, metadata: dict[str, Any], message: str, notification_id: str) -> None:
self.send_notification(
metadata=metadata,
notification_id=notification_id,
message=message,
block="report_orchestration",
level=NotificationLevel.INFO
)
def send_error_report(self, message: str, notification_id: str,
def send_error_report(self, metadata: dict[str, Any], message: str, notification_id: str,
attachment: dict[str, Any]) -> None:
self.notification_handler.build_and_send_notification(
notification_id,
message,
"report_orchestration",
NotificationLevel.ERROR,
self.send_notification(
metadata=metadata,
notification_id=notification_id,
message=message,
block="report_orchestration",
level=NotificationLevel.ERROR,
attachment_content=json.dumps(attachment, indent=4, sort_keys=True)
)
@@ -332,6 +336,8 @@ class Formatters(BaseActivity):
self.logger.info("Reporting orchestration...")
metadata = input_data.get("metadata", {})
created_schedules = input_data['created_schedules']
updated_schedules = input_data['updated_schedules']
deleted_schedules = input_data['deleted_schedules']
@@ -343,15 +349,17 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
f"Created schedules: \n {', '.join(success_keys)}",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES"
metadata=metadata,
message=f"Created schedules: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES"
)
if len(error_keys) > 0:
self.send_error_report(
f"Failed to create schedules: \n {', '.join(error_keys)}",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES",
created_schedules
metadata=metadata,
message=f"Failed to create schedules: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES",
attachment=created_schedules
)
# Send report for updated schedules
@@ -361,15 +369,17 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
f"Updated schedules: \n {', '.join(success_keys)}",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
metadata=metadata,
message=f"Updated schedules: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
)
if len(error_keys) > 0:
self.send_error_report(
f"Failed to update schedules: \n {', '.join(error_keys)}",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
updated_schedules
metadata=metadata,
message=f"Failed to update schedules: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
attachment=updated_schedules
)
if len(deleted_schedules) > 0:
@@ -378,15 +388,17 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
f"Deleted schedules: \n {', '.join(success_keys)}",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES"
metadata=metadata,
message=f"Deleted schedules: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES"
)
if len(error_keys) > 0:
self.send_error_report(
f"Failed to delete schedules: \n {', '.join(error_keys)}",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES",
deleted_schedules
metadata=metadata,
message=f"Failed to delete schedules: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES",
attachment=deleted_schedules
)
@activity.defn(name="report_slot_orchestration")
@@ -404,6 +416,8 @@ class Formatters(BaseActivity):
self.logger.info("Reporting orchestration...")
metadata = input_data.get("metadata", {})
inserted_slots = input_data['inserted_slots']
deleted_slots = input_data['deleted_slots']
@@ -412,15 +426,17 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
f"Inserted slots: \n {', '.join(success_keys)}",
"REPORT_ORCHESTRATION_INSERTED_SLOTS"
metadata=metadata,
message=f"Inserted slots: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS"
)
if len(error_keys) > 0:
self.send_error_report(
f"Failed to insert slots: \n {', '.join(error_keys)}",
"REPORT_ORCHESTRATION_INSERTED_SLOTS",
inserted_slots
metadata=metadata,
message=f"Failed to insert slots: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS",
attachment=inserted_slots
)
if len(deleted_slots) > 0:
@@ -428,13 +444,15 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
f"Deleted slots: \n {', '.join(success_keys)}",
"REPORT_ORCHESTRATION_DELETED_SLOTS"
metadata=metadata,
message=f"Deleted slots: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS"
)
if len(error_keys) > 0:
self.send_error_report(
f"Failed to delete slots: \n {', '.join(error_keys)}",
"REPORT_ORCHESTRATION_DELETED_SLOTS",
deleted_slots
metadata=metadata,
message=f"Failed to delete slots: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS",
attachment=deleted_slots
)

64
orchestrator/activities/mongo_db.py
View File

@@ -92,6 +92,7 @@ class MongoDB(BaseActivity):
"""
query = input_data.get("query", {})
metadata = input_data.get("metadata", {})
collection_name = query.get("collection")
if not collection_name:
@@ -120,7 +121,8 @@ class MongoDB(BaseActivity):
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id="MONGODB_QUERY_ERROR",
message=f"Failed to execute MongoDB query: {e}",
block="load_query_from_mongodb",
@@ -146,6 +148,7 @@ class MongoDB(BaseActivity):
"""
query = input_data.get("query", {})
metadata = input_data.get("metadata", {})
collection_name = query.get("collection")
if not collection_name:
@@ -177,7 +180,8 @@ class MongoDB(BaseActivity):
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id="MONGODB_AGGREGATION_ERROR",
message=f"Failed to execute MongoDB aggregation: {e}",
block="aggregate_documents_in_mongodb",
@@ -196,6 +200,7 @@ class MongoDB(BaseActivity):
- updated_pipelines (list): List of updated pipelines.
"""
updated_pipelines = input_data.get("updated_pipelines", [])
metadata = input_data.get("metadata", {})
now = datetime.now().strftime(DEFAULT_DATE_FORMAT)
collection = self.database["orchestrated_schedules"]
@@ -205,10 +210,24 @@ class MongoDB(BaseActivity):
for pipeline in updated_pipelines if pipeline["success"]
]
data_filter = {"$or": argument} if argument else {}
collection.update_many(
data_filter,
{"$set": {"updated_at": now}}
)
try:
collection.update_many(
data_filter,
{"$set": {"updated_at": now}}
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_UPDATE_ERROR",
message=f"Failed to update pipelines timestamps: {e}",
block="update_pipelines_timestamps",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
@activity.defn(name="create_pipelines_timestamps")
async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
@@ -218,6 +237,7 @@ class MongoDB(BaseActivity):
- created_pipelines (list): List of created pipelines.
"""
created_pipelines = input_data.get("created_pipelines", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
now = datetime.now().strftime(DEFAULT_DATE_FORMAT)
@@ -229,7 +249,21 @@ class MongoDB(BaseActivity):
for pipeline in created_pipelines if pipeline["success"]
]
data_filter = argument if argument else {}
collection.insert_many(data_filter)
try:
collection.insert_many(data_filter)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_INSERT_ERROR",
message=f"Failed to create pipelines timestamps: {e}",
block="create_pipelines_timestamps",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
@activity.defn(name="delete_pipelines_timestamps")
async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
@@ -239,6 +273,7 @@ class MongoDB(BaseActivity):
- deleted_pipelines (list): List of deleted pipelines.
"""
deleted_pipelines = input_data.get("deleted_pipelines", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
argument = [
@@ -248,4 +283,17 @@ class MongoDB(BaseActivity):
]
data_filter = {"$or": argument} if argument else {}
collection.delete_many(data_filter)
try:
collection.delete_many(data_filter)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_DELETE_ERROR",
message=f"Failed to delete pipelines timestamps: {e}",
block="delete_pipelines_timestamps",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e

0
orchestrator/activities/notification.py
View File

80
orchestrator/activities/slot_manager.py
View File

@@ -2,10 +2,12 @@ from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
import traceback
import json
from logging import Logger
from sientia_do.temporal.activities.redis_base import Redis
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
class SlotManager(Redis):
@@ -18,7 +20,7 @@ class SlotManager(Redis):
password, logger, notification_handler)
@activity.defn(name="load_opc_slots")
async def load_opc_slots(self) -> dict[str, Any]:
async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Load all OPC slots from Redis
@@ -26,32 +28,45 @@ class SlotManager(Redis):
dict[str, Any]: A dictionary of OPC slots
"""
metadata = input_data.get("metadata", {})
self.logger.info("Loading OPC slots...")
opc_slots = {}
slot_keys = self.redis_client.keys("slot:opc_tags:*")
try:
self.logger.debug(f"Slot keys: {slot_keys}")
slot_keys = self.redis_client.keys("slot:opc_tags:*")
if slot_keys:
if isinstance(slot_keys[0], bytes):
decoded_keys = [key.decode('utf-8') for key in slot_keys]
else:
decoded_keys = slot_keys
self.logger.debug(f"Slot keys: {slot_keys}")
for key in decoded_keys:
opc_slots[key] = self.get(key)
if slot_keys:
if isinstance(slot_keys[0], bytes):
decoded_keys = [key.decode('utf-8') for key in slot_keys]
else:
decoded_keys = slot_keys
for key in decoded_keys:
opc_slots[key] = self.get(key)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="REDIS_GET_ERROR",
message=f"Failed to load OPC slots: {e}",
block="load_opc_slots",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
self.logger.info(f"Loaded {len(opc_slots)} OPC slots")
self.logger.debug(
f"Loaded: \n {json.dumps(opc_slots, indent=4, sort_keys=True)}")
return opc_slots
@activity.defn(name="load_active_ingestors")
async def load_active_ingestors(self) -> list[str]:
async def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
"""
Load all active ingestors from Redis
@@ -59,23 +74,40 @@ class SlotManager(Redis):
list[str]: A list of active ingestors
"""
metadata = input_data.get("metadata", {})
self.logger.info("Loading active ingestors...")
active_ingestors = self.redis_client.keys("heartbeat:ingestor:*")
try:
self.logger.info(f"Loaded {len(active_ingestors)} active ingestors")
active_ingestors = self.redis_client.keys("heartbeat:ingestor:*")
self.logger.debug(f"Active ingestors: \n {active_ingestors}")
self.logger.info(
f"Loaded {len(active_ingestors)} active ingestors")
ingestors = []
self.logger.debug(f"Active ingestors: \n {active_ingestors}")
for ingestor in active_ingestors:
if isinstance(ingestor, bytes):
ingestors.append(ingestor.decode('utf-8'))
else:
ingestors.append(ingestor)
ingestors = []
return ingestors
for ingestor in active_ingestors:
if isinstance(ingestor, bytes):
ingestors.append(ingestor.decode('utf-8'))
else:
ingestors.append(ingestor)
return ingestors
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="REDIS_GET_ERROR",
message=f"Failed to load active ingestors: {e}",
block="load_active_ingestors",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
@activity.defn(name="update_slots")
async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:

117
orchestrator/activities/temporal_manager.py
View File

@@ -4,10 +4,12 @@ from temporalio.client import (
from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.notifications.models import NotificationLevel
from google.protobuf.json_format import MessageToDict
import base64
from datetime import timedelta
@@ -62,7 +64,7 @@ class TemporalManager(BaseActivity):
}
@activity.defn(name="load_schedule")
async def load_schedule(self) -> dict[str, Any]:
async def load_schedule(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Load all orchestrated schedules from Temporal. Filters by search attribute
"Orchestrated" set to "true" and returns a dictionary of schedule_id:
@@ -72,50 +74,67 @@ class TemporalManager(BaseActivity):
dict[str, Any]: A dictionary of orchestrated schedules
"""
metadata = input_data.get("metadata", {})
self.logger.info("Getting orchestrated schedules...")
orchestrated_schedules = {}
for namespace, client in self.temporal_clients.items():
orchestrated_schedules[namespace] = {}
try:
self.logger.info(
f"Getting orchestrated schedules for {namespace}")
for namespace, client in self.temporal_clients.items():
orchestrated_schedules[namespace] = {}
async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]:
schedule_id = schedule.id
self.logger.info(
f"Getting orchestrated schedules for {namespace}")
self.logger.debug(f"Schedule id: {schedule_id}")
async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]:
schedule_id = schedule.id
handle = client.get_schedule_handle(
schedule_id)
self.logger.debug(f"Schedule id: {schedule_id}")
self.logger.debug("Handle acquired")
handle = client.get_schedule_handle(
schedule_id)
self.schedule_handles[namespace][schedule_id] = handle
self.logger.debug("Handle acquired")
self.logger.debug("Describing schedule...")
self.schedule_handles[namespace][schedule_id] = handle
desc = await handle.describe(
rpc_timeout=timedelta(seconds=60)
)
self.logger.debug("Describing schedule...")
self.logger.debug("Parsing args...")
desc = await handle.describe(
rpc_timeout=timedelta(seconds=60)
)
for arg in desc.schedule.action.args:
data = MessageToDict(arg)['data']
data = base64.b64decode(data).decode('utf-8')
self.logger.debug("Parsing args...")
frequency = desc.schedule.spec.intervals[0].every.seconds
for arg in desc.schedule.action.args:
data = MessageToDict(arg)['data']
data = base64.b64decode(data).decode('utf-8')
orchestrated_schedules[namespace][schedule_id] = {
'frequency': frequency,
'data': json.loads(data),
}
frequency = desc.schedule.spec.intervals[0].every.seconds
await sleep(0.1)
orchestrated_schedules[namespace][schedule_id] = {
'frequency': frequency,
'data': json.loads(data),
}
await sleep(0.1)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="TEMPORAL_LOAD_SCHEDULE_ERROR",
message=f"Failed to load orchestrated schedules: {e}",
block="load_schedule",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
self.logger.info(
f"Found {len(orchestrated_schedules[self.scouter_namespace]) + len(orchestrated_schedules[self.laborious_namespace])} orchestrated schedules")
@@ -137,27 +156,43 @@ class TemporalManager(BaseActivity):
"""
self.logger.info("Getting orchestrated schedules...")
metadata = input_data.get("metadata", {})
orchestrated_schedules = input_data.get('orchestrated_schedules', {})
for namespace, client in self.temporal_clients.items():
schedules = orchestrated_schedules.get(namespace, {})
try:
schedules = orchestrated_schedules.get(namespace, {})
self.logger.info(
f"Getting orchestrated schedules for {namespace}")
self.logger.info(
f"Getting orchestrated schedules for {namespace}")
async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]:
schedule_id = schedule.id
async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]:
schedule_id = schedule.id
if schedule_id not in schedules:
self.logger.info(
f"Schedule {schedule_id} not found in mongo db, cleaning up")
if schedule_id not in schedules:
self.logger.info(
f"Schedule {schedule_id} not found in mongo db, cleaning up")
handle = client.get_schedule_handle(
schedule_id)
handle = client.get_schedule_handle(
schedule_id)
await handle.delete()
await handle.delete()
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="TEMPORAL_NORMALIZE_SCHEDULES_ERROR",
message=f"Failed to normalize schedules: {e}",
block="normalize_schedules",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
raise e
@activity.defn(name="create_schedules")
async def create_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:

34
orchestrator/workflows/orchestrator.py
View File

@@ -14,9 +14,19 @@ class Orchestrator:
input_data['workflow_name'] = 'orchestrator'
metadata = {
'metadata': {
'schedule_name': input_data.get('schedule_name', 'orchestrator'),
'model_name': '-',
'model_id': '-',
'workflow_name': input_data['workflow_name']
}
}
pipeline_config_handler = workflow.start_local_activity_method(
Activities.aggregate_documents_in_mongodb,
{
**metadata,
'query': input_data['pipelines_query']
},
retry_policy=retry_policy,
@@ -26,6 +36,7 @@ class Orchestrator:
opc_servers_handler = workflow.start_local_activity_method(
Activities.find_documents_in_mongodb,
{
**metadata,
'query': input_data['opc_servers_query']
},
retry_policy=retry_policy,
@@ -35,6 +46,7 @@ class Orchestrator:
orchestrated_schedules_handler = workflow.start_local_activity_method(
Activities.find_documents_in_mongodb,
{
**metadata,
'query': {
'collection': 'orchestrated_schedules'
}
@@ -45,12 +57,18 @@ class Orchestrator:
current_slot_config_handler = workflow.start_local_activity_method(
Activities.load_opc_slots,
{
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
active_ingestors_handler = workflow.start_local_activity_method(
Activities.load_active_ingestors,
{
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
@@ -64,6 +82,7 @@ class Orchestrator:
formatted_orchestrated_schedules_handler = workflow.start_local_activity_method(
Activities.format_schedule_config,
{
**metadata,
'schedule_config': orchestrated_schedules
},
retry_policy=retry_policy,
@@ -73,6 +92,7 @@ class Orchestrator:
schedules_config_handler = workflow.start_local_activity_method(
Activities.process_schedules,
{
**metadata,
'pipelines': pipeline_config
},
retry_policy=retry_policy,
@@ -82,6 +102,7 @@ class Orchestrator:
slot_config_handler = workflow.start_local_activity_method(
Activities.process_slots,
{
**metadata,
'opc_servers': opc_servers,
'active_ingestors': active_ingestors,
'pipelines': pipeline_config,
@@ -97,6 +118,7 @@ class Orchestrator:
schedule_actions_handler = workflow.start_local_activity_method(
Activities.create_schedule_config,
{
**metadata,
'current_schedule_config': formatted_orchestrated_schedules,
'schedule_config': schedules_config
},
@@ -107,6 +129,7 @@ class Orchestrator:
slot_actions_handler = workflow.start_local_activity_method(
Activities.create_slot_config,
{
**metadata,
'current_slot_config': current_slot_config,
'slot_config': slot_config
},
@@ -117,6 +140,7 @@ class Orchestrator:
normalize_schedules_handler = workflow.start_local_activity_method(
Activities.normalize_schedules,
{
**metadata,
'orchestrated_schedules': formatted_orchestrated_schedules
},
retry_policy=retry_policy,
@@ -130,6 +154,7 @@ class Orchestrator:
slot_deletion_report_handler = workflow.start_activity_method(
Activities.delete_slots,
{
**metadata,
'to_delete': slot_actions['to_delete']
},
retry_policy=retry_policy,
@@ -139,6 +164,7 @@ class Orchestrator:
slot_insertion_report_handler = workflow.start_activity_method(
Activities.update_slots,
{
**metadata,
'to_insert': slot_actions['to_insert']
},
retry_policy=retry_policy,
@@ -148,6 +174,7 @@ class Orchestrator:
schedule_deletion_report_handler = workflow.start_activity_method(
Activities.delete_schedules,
{
**metadata,
'schedules': schedule_actions['to_delete']
},
retry_policy=retry_policy,
@@ -157,6 +184,7 @@ class Orchestrator:
schedule_insertion_report_handler = workflow.start_activity_method(
Activities.create_schedules,
{
**metadata,
'schedules': schedule_actions['to_create']
},
retry_policy=retry_policy,
@@ -166,6 +194,7 @@ class Orchestrator:
schedule_update_report_handler = workflow.start_activity_method(
Activities.update_schedules,
{
**metadata,
'schedules': schedule_actions['to_update']
},
retry_policy=retry_policy,
@@ -183,6 +212,7 @@ class Orchestrator:
schedule_report_handler = workflow.start_activity_method(
Activities.report_schedule_orchestration,
{
**metadata,
'created_schedules': schedule_insertion_report,
'updated_schedules': schedule_update_report,
'deleted_schedules': schedule_deletion_report
@@ -196,6 +226,7 @@ class Orchestrator:
slot_report_handler = workflow.start_activity_method(
Activities.report_slot_orchestration,
{
**metadata,
'inserted_slots': slot_insertion_report,
'deleted_slots': slot_deletion_report
},
@@ -208,6 +239,7 @@ class Orchestrator:
update_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.update_pipelines_timestamps,
{
**metadata,
'updated_pipelines': schedule_update_report
},
retry_policy=retry_policy,
@@ -219,6 +251,7 @@ class Orchestrator:
create_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.create_pipelines_timestamps,
{
**metadata,
'created_pipelines': schedule_insertion_report
},
retry_policy=retry_policy,
@@ -230,6 +263,7 @@ class Orchestrator:
delete_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.delete_pipelines_timestamps,
{
**metadata,
'deleted_pipelines': schedule_deletion_report
},
retry_policy=retry_policy,

2
requirements.txt
View File

@@ -4,4 +4,4 @@ sqlalchemy
redis
couchbase
pymongo
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.1
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.3.0

111
tests/orchestrator/activities/test_formatters.py
View File

@@ -8,13 +8,26 @@ from orchestrator.utils.orchestrator_functions import build_tag_config
@fixture
def formatters():
return Formatters(
formatters = Formatters(
scouter_namespace="scouter",
laborious_namespace="laborious",
logger=MagicMock(),
notification_handler=MagicMock()
)
formatters.send_notification = MagicMock()
return formatters
metadata = {
"metadata": {
"schedule_name": "test_schedule_name",
"workflow_name": "test_workflow_name",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
}
@mark.asyncio
@patch("orchestrator.activities.formatters.scouter",
@@ -374,23 +387,33 @@ async def test_create_slot_config(formatters):
def test_send_success_report(formatters):
formatters.send_success_report("test_message", "test_notification_id")
formatters.notification_handler.build_and_send_notification.assert_called_once_with(
"test_notification_id",
"test_message",
"report_orchestration",
NotificationLevel.INFO
formatters.send_success_report(
metadata=metadata,
message="test_message",
notification_id="test_notification_id"
)
formatters.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="test_notification_id",
message="test_message",
block="report_orchestration",
level=NotificationLevel.INFO
)
def test_send_error_report(formatters):
formatters.send_error_report(
"test_message", "test_notification_id", {"test": "test"})
formatters.notification_handler.build_and_send_notification.assert_called_once_with(
"test_notification_id",
"test_message",
"report_orchestration",
NotificationLevel.ERROR,
metadata=metadata,
message="test_message",
notification_id="test_notification_id",
attachment={"test": "test"}
)
formatters.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="test_notification_id",
message="test_message",
block="report_orchestration",
level=NotificationLevel.ERROR,
attachment_content=json.dumps(
{"test": "test"}, indent=4, sort_keys=True)
)
@@ -446,6 +469,7 @@ async def test_report_schedule_orchestration(formatters):
formatters.send_error_report = MagicMock()
input_data = {
**metadata,
"created_schedules": [
{
"namespace": "test_namespace",
@@ -496,33 +520,39 @@ async def test_report_schedule_orchestration(formatters):
])
formatters.send_success_report.assert_has_calls([
call(
"Created schedules: \n test_namespace/test_schedule_name_to_create",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES"
metadata=metadata['metadata'],
message="Created schedules: \n test_namespace/test_schedule_name_to_create",
notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES"
),
call(
"Updated schedules: \n test_namespace/test_schedule_name_to_update",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
metadata=metadata['metadata'],
message="Updated schedules: \n test_namespace/test_schedule_name_to_update",
notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
),
call(
"Deleted schedules: \n test_namespace/test_schedule_name_to_delete",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES"
metadata=metadata['metadata'],
message="Deleted schedules: \n test_namespace/test_schedule_name_to_delete",
notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES"
)
])
formatters.send_error_report.assert_has_calls([
call(
"Failed to create schedules: \n test_namespace/test_schedule_name_to_create_error: test_error",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES",
input_data['created_schedules']
metadata=metadata['metadata'],
message="Failed to create schedules: \n test_namespace/test_schedule_name_to_create_error: test_error",
notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES",
attachment=input_data['created_schedules']
),
call(
"Failed to update schedules: \n test_namespace/test_schedule_name_to_update_error: test_error",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
input_data['updated_schedules']
metadata=metadata['metadata'],
message="Failed to update schedules: \n test_namespace/test_schedule_name_to_update_error: test_error",
notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
attachment=input_data['updated_schedules']
),
call(
"Failed to delete schedules: \n test_namespace/test_schedule_name_to_delete_error: test_error",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES",
input_data['deleted_schedules']
metadata=metadata['metadata'],
message="Failed to delete schedules: \n test_namespace/test_schedule_name_to_delete_error: test_error",
notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES",
attachment=input_data['deleted_schedules']
)
])
@@ -536,6 +566,7 @@ async def test_report_slot_orchestration(formatters):
formatters.send_error_report = MagicMock()
input_data = {
**metadata,
"inserted_slots": {
"test_slot_name_to_create": {
"success": True
@@ -564,23 +595,27 @@ async def test_report_slot_orchestration(formatters):
])
formatters.send_success_report.assert_has_calls([
call(
"Inserted slots: \n test_slot_name_to_create",
"REPORT_ORCHESTRATION_INSERTED_SLOTS"
metadata=metadata['metadata'],
message="Inserted slots: \n test_slot_name_to_create",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS"
),
call(
"Deleted slots: \n test_slot_name_to_delete",
"REPORT_ORCHESTRATION_DELETED_SLOTS"
metadata=metadata['metadata'],
message="Deleted slots: \n test_slot_name_to_delete",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS"
)
])
formatters.send_error_report.assert_has_calls([
call(
"Failed to insert slots: \n test_slot_name_to_create_error",
"REPORT_ORCHESTRATION_INSERTED_SLOTS",
input_data['inserted_slots']
metadata=metadata['metadata'],
message="Failed to insert slots: \n test_slot_name_to_create_error",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS",
attachment=input_data['inserted_slots']
),
call(
"Failed to delete slots: \n test_slot_name_to_delete_error",
"REPORT_ORCHESTRATION_DELETED_SLOTS",
input_data['deleted_slots']
metadata=metadata['metadata'],
message="Failed to delete slots: \n test_slot_name_to_delete_error",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS",
attachment=input_data['deleted_slots']
)
])

28
tests/orchestrator/activities/test_mongo_db.py
View File

@@ -1,3 +1,4 @@
from curses import meta
from unittest.mock import MagicMock, patch, ANY
from pytest import fixture, mark
from orchestrator.activities.mongo_db import clear_mongo_id
@@ -50,7 +51,7 @@ def test_clear_mongo_id():
@fixture
@patch("orchestrator.activities.mongo_db.MongoClient")
def mongo_db(mongo_mock):
return (
mongo = (
MongoDB(
connection_string="mongodb://localhost:27017",
database_name="test_db",
@@ -58,6 +59,8 @@ def mongo_db(mongo_mock):
notification_handler=MagicMock()
)
)
mongo.send_notification = MagicMock()
return mongo
@patch("orchestrator.activities.mongo_db.MongoClient")
@@ -117,6 +120,15 @@ async def test_find_documents_in_mongodb_success(mongo_db):
{"name": {"$exists": True}}, {"_id": 0}
)
metadata = {
"metadata": {
"schedule_name": "test_schedule_name",
"workflow_name": "test_workflow_name",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
}
@mark.asyncio
async def test_find_documents_in_mongodb_failure(mongo_db):
@@ -130,11 +142,13 @@ async def test_find_documents_in_mongodb_failure(mongo_db):
try:
await mongo_db.find_documents_in_mongodb(
{
"query": input_data
"query": input_data,
**metadata
})
except Exception as e:
assert str(e) == "Error"
mongo_db.notification_handler.build_and_send_notification.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MONGODB_QUERY_ERROR",
message="Failed to execute MongoDB query: Error",
level=NotificationLevel.ERROR,
@@ -202,15 +216,17 @@ async def test_aggregate_documents_in_mongodb_failure(mongo_db):
try:
await mongo_db.aggregate_documents_in_mongodb(
{
"query": input_data
"query": input_data,
**metadata
})
except Exception as e:
assert str(e) == "Error"
mongo_db.notification_handler.build_and_send_notification.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MONGODB_AGGREGATION_ERROR",
message="Failed to execute MongoDB aggregation: Error",
level=NotificationLevel.ERROR,
block="aggregate_documents_in_mongodb",
level=NotificationLevel.ERROR,
attachment_content=ANY
)

17
tests/orchestrator/activities/test_slot_manager.py
View File

@@ -2,6 +2,15 @@ from unittest.mock import MagicMock, patch, call
from pytest import mark, fixture
from orchestrator.activities.slot_manager import SlotManager
metadata = {
"metadata": {
"schedule_name": "test_schedule_name",
"workflow_name": "test_workflow_name",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
}
@fixture
@patch("orchestrator.activities.slot_manager.Redis.__init__")
@@ -26,7 +35,7 @@ def slot_manager(_redis_mock):
@mark.asyncio
async def test_load_opc_slots_no_slot_keys(slot_manager):
slot_manager.redis_client.keys.return_value = []
assert await slot_manager.load_opc_slots() == {}
assert await slot_manager.load_opc_slots(metadata) == {}
@mark.asyncio
@@ -42,7 +51,7 @@ async def test_load_opc_slots(slot_manager):
]
)
response = await slot_manager.load_opc_slots()
response = await slot_manager.load_opc_slots(metadata)
assert response == {
"slot:opc_tags:1": "value1",
@@ -64,7 +73,7 @@ async def test_load_opc_slots_no_decode(slot_manager):
]
)
response = await slot_manager.load_opc_slots()
response = await slot_manager.load_opc_slots(metadata)
assert response == {
"slot:opc_tags:1": "value1",
@@ -78,7 +87,7 @@ async def test_load_active_ingestors(slot_manager):
slot_manager.redis_client.keys.return_value = [
b"heartbeat:ingestor:1", b"heartbeat:ingestor:2", "heartbeat:ingestor:3"]
response = await slot_manager.load_active_ingestors()
response = await slot_manager.load_active_ingestors(metadata)
assert response == ["heartbeat:ingestor:1",
"heartbeat:ingestor:2", "heartbeat:ingestor:3"]

11
tests/orchestrator/activities/test_temporal_manager.py
View File

@@ -7,6 +7,15 @@ import pytest_asyncio
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.utils.converters import parse_frequency
metadata = {
"metadata": {
"schedule_name": "test_schedule_name",
"workflow_name": "test_workflow_name",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
}
@fixture
@patch("orchestrator.activities.temporal_manager.Client.connect")
@@ -113,7 +122,7 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
temporal_manager.temporal_clients['laborious'].get_schedule_handle.return_value.describe \
.return_value.schedule.spec = describe_mock
response = await temporal_manager.load_schedule()
response = await temporal_manager.load_schedule(metadata)
temporal_manager.temporal_clients['scouter'].list_schedules.assert_awaited_once(
)

71
tests/orchestrator/workflows/test_orchestrator.py
View File

@@ -9,6 +9,16 @@ def orchestrator():
return Orchestrator()
metadata = {
'metadata': {
'schedule_name': 'test-schedule-name',
'workflow_name': 'orchestrator',
'model_name': '-',
'model_id': '-',
}
}
@mark.asyncio
@patch("orchestrator.workflows.orchestrator.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, orchestrator):
@@ -24,7 +34,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.aggregate_documents_in_mongodb,
{
"query": input_data["pipelines_query"]
"query": input_data["pipelines_query"],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -35,7 +46,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.find_documents_in_mongodb,
{
"query": input_data["opc_servers_query"]
"query": input_data["opc_servers_query"],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -48,7 +60,8 @@ async def test_run(workflow_mock, orchestrator):
{
"query": {
"collection": "orchestrated_schedules"
}
},
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -58,6 +71,9 @@ async def test_run(workflow_mock, orchestrator):
workflow_mock.start_local_activity_method.assert_has_calls([
call(
Activities.load_opc_slots,
{
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
@@ -66,6 +82,9 @@ async def test_run(workflow_mock, orchestrator):
workflow_mock.start_local_activity_method.assert_has_calls([
call(
Activities.load_active_ingestors,
{
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
@@ -75,7 +94,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.format_schedule_config,
{
'schedule_config': workflow_mock.start_local_activity_method.return_value
'schedule_config': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -86,7 +106,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.process_schedules,
{
'pipelines': workflow_mock.start_local_activity_method.return_value
'pipelines': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -100,6 +121,7 @@ async def test_run(workflow_mock, orchestrator):
'opc_servers': workflow_mock.start_local_activity_method.return_value,
'active_ingestors': workflow_mock.start_local_activity_method.return_value,
'pipelines': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -111,7 +133,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.create_schedule_config,
{
'current_schedule_config': workflow_mock.start_local_activity_method.return_value,
'schedule_config': workflow_mock.start_local_activity_method.return_value
'schedule_config': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -123,7 +146,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.create_slot_config,
{
'current_slot_config': workflow_mock.start_local_activity_method.return_value,
'slot_config': workflow_mock.start_local_activity_method.return_value
'slot_config': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -134,7 +158,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.normalize_schedules,
{
'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value
'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -146,7 +171,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.delete_slots,
{
'to_delete':
workflow_mock.start_local_activity_method.return_value['to_delete']
workflow_mock.start_local_activity_method.return_value['to_delete'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -158,7 +184,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.update_slots,
{
'to_insert':
workflow_mock.start_local_activity_method.return_value['to_insert']
workflow_mock.start_local_activity_method.return_value['to_insert'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -170,7 +197,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.delete_schedules,
{
'schedules':
workflow_mock.start_local_activity_method.return_value['to_delete']
workflow_mock.start_local_activity_method.return_value['to_delete'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -182,7 +210,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.create_schedules,
{
'schedules':
workflow_mock.start_local_activity_method.return_value['to_create']
workflow_mock.start_local_activity_method.return_value['to_create'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -194,7 +223,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.update_schedules,
{
'schedules':
workflow_mock.start_local_activity_method.return_value['to_update']
workflow_mock.start_local_activity_method.return_value['to_update'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -207,7 +237,8 @@ async def test_run(workflow_mock, orchestrator):
{
'created_schedules': workflow_mock.start_activity_method.return_value,
'updated_schedules': workflow_mock.start_activity_method.return_value,
'deleted_schedules': workflow_mock.start_activity_method.return_value
'deleted_schedules': workflow_mock.start_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -219,7 +250,8 @@ async def test_run(workflow_mock, orchestrator):
Activities.report_slot_orchestration,
{
'inserted_slots': workflow_mock.start_activity_method.return_value,
'deleted_slots': workflow_mock.start_activity_method.return_value
'deleted_slots': workflow_mock.start_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -230,7 +262,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.update_pipelines_timestamps,
{
'updated_pipelines': workflow_mock.start_activity_method.return_value
'updated_pipelines': workflow_mock.start_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -241,7 +274,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.delete_pipelines_timestamps,
{
'deleted_pipelines': workflow_mock.start_activity_method.return_value
'deleted_pipelines': workflow_mock.start_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -252,7 +286,8 @@ async def test_run(workflow_mock, orchestrator):
call(
Activities.create_pipelines_timestamps,
{
'created_pipelines': workflow_mock.start_activity_method.return_value
'created_pipelines': workflow_mock.start_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY