SIENTIAPDE-1150

fix: implement schedule normalization and timestamp management in MongoDB activities, enhance orchestration workflow with new formatting and normalization methods
This commit is contained in:
vitor-aignosi
2025-07-14 12:59:52 -03:00
parent e3a88eb1d5
commit 43206b578d
10 changed files with 648 additions and 135 deletions

View File

@@ -1,15 +1,19 @@
from sientia_do.notifications.models import NotificationLevel
from temporalio import activity, workflow from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import json import json
from typing import Any from typing import Any
from logging import Logger from logging import Logger
from datetime import datetime
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.notifications.models import NotificationLevel
from orchestrator.utils.orchestrator_functions import ( from orchestrator.utils.orchestrator_functions import (
scouter, predictions_batch, gather_read_tags, build_tag_config scouter, predictions_batch, gather_read_tags, build_tag_config
) )
from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT
from math import ceil from math import ceil
@@ -119,29 +123,48 @@ class Formatters(BaseActivity):
return slot_config return slot_config
def compare_configs(self, @activity.defn(name="format_schedule_config")
async def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Formats the schedule config to a dictionary with the schedule name as the key.
input_data:
- schedule_config (list[dict[str, Any]]): The schedule config to format.
Returns:
- dict[str, Any]: The formatted schedule config.
"""
schedule_config = input_data['schedule_config']
config = {}
for schedule in schedule_config:
namespace = schedule['namespace']
schedule_name = schedule['schedule_name']
updated_at = schedule['updated_at']
if namespace not in config:
config[namespace] = {}
config[namespace][schedule_name] = updated_at
return config
def compare_config_timestamps(self,
schedules: dict[str, Any], current_schedules: dict[str, Any], schedules: dict[str, Any], current_schedules: dict[str, Any],
to_update: dict[str, Any], to_create: dict[str, Any], to_update: dict[str, Any], to_create: dict[str, Any],
namespace: str): namespace: str):
""" """
Compares the schedules and current schedules, and updates the Compares the timestamps of the schedule and the current schedule.
to_update and to_create dictionaries.
""" """
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
if schedule_name in current_schedules: if schedule_name in current_schedules:
old_config = current_schedules[schedule_name]['data'] update_timestamp = schedule.get(
'updated_at', datetime.now().strftime(DEFAULT_DATE_FORMAT))
self.logger.debug(f"Comparing {schedule_name}:") old_timestamp = current_schedules[schedule_name]
self.logger.debug(json.dumps(
old_config, indent=4, sort_keys=True))
self.logger.debug(json.dumps(
schedule, indent=4, sort_keys=True))
if schedule != old_config: if update_timestamp > old_timestamp:
to_update[namespace][schedule_name] = schedule to_update[namespace][schedule_name] = schedule
else:
elif schedule_name not in current_schedules:
to_create[namespace][schedule_name] = schedule to_create[namespace][schedule_name] = schedule
@activity.defn(name="create_schedule_config") @activity.defn(name="create_schedule_config")
@@ -183,7 +206,7 @@ class Formatters(BaseActivity):
for namespace, schedules in schedule_config.items(): for namespace, schedules in schedule_config.items():
current_schedules = current_schedule_config.get(namespace, {}) current_schedules = current_schedule_config.get(namespace, {})
self.compare_configs( self.compare_config_timestamps(
schedules, current_schedules, to_update, to_create, namespace) schedules, current_schedules, to_update, to_create, namespace)
for namespace, schedules in current_schedule_config.items(): for namespace, schedules in current_schedule_config.items():
@@ -264,6 +287,15 @@ class Formatters(BaseActivity):
attachment_content=json.dumps(attachment, indent=4, sort_keys=True) attachment_content=json.dumps(attachment, indent=4, sort_keys=True)
) )
def parse_report_schedule(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]:
success_keys = [f"{value['namespace']}/{value['schedule_name']}"
for value in input_data if value['success']]
error_keys = [f"{value['namespace']}/{value['schedule_name']}: {value['message']}"
for value in input_data if not value['success']]
return success_keys, error_keys
def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]: def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]:
success_keys = [key for key, value success_keys = [key for key, value
in input_data.items() if value['success']] in input_data.items() if value['success']]
@@ -295,7 +327,8 @@ class Formatters(BaseActivity):
# Send report for created schedules # Send report for created schedules
if len(created_schedules) > 0: if len(created_schedules) > 0:
success_keys, error_keys = self.parse_report(created_schedules) success_keys, error_keys = self.parse_report_schedule(
created_schedules)
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(
@@ -312,7 +345,8 @@ class Formatters(BaseActivity):
# Send report for updated schedules # Send report for updated schedules
if len(updated_schedules) > 0: if len(updated_schedules) > 0:
success_keys, error_keys = self.parse_report(updated_schedules) success_keys, error_keys = self.parse_report_schedule(
updated_schedules)
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(
@@ -328,7 +362,8 @@ class Formatters(BaseActivity):
) )
if len(deleted_schedules) > 0: if len(deleted_schedules) > 0:
success_keys, error_keys = self.parse_report(deleted_schedules) success_keys, error_keys = self.parse_report_schedule(
deleted_schedules)
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(

View File

@@ -1,5 +1,8 @@
import datetime
from temporalio import workflow, activity from temporalio import workflow, activity
from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from typing import Any from typing import Any
import traceback import traceback
@@ -184,3 +187,56 @@ class MongoDB(BaseActivity):
self.logger.error(trace) self.logger.error(trace)
raise e raise e
@activity.defn(name="update_pipelines_timestamps")
async def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Update the timestamps of the pipelines in the MongoDB collection.
input_data:
- updated_pipelines (list): List of updated pipelines.
"""
updated_pipelines = input_data.get("updated_pipelines", [])
now = datetime.now().strftime(DEFAULT_DATE_FORMAT)
self.database["pipelines"].update_many(
{"$or": [
{"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"]}
for pipeline in updated_pipelines
]},
{"$set": {"updated_at": now}}
)
@activity.defn(name="create_pipelines_timestamps")
async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Create the timestamps of the pipelines in the MongoDB collection.
input_data:
- created_pipelines (list): List of created pipelines.
"""
created_pipelines = input_data.get("created_pipelines", [])
now = datetime.now().strftime(DEFAULT_DATE_FORMAT)
self.database["pipelines"].insert_many([{
"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"],
"updated_at": now
} for pipeline in created_pipelines])
@activity.defn(name="delete_pipelines_timestamps")
async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Delete the timestamps of the pipelines in the MongoDB collection.
input_data:
- deleted_pipelines (list): List of deleted pipelines.
"""
deleted_pipelines = input_data.get("deleted_pipelines", [])
self.database["pipelines"].delete_many({
"$or": [
{"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"]}
for pipeline in deleted_pipelines
]
})

View File

@@ -129,6 +129,37 @@ class TemporalManager(BaseActivity):
return orchestrated_schedules return orchestrated_schedules
@activity.defn(name="normalize_schedules")
async def normalize_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules".
input_data:
- orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare.
"""
self.logger.info("Getting orchestrated schedules...")
orchestrated_schedules = input_data.get('orchestrated_schedules', {})
for namespace, client in self.temporal_clients.items():
schedules = orchestrated_schedules.get(namespace, {})
self.logger.info(
"Getting orchestrated schedules for %s", 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
if schedule_id not in schedules:
self.logger.info(
"Schedule %s not found in mongo db, cleaning up", schedule_id)
handle = client.get_schedule_handle(
schedule_id)
await handle.delete()
@activity.defn(name="create_schedules") @activity.defn(name="create_schedules")
async def create_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: async def create_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
@@ -145,7 +176,7 @@ class TemporalManager(BaseActivity):
schedules_to_create = input_data['schedules'] schedules_to_create = input_data['schedules']
report = {} report = []
for namespace, schedules in schedules_to_create.items(): for namespace, schedules in schedules_to_create.items():
client = self.temporal_clients.get(namespace) client = self.temporal_clients.get(namespace)
@@ -198,17 +229,21 @@ class TemporalManager(BaseActivity):
search_attributes=search_attributes search_attributes=search_attributes
) )
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True, "success": True,
"message": "Schedule created successfully" "message": "Schedule created successfully"
} })
except Exception as e: except Exception as e:
self.logger.error("Failed to create schedule %s: %s", self.logger.error("Failed to create schedule %s: %s",
schedule_name, str(e)) schedule_name, str(e))
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False, "success": False,
"message": str(e) "message": str(e)
} })
self.logger.info(f"Processed {len(schedules_to_create)} schedules") self.logger.info(f"Processed {len(schedules_to_create)} schedules")
@@ -233,18 +268,19 @@ class TemporalManager(BaseActivity):
schedules_to_update = input_data['schedules'] schedules_to_update = input_data['schedules']
report = {} report = []
for namespace, schedules in schedules_to_update.items(): for namespace, schedules in schedules_to_update.items():
schedule_handles = self.schedule_handles.get(namespace, None) client = self.temporal_clients.get(namespace)
if schedule_handles is None: if not client:
raise ValueError( raise ValueError(
f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") f"Temporal client for {namespace} not found, clients: {self.temporal_clients}")
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
try: try:
handler = schedule_handles.get(schedule_name) handler = client.get_schedule_handle(
schedule_name)
if not handler: if not handler:
raise ValueError(f"Schedule {schedule_name} not found") raise ValueError(f"Schedule {schedule_name} not found")
@@ -276,17 +312,21 @@ class TemporalManager(BaseActivity):
del update_schedule del update_schedule
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True, "success": True,
"message": "Schedule updated successfully" "message": "Schedule updated successfully"
} })
except Exception as e: except Exception as e:
self.logger.error("Failed to update schedule %s: %s", self.logger.error("Failed to update schedule %s: %s",
schedule_name, str(e)) schedule_name, str(e))
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False, "success": False,
"message": str(e) "message": str(e)
} })
self.logger.info(f"Processed {len(schedules_to_update)} schedules") self.logger.info(f"Processed {len(schedules_to_update)} schedules")
@@ -311,18 +351,19 @@ class TemporalManager(BaseActivity):
schedules_to_delete = input_data['schedules'] schedules_to_delete = input_data['schedules']
report = {} report = []
for namespace, schedules in schedules_to_delete.items(): for namespace, schedules in schedules_to_delete.items():
schedule_handles = self.schedule_handles.get(namespace, None) client = self.temporal_clients.get(namespace)
if schedule_handles is None: if not client:
raise ValueError( raise ValueError(
f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") f"Temporal client for {namespace} not found, clients: {self.temporal_clients}")
for schedule_name in schedules: for schedule_name in schedules:
try: try:
handler = schedule_handles.get(schedule_name) handler = client.get_schedule_handle(
schedule_name)
if not handler: if not handler:
raise ValueError(f"Schedule {schedule_name} not found") raise ValueError(f"Schedule {schedule_name} not found")
@@ -331,17 +372,21 @@ class TemporalManager(BaseActivity):
del self.schedule_handles[namespace][schedule_name] del self.schedule_handles[namespace][schedule_name]
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True, "success": True,
"message": "Schedule deleted successfully" "message": "Schedule deleted successfully"
} })
except Exception as e: except Exception as e:
self.logger.error("Failed to delete schedule %s: %s", self.logger.error("Failed to delete schedule %s: %s",
schedule_name, str(e)) schedule_name, str(e))
report[schedule_name] = { report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False, "success": False,
"message": str(e) "message": str(e)
} })
self.logger.info(f"Processed {len(schedules_to_delete)} schedules") self.logger.info(f"Processed {len(schedules_to_delete)} schedules")

View File

@@ -0,0 +1 @@
DEFAULT_DATE_FORMAT = '%Y-%m-%d %H:%M:%S.%f'

View File

@@ -33,7 +33,12 @@ class Orchestrator:
) )
orchestrated_schedules_handler = workflow.execute_local_activity_method( orchestrated_schedules_handler = workflow.execute_local_activity_method(
Activities.load_schedule, Activities.find_documents_in_mongodb,
{
'query': {
'collection': 'orchestrated_schedules'
}
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=600) start_to_close_timeout=timedelta(seconds=600)
) )
@@ -56,6 +61,15 @@ class Orchestrator:
opc_servers = await opc_servers_handler opc_servers = await opc_servers_handler
active_ingestors = await active_ingestors_handler active_ingestors = await active_ingestors_handler
formatted_orchestrated_schedules_handler = workflow.execute_local_activity_method(
Activities.format_schedule_config,
{
'schedule_config': orchestrated_schedules
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
schedules_config_handler = workflow.execute_local_activity_method( schedules_config_handler = workflow.execute_local_activity_method(
Activities.process_schedules, Activities.process_schedules,
{ {
@@ -78,11 +92,12 @@ class Orchestrator:
schedules_config = await schedules_config_handler schedules_config = await schedules_config_handler
slot_config = await slot_config_handler slot_config = await slot_config_handler
formatted_orchestrated_schedules = await formatted_orchestrated_schedules_handler
schedule_actions_handler = workflow.execute_local_activity_method( schedule_actions_handler = workflow.execute_local_activity_method(
Activities.create_schedule_config, Activities.create_schedule_config,
{ {
'current_schedule_config': orchestrated_schedules, 'current_schedule_config': formatted_orchestrated_schedules,
'schedule_config': schedules_config 'schedule_config': schedules_config
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
@@ -99,8 +114,18 @@ class Orchestrator:
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60)
) )
normalize_schedules_handler = workflow.execute_local_activity_method(
Activities.normalize_schedules,
{
'orchestrated_schedules': formatted_orchestrated_schedules
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
schedule_actions = await schedule_actions_handler schedule_actions = await schedule_actions_handler
slot_actions = await slot_actions_handler slot_actions = await slot_actions_handler
await normalize_schedules_handler
slot_deletion_report_handler = workflow.execute_activity_method( slot_deletion_report_handler = workflow.execute_activity_method(
Activities.delete_slots, Activities.delete_slots,

View File

@@ -110,7 +110,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 14, "execution_count": 3,
"id": "7d01f160", "id": "7d01f160",
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
@@ -128,9 +128,6 @@
" namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')\n", " namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')\n",
")\n", ")\n",
"\n", "\n",
"manager = TemporalManager(temporal_client=temporal_client,\n",
" logger=logger,\n",
" notification_handler=MagicMock())\n",
" " " "
] ]
}, },
@@ -382,7 +379,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 17, "execution_count": 4,
"id": "f81b3728", "id": "f81b3728",
"metadata": {}, "metadata": {},
"outputs": [ "outputs": [
@@ -390,12 +387,112 @@
"name": "stdout", "name": "stdout",
"output_type": "stream", "output_type": "stream",
"text": [ "text": [
"{'_client': <temporalio.client.Client object at 0x770086544f90>, 'id': 'orchestrator'}\n" "{'id': 'orchestrator', 'schedule': ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='orchestrator'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=3600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), 'info': ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 9, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 9, 0, 0, 165083, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T09:00:00Z', first_execution_run_id='01980829-9700-771a-9c9d-2450eecc9af4')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 10, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 10, 0, 0, 125309, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T10:00:00Z', first_execution_run_id='01980860-8578-759d-97fa-dbdfc784f88a')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 11, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 11, 0, 0, 124356, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T11:00:00Z', first_execution_run_id='01980897-73f7-7e83-a7a6-20eaac38e7d4')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 12, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 12, 0, 0, 167329, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T12:00:00Z', first_execution_run_id='019808ce-62a2-75d6-b11d-c5387cfa258f')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 13, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 13, 0, 0, 167575, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T13:00:00Z', first_execution_run_id='01980905-5122-7807-90ae-3782d0332e79'))], next_action_times=[datetime.datetime(2025, 7, 14, 14, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 15, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 16, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 17, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 18, 0, tzinfo=datetime.timezone.utc)]), 'typed_search_attributes': TypedSearchAttributes(search_attributes=[]), 'search_attributes': {}, 'data_converter': DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x787d51390750>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x787d5277b590>), 'raw_entry': schedule_id: \"orchestrator\"\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 3600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"orchestrator\"\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1752483600\n",
" }\n",
" actual_time {\n",
" seconds: 1752483600\n",
" nanos: 165083498\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"orchestrator-2025-07-14T09:00:00Z\"\n",
" run_id: \"01980829-9700-771a-9c9d-2450eecc9af4\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1752487200\n",
" }\n",
" actual_time {\n",
" seconds: 1752487200\n",
" nanos: 125309119\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"orchestrator-2025-07-14T10:00:00Z\"\n",
" run_id: \"01980860-8578-759d-97fa-dbdfc784f88a\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1752490800\n",
" }\n",
" actual_time {\n",
" seconds: 1752490800\n",
" nanos: 124356906\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"orchestrator-2025-07-14T11:00:00Z\"\n",
" run_id: \"01980897-73f7-7e83-a7a6-20eaac38e7d4\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1752494400\n",
" }\n",
" actual_time {\n",
" seconds: 1752494400\n",
" nanos: 167329351\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"orchestrator-2025-07-14T12:00:00Z\"\n",
" run_id: \"019808ce-62a2-75d6-b11d-c5387cfa258f\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1752498000\n",
" }\n",
" actual_time {\n",
" seconds: 1752498000\n",
" nanos: 167575982\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"orchestrator-2025-07-14T13:00:00Z\"\n",
" run_id: \"01980905-5122-7807-90ae-3782d0332e79\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n",
" }\n",
" future_action_times {\n",
" seconds: 1752501600\n",
" }\n",
" future_action_times {\n",
" seconds: 1752505200\n",
" }\n",
" future_action_times {\n",
" seconds: 1752508800\n",
" }\n",
" future_action_times {\n",
" seconds: 1752512400\n",
" }\n",
" future_action_times {\n",
" seconds: 1752516000\n",
" }\n",
"}\n",
"}\n",
"{'_client': <temporalio.client.Client object at 0x787d53dd3490>, 'id': 'orchestrator'}\n"
] ]
} }
], ],
"source": [ "source": [
"async for schedule in await temporal_client.list_schedules():\n", "async for schedule in await temporal_client.list_schedules():\n",
" print(vars(schedule))\n",
" id = schedule.id\n", " id = schedule.id\n",
"\n", "\n",
" handle = temporal_client.get_schedule_handle(id)\n", " handle = temporal_client.get_schedule_handle(id)\n",

View File

@@ -247,32 +247,51 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
} }
@mark.asyncio
async def test_format_schedule_config(formatters):
input_data = {
"schedule_config": [
{"namespace": "test_namespace1", "schedule_name": "test1",
"updated_at": "2021-01-01"},
{"namespace": "test_namespace2", "schedule_name": "test2",
"updated_at": "2021-01-02"}
]
}
result = await formatters.format_schedule_config(input_data)
assert result == {
"test_namespace1": {
"test1": "2021-01-01"
},
"test_namespace2": {
"test2": "2021-01-02"
}
}
@mark.asyncio @mark.asyncio
async def test_create_schedule_config(formatters): async def test_create_schedule_config(formatters):
input_data = { input_data = {
"current_schedule_config": { "current_schedule_config": {
"scouter": { "scouter": {
"test_schedule_name_to_delete": { "test_schedule_name_to_delete": '2021-01-01',
"frequency": 60, "test_schedule_name_to_update": '2021-01-02'
"data": {"test": "test"}
},
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test"}
}
} }
}, },
"schedule_config": { "schedule_config": {
"laborious": { "laborious": {
"test_schedule_name_to_create": { "test_schedule_name_to_create": {
"frequency": 60, "frequency": 60,
"data": {"test": "test"} "data": {"test": "test"},
"updated_at": "2021-01-03"
} }
}, },
"scouter": { "scouter": {
"test_schedule_name_to_update": { "test_schedule_name_to_update": {
"frequency": 60, "frequency": 60,
"data": {"test": "test2"} "data": {"test": "test2"},
"updated_at": "2021-01-04"
} }
} }
} }
@@ -285,7 +304,8 @@ async def test_create_schedule_config(formatters):
"laborious": { "laborious": {
"test_schedule_name_to_create": { "test_schedule_name_to_create": {
"frequency": 60, "frequency": 60,
"data": {"test": "test"} "data": {"test": "test"},
"updated_at": "2021-01-03"
} }
}, },
"scouter": {} "scouter": {}
@@ -294,7 +314,8 @@ async def test_create_schedule_config(formatters):
"scouter": { "scouter": {
"test_schedule_name_to_update": { "test_schedule_name_to_update": {
"frequency": 60, "frequency": 60,
"data": {"test": "test2"} "data": {"test": "test2"},
"updated_at": "2021-01-04"
} }
}, },
"laborious": {} "laborious": {}
@@ -369,95 +390,129 @@ def test_send_error_report(formatters):
def test_parse_report(formatters): def test_parse_report(formatters):
input_data = { input_data = {
"test_schedule_name_to_create": { "test_key": {
"success": True "success": True
}, },
"test_schedule_name_to_create_error": { "test_key2": {
"success": False, "success": False,
"error": "test_error" "message": "test_error"
} }
} }
result = formatters.parse_report(input_data) result = formatters.parse_report(input_data)
assert result == ( assert result == (
["test_schedule_name_to_create"], ["test_key"],
["test_schedule_name_to_create_error"] ["test_key2"]
)
def test_parse_report_schedule(formatters):
input_data = [
{
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_create",
"success": True
},
{
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_create_error",
"success": False,
"message": "test_error"
}
]
result = formatters.parse_report_schedule(input_data)
assert result == (
["test_namespace/test_schedule_name_to_create"],
["test_namespace/test_schedule_name_to_create_error: test_error"]
) )
@mark.asyncio @mark.asyncio
async def test_report_schedule_orchestration(formatters): async def test_report_schedule_orchestration(formatters):
formatters.parse_report = MagicMock( formatters.parse_report_schedule = MagicMock(
side_effect=formatters.parse_report side_effect=formatters.parse_report_schedule
) )
formatters.send_success_report = MagicMock() formatters.send_success_report = MagicMock()
formatters.send_error_report = MagicMock() formatters.send_error_report = MagicMock()
input_data = { input_data = {
"created_schedules": { "created_schedules": [
"test_schedule_name_to_create": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_create",
"success": True "success": True
}, },
"test_schedule_name_to_create_error": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_create_error",
"success": False, "success": False,
"error": "test_error" "message": "test_error"
} }
}, ],
"updated_schedules": { "updated_schedules": [
"test_schedule_name_to_update": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_update",
"success": True "success": True
}, },
"test_schedule_name_to_update_error": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_update_error",
"success": False, "success": False,
"error": "test_error" "message": "test_error"
} }
}, ],
"deleted_schedules": { "deleted_schedules": [
"test_schedule_name_to_delete": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_delete",
"success": True "success": True
}, },
"test_schedule_name_to_delete_error": { {
"namespace": "test_namespace",
"schedule_name": "test_schedule_name_to_delete_error",
"success": False, "success": False,
"error": "test_error" "message": "test_error"
}
} }
]
} }
await formatters.report_schedule_orchestration(input_data) await formatters.report_schedule_orchestration(input_data)
formatters.parse_report.assert_has_calls([ formatters.parse_report_schedule.assert_has_calls([
call(input_data['created_schedules']), call(input_data['created_schedules']),
call(input_data['updated_schedules']), call(input_data['updated_schedules']),
call(input_data['deleted_schedules']) call(input_data['deleted_schedules'])
]) ])
formatters.send_success_report.assert_has_calls([ formatters.send_success_report.assert_has_calls([
call( call(
"Created schedules: \n test_schedule_name_to_create", "Created schedules: \n test_namespace/test_schedule_name_to_create",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES" "REPORT_ORCHESTRATION_CREATED_SCHEDULES"
), ),
call( call(
"Updated schedules: \n test_schedule_name_to_update", "Updated schedules: \n test_namespace/test_schedule_name_to_update",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES" "REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
), ),
call( call(
"Deleted schedules: \n test_schedule_name_to_delete", "Deleted schedules: \n test_namespace/test_schedule_name_to_delete",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES" "REPORT_ORCHESTRATION_DELETED_SCHEDULES"
) )
]) ])
formatters.send_error_report.assert_has_calls([ formatters.send_error_report.assert_has_calls([
call( call(
"Failed to create schedules: \n test_schedule_name_to_create_error", "Failed to create schedules: \n test_namespace/test_schedule_name_to_create_error: test_error",
"REPORT_ORCHESTRATION_CREATED_SCHEDULES", "REPORT_ORCHESTRATION_CREATED_SCHEDULES",
input_data['created_schedules'] input_data['created_schedules']
), ),
call( call(
"Failed to update schedules: \n test_schedule_name_to_update_error", "Failed to update schedules: \n test_namespace/test_schedule_name_to_update_error: test_error",
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES", "REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
input_data['updated_schedules'] input_data['updated_schedules']
), ),
call( call(
"Failed to delete schedules: \n test_schedule_name_to_delete_error", "Failed to delete schedules: \n test_namespace/test_schedule_name_to_delete_error: test_error",
"REPORT_ORCHESTRATION_DELETED_SCHEDULES", "REPORT_ORCHESTRATION_DELETED_SCHEDULES",
input_data['deleted_schedules'] input_data['deleted_schedules']
) )

View File

@@ -244,3 +244,58 @@ async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
else: else:
assert False, "Expected a ValueError to be raised" assert False, "Expected a ValueError to be raised"
@mark.asyncio
@patch("orchestrator.activities.mongo_db.datetime")
async def test_update_pipelines_timestamps_success(datetime_mock, mongo_db):
input_data = {"updated_pipelines": [
{"schedule_name": "test1", "namespace": "test1"},
{"schedule_name": "test2", "namespace": "test2"}
]}
mongo_db.database["pipelines"].update_many.return_value = MagicMock()
await mongo_db.update_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].update_many.assert_called_once_with(
{"$or": [
{"schedule_name": "test1", "namespace": "test1"},
{"schedule_name": "test2", "namespace": "test2"}
]},
{"$set": {
"updated_at": datetime_mock.now.return_value.strftime.return_value}}
)
@mark.asyncio
@patch("orchestrator.activities.mongo_db.datetime")
async def test_create_pipelines_timestamps_success(datetime_mock, mongo_db):
input_data = {"created_pipelines": [
{"schedule_name": "test1", "namespace": "test1"},
{"schedule_name": "test2", "namespace": "test2"}
]}
mongo_db.database["pipelines"].insert_many.return_value = MagicMock()
await mongo_db.create_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].insert_many.assert_called_once_with(
[
{"schedule_name": "test1", "namespace": "test1",
"updated_at": datetime_mock.now.return_value.strftime.return_value},
{"schedule_name": "test2", "namespace": "test2",
"updated_at": datetime_mock.now.return_value.strftime.return_value}
]
)
@mark.asyncio
@patch("orchestrator.activities.mongo_db.datetime")
async def test_delete_pipelines_timestamps_success(datetime_mock, mongo_db):
input_data = {"deleted_pipelines": [
{"schedule_name": "test1", "namespace": "test1"},
{"schedule_name": "test2", "namespace": "test2"}
]}
mongo_db.database["pipelines"].delete_many.return_value = MagicMock()
await mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].delete_many.assert_called_once_with(
{"$or": [
{"schedule_name": "test1", "namespace": "test1"},
{"schedule_name": "test2", "namespace": "test2"}
]}
)

View File

@@ -41,12 +41,7 @@ async def test_connect_to_temporal(connect_mock, temporal_manager):
]) ])
@mark.asyncio async def async_iter():
@patch("orchestrator.activities.temporal_manager.MessageToDict",
return_value={"data": base64.b64encode(json.dumps({"test": "test"}).encode('utf-8'))})
async def test_load_schedule(_mock_message_to_dict, temporal_manager):
# Create async iterator mock
async def async_iter():
yield MagicMock( yield MagicMock(
id="test-schedule-id", id="test-schedule-id",
search_attributes={ search_attributes={
@@ -66,6 +61,13 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
} }
) )
@mark.asyncio
@patch("orchestrator.activities.temporal_manager.MessageToDict",
return_value={"data": base64.b64encode(json.dumps({"test": "test"}).encode('utf-8'))})
async def test_load_schedule(_mock_message_to_dict, temporal_manager):
# Create async iterator mock
handle = MagicMock( handle = MagicMock(
describe=AsyncMock( describe=AsyncMock(
return_value=MagicMock( return_value=MagicMock(
@@ -134,6 +136,43 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
} }
@mark.asyncio
async def test_normalize_schedules(temporal_manager):
input_data = {
"orchestrated_schedules": {
"scouter": {"test-scouter": "2021-01-01"},
"laborious": {"test-schedule-id1": "2021-01-01"}
}
}
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
return_value=async_iter()
)
temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock(
return_value=async_iter()
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
return_value=MagicMock(
delete=AsyncMock()
)
)
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
return_value=MagicMock(
delete=AsyncMock()
)
)
await temporal_manager.normalize_schedules(input_data)
temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([
call("test-schedule-id"),
])
temporal_manager.temporal_clients['scouter'].get_schedule_handle.return_value.delete.assert_awaited_once(
)
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.parse_frequency", @patch("orchestrator.activities.temporal_manager.parse_frequency",
side_effect=parse_frequency) side_effect=parse_frequency)
@@ -308,20 +347,26 @@ async def test_create_schedule(
) )
]) ])
assert report == { assert report == [
"test-schedule": { {
"schedule_name": "test-schedule",
"namespace": "scouter",
"success": True, "success": True,
"message": "Schedule created successfully" "message": "Schedule created successfully"
}, },
"test-schedule-invalid-frequency": { {
"schedule_name": "test-schedule-invalid-frequency",
"namespace": "scouter",
"success": False, "success": False,
"message": "Invalid frequency" "message": "Invalid frequency"
}, },
"test-schedule-laborious": { {
"schedule_name": "test-schedule-laborious",
"namespace": "laborious",
"success": True, "success": True,
"message": "Schedule created successfully" "message": "Schedule created successfully"
} }
} ]
@mark.asyncio @mark.asyncio
@@ -397,29 +442,72 @@ async def test_update_schedules(
} }
} }
handler_scouter = MagicMock(
update=AsyncMock(
update=AsyncMock(
side_effect=lambda f: f(input_mock)
)
)
)
handler_laborious = MagicMock(
update=AsyncMock(
update=AsyncMock(
side_effect=lambda f: f(input_mock)
)
)
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
side_effect=[
handler_scouter,
None
]
)
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
side_effect=[
handler_laborious
]
)
report = await temporal_manager.update_schedules(input_data) report = await temporal_manager.update_schedules(input_data)
temporal_manager.schedule_handles['scouter']['test-schedule'].update.assert_called_once() temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([
temporal_manager.schedule_handles['laborious']['test-schedule-laborious'].update.assert_called_once() call("test-schedule"),
call("test-schedule_no_handler")
])
temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls([
call("test-schedule-laborious")
])
assert report == { handler_scouter.update.assert_called_once()
"test-schedule": { handler_laborious.update.assert_called_once()
assert report == [
{
"schedule_name": "test-schedule",
"namespace": "scouter",
"success": True, "success": True,
"message": "Schedule updated successfully" "message": "Schedule updated successfully"
}, },
"test-schedule_no_handler": { {
"schedule_name": "test-schedule_no_handler",
"namespace": "scouter",
"success": False, "success": False,
"message": "Schedule test-schedule_no_handler not found" "message": "Schedule test-schedule_no_handler not found"
}, },
"test-schedule-laborious": { {
"schedule_name": "test-schedule-laborious",
"namespace": "laborious",
"success": True, "success": True,
"message": "Schedule updated successfully" "message": "Schedule updated successfully"
} }
} ]
@mark.asyncio @mark.asyncio
async def test_update_schedules_with_no_handle(temporal_manager): async def test_update_schedules_with_no_client(temporal_manager):
temporal_manager.temporal_clients = {} temporal_manager.temporal_clients = {}
input_data = { input_data = {
"schedules": { "schedules": {
@@ -436,7 +524,7 @@ async def test_update_schedules_with_no_handle(temporal_manager):
await temporal_manager.update_schedules(input_data) await temporal_manager.update_schedules(input_data)
except Exception as e: except Exception as e:
assert str( assert str(
e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}"
@mark.asyncio @mark.asyncio
@@ -464,28 +552,57 @@ async def test_delete_schedules(temporal_manager):
] ]
} }
} }
handler_scouter = MagicMock(
delete=AsyncMock()
)
handler_laborious = MagicMock(
delete=AsyncMock()
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
side_effect=[
handler_scouter,
None
]
)
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
side_effect=[
handler_laborious
]
)
report = await temporal_manager.delete_schedules(input_data) report = await temporal_manager.delete_schedules(input_data)
assert report == { handler_scouter.delete.assert_called_once()
"test-schedule": { handler_laborious.delete.assert_called_once()
assert report == [
{
"schedule_name": "test-schedule",
"namespace": "scouter",
"success": True, "success": True,
"message": "Schedule deleted successfully" "message": "Schedule deleted successfully"
}, },
"test-schedule_no_handler": { {
"schedule_name": "test-schedule_no_handler",
"namespace": "scouter",
"success": False, "success": False,
"message": "Schedule test-schedule_no_handler not found" "message": "Schedule test-schedule_no_handler not found"
}, },
"test-schedule-laborious": { {
"schedule_name": "test-schedule-laborious",
"namespace": "laborious",
"success": True, "success": True,
"message": "Schedule deleted successfully" "message": "Schedule deleted successfully"
} }
} ]
@mark.asyncio @mark.asyncio
async def test_delete_schedules_with_no_handle(temporal_manager): async def test_delete_schedules_with_no_client(temporal_manager):
temporal_manager.schedule_handles = {} temporal_manager.temporal_clients = {}
input_data = { input_data = {
"schedules": { "schedules": {
"abc": [ "abc": [
@@ -498,4 +615,4 @@ async def test_delete_schedules_with_no_handle(temporal_manager):
await temporal_manager.delete_schedules(input_data) await temporal_manager.delete_schedules(input_data)
except Exception as e: except Exception as e:
assert str( assert str(
e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}"

View File

@@ -44,7 +44,12 @@ async def test_run(workflow_mock, orchestrator):
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call( call(
Activities.load_schedule, Activities.find_documents_in_mongodb,
{
"query": {
"collection": "orchestrated_schedules"
}
},
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY
) )
@@ -66,6 +71,17 @@ async def test_run(workflow_mock, orchestrator):
) )
]) ])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_schedule_config,
{
'schedule_config': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call( call(
Activities.process_schedules, Activities.process_schedules,
@@ -114,6 +130,17 @@ async def test_run(workflow_mock, orchestrator):
) )
]) ])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.normalize_schedules,
{
'orchestrated_schedules': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls([
call( call(
Activities.delete_slots, Activities.delete_slots,