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

SIENTIAPDE-1231

Browse Source 
Refactor orchestrator activities and configuration files

- Removed unused Temporal timeout environment variables from values.yaml.
- Reorganized import statements in various activity files for better readability.
- Updated logging messages to use consistent formatting across activities.
- Enhanced test cases to ensure proper initialization and shutdown of orchestrator activities.
- Improved overall code structure and readability by applying consistent formatting and style adjustments.
This commit is contained in:
vitor-aignosi
2025-10-15 17:20:52 -03:00
parent 310cae4c2f
commit f3a6af5603
39 changed files with 4988 additions and 3826 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
orchestrator/activities/activities.py
View File

@@ -1,21 +1,22 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.formatters import Formatters
from orchestrator.activities.email import Email
from orchestrator.activities.mongo_db import MongoDB
from typing import Any
from logging import Logger
from sientia_do.temporal.activities.postgres import Postgres
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.temporal.activities.postgres import Postgres
from orchestrator.activities.email import Email
from orchestrator.activities.formatters import Formatters
from orchestrator.activities.mongo_db import MongoDB
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.temporal_manager import TemporalManager
class Activities( # Couchbase,
TemporalManager, SlotManager, Formatters, MongoDB, Email,
Postgres):
TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres
):
"""
Central activities orchestrator for Temporal workflow operations.
@@ -34,16 +35,17 @@ class Activities( # Couchbase,
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self,
temporal_config: dict[str, Any],
# couchbase_config: dict[str, Any],
redis_config: dict[str, Any],
mongodb_config: dict[str, Any],
email_config: dict[str, Any],
postgres_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
def __init__(
self,
temporal_config: dict[str, Any],
# couchbase_config: dict[str, Any],
redis_config: dict[str, Any],
mongodb_config: dict[str, Any],
email_config: dict[str, Any],
postgres_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
):
# Initialize parent classes
# Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
# username=couchbase_config['username'],
@@ -51,55 +53,64 @@ class Activities( # Couchbase,
# logger=logger,
# notification_handler=notification_handler)
TemporalManager.__init__(self,
host=temporal_config['temporal_host'],
scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'],
task_timeout=temporal_config['temporal_task_timeout_minutes'],
run_timeout=temporal_config['temporal_run_timeout_minutes'],
execution_timeout=temporal_config['temporal_execution_timeout_minutes'],
logger=logger,
notification_handler=notification_handler)
TemporalManager.__init__(
self,
host=temporal_config['temporal_host'],
scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger,
notification_handler=notification_handler,
)
SlotManager.__init__(self,
host=redis_config['host'],
port=redis_config['port'],
username=redis_config['username'],
password=redis_config['password'],
logger=logger,
notification_handler=notification_handler)
SlotManager.__init__(
self,
host=redis_config['host'],
port=redis_config['port'],
username=redis_config['username'],
password=redis_config['password'],
logger=logger,
notification_handler=notification_handler,
)
Formatters.__init__(self,
scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger,
notification_handler=notification_handler)
Formatters.__init__(
self,
scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger,
notification_handler=notification_handler,
)
MongoDB.__init__(self,
connection_string=mongodb_config['connection_string'],
database_name=mongodb_config['database_name'],
ttl_index_seconds=mongodb_config['ttl_index_seconds'],
logger=logger,
notification_handler=notification_handler)
MongoDB.__init__(
self,
connection_string=mongodb_config['connection_string'],
database_name=mongodb_config['database_name'],
ttl_index_seconds=mongodb_config['ttl_index_seconds'],
logger=logger,
notification_handler=notification_handler,
)
Email.__init__(self,
sender_email=email_config['sender_email'],
sender_password=email_config['sender_password'],
smtp_server=email_config['smtp_server'],
smtp_port=email_config['smtp_port'],
logger=logger,
notification_handler=notification_handler)
Email.__init__(
self,
sender_email=email_config['sender_email'],
sender_password=email_config['sender_password'],
smtp_server=email_config['smtp_server'],
smtp_port=email_config['smtp_port'],
logger=logger,
notification_handler=notification_handler,
)
Postgres.__init__(self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler)
Postgres.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler,
)
def shutdown(self):
"""

53
orchestrator/activities/couchbase.py
View File

@@ -1,11 +1,12 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from logging import Logger
from datetime import timedelta
import json
import traceback
from datetime import timedelta
from logging import Logger
from typing import Any
from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster
from couchbase.options import ClusterOptions
@@ -34,44 +35,43 @@ class Couchbase(BaseActivity):
but maintained for potential future use.
"""
def __init__(self, connection_string: str, username: str,
password: str, logger: Logger,
notification_handler: NotificationHandler):
def __init__(
self,
connection_string: str,
username: str,
password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.connection_string = connection_string
self.username = username
self.password = password
logger.info("Initializing Couchbase connection...")
logger.info('Initializing Couchbase connection...')
self.cluster = Cluster(
connection_string,
ClusterOptions(
authenticator=PasswordAuthenticator(
username=username,
password=password
)
)
authenticator=PasswordAuthenticator(username=username, password=password)
),
)
logger.info("Awaiting Couchbase connection...")
logger.info('Awaiting Couchbase connection...')
self.cluster.wait_until_ready(timeout=timedelta(seconds=10))
logger.info("Couchbase connection ready")
logger.info('Couchbase connection ready')
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
def shutdown(self):
try:
self.cluster.close()
except Exception as e:
self.logger.error(f"Failed to close Couchbase connection: {e}")
self.logger.error(f'Failed to close Couchbase connection: {e}')
def __del__(self):
self.shutdown()
@activity.defn(name="load_query_from_couchbase")
@activity.defn(name='load_query_from_couchbase')
async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Load a query from couchbase
@@ -84,16 +84,16 @@ class Couchbase(BaseActivity):
"""
query = input_data['query']
self.logger.info(f"Executing couchbase query: {query}")
self.logger.info(f'Executing couchbase query: {query}')
try:
result = self.cluster.query(query)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="COUCHBASE_LOAD_QUERY_ERROR",
message=f"Failed to execute couchbase query: {e}",
block="load_query_from_couchbase",
notification_id='COUCHBASE_LOAD_QUERY_ERROR',
message=f'Failed to execute couchbase query: {e}',
block='load_query_from_couchbase',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
@@ -106,8 +106,7 @@ class Couchbase(BaseActivity):
for row in result.rows():
rows.append(row)
self.logger.info("Fetched %d rows from couchbase", len(rows))
self.logger.debug("Rows: \n %s",
json.dumps(rows, indent=4, sort_keys=True))
self.logger.info('Fetched %d rows from couchbase', len(rows))
self.logger.debug('Rows: \n %s', json.dumps(rows, indent=4, sort_keys=True))
return rows

117
orchestrator/activities/email.py
View File

@@ -1,21 +1,23 @@
from smtplib import SMTPServerDisconnected
from temporalio import workflow, activity
from temporalio import activity, workflow
from orchestrator import metrics
with workflow.unsafe.imports_passed_through():
import traceback
import smtplib
from typing import Any
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.observability.logger import Logger
from sientia_do.notifications.handlers import NotificationHandler
from orchestrator.utils.email_builder import EmailBuilder
import traceback
from email import encoders
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
from typing import Any
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from orchestrator.utils.email_builder import EmailBuilder
class Email(BaseActivity):
@@ -35,10 +37,15 @@ class Email(BaseActivity):
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self, sender_email: str, sender_password: str,
smtp_server: str, smtp_port: int,
logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
sender_email: str,
sender_password: str,
smtp_server: str,
smtp_port: int,
logger: Logger,
notification_handler: NotificationHandler,
):
self.email_builder = EmailBuilder(logger=logger)
self.sender_email = sender_email
@@ -46,7 +53,7 @@ class Email(BaseActivity):
self.smtp_port = smtp_port
self.smtp_server = smtp_server
logger.info(f"Initializing Email with {smtp_server}:{smtp_port}")
logger.info(f'Initializing Email with {smtp_server}:{smtp_port}')
if smtp_server is not None:
self.server = smtplib.SMTP(smtp_server, smtp_port, timeout=20)
@@ -55,9 +62,7 @@ class Email(BaseActivity):
self.server.starttls()
self.server.login(self.sender_email, self.sender_password)
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
def shutdown(self):
"""
@@ -65,7 +70,7 @@ class Email(BaseActivity):
"""
self.server.quit()
@activity.defn(name="build_email_html")
@activity.defn(name='build_email_html')
async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Build HTML email content for configured receiver groups.
@@ -91,18 +96,14 @@ class Email(BaseActivity):
receiver_groups = input_data['receiver_groups']
mail_type = input_data['mail_type']
self.info(f"Building email html for {mail_type} mail type.",
metadata=metadata)
self.info(f'Building email html for {mail_type} mail type.', metadata=metadata)
for group_name, group_config in receiver_groups.items():
html = self.email_builder.build_email(
group_config['notifications'], mail_type)
for _group_name, group_config in receiver_groups.items():
html = self.email_builder.build_email(group_config['notifications'], mail_type)
group_config['html'] = html
self.info(f"Email html built for {mail_type} mail type.",
metadata=metadata)
self.info(f'Email html built for {mail_type} mail type.', metadata=metadata)
return receiver_groups
@@ -124,17 +125,12 @@ class Email(BaseActivity):
try:
# Create the attachment as a MIMEBase object
part = MIMEBase('application', 'octet-stream')
part.set_payload(
attachment['attachment_content'].encode('utf-8'))
part.set_payload(attachment['attachment_content'].encode('utf-8'))
encoders.encode_base64(part)
part.add_header(
'Content-Disposition',
f'attachment; filename="{att_name}"'
)
part.add_header('Content-Disposition', f'attachment; filename="{att_name}"')
msg.attach(part)
except Exception as e:
self.logger.error(
f"Failed to attach content of {att_name}: {e}")
self.logger.error(f'Failed to attach content of {att_name}: {e}')
raise e
@@ -152,30 +148,27 @@ class Email(BaseActivity):
Exception: If email sending fails after reconnection attempts.
"""
try:
self.server.sendmail(
self.sender_email, receivers, msg.as_string())
self.server.sendmail(self.sender_email, receivers, msg.as_string())
except SMTPServerDisconnected as e:
self.logger.error(f"SMTP server disconnected: {e}")
self.logger.info(
f"Reconnecting to {self.smtp_server}:{self.smtp_port}")
self.logger.error(f'SMTP server disconnected: {e}')
self.logger.info(f'Reconnecting to {self.smtp_server}:{self.smtp_port}')
if self.server:
try:
self.server.quit()
except SMTPServerDisconnected as e:
self.logger.info(f"Server already disconnected: {e}")
self.logger.info(f'Server already disconnected: {e}')
except Exception as e:
self.logger.error(f"Failed to quit server: {e}")
self.logger.error(f'Failed to quit server: {e}')
raise e
self.server = smtplib.SMTP(
self.smtp_server, self.smtp_port, timeout=20)
self.server = smtplib.SMTP(self.smtp_server, self.smtp_port, timeout=20)
if self.sender_password:
self.server.starttls()
self.server.login(self.sender_email, self.sender_password)
self.server.sendmail(self.sender_email, receivers, msg.as_string())
@activity.defn(name="send_email")
@activity.defn(name='send_email')
async def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Send email notifications to configured receiver groups.
@@ -202,40 +195,38 @@ class Email(BaseActivity):
mail_type = input_data['mail_type']
if self.smtp_server is None:
self.info(f"Skipping email sending for {mail_type} mail type.",
metadata=metadata)
self.info(f'Skipping email sending for {mail_type} mail type.', metadata=metadata)
return {}
self.info(f"Sending email for {mail_type} mail type.",
metadata=metadata)
self.info(f'Sending email for {mail_type} mail type.', metadata=metadata)
for group_name, group_config in receiver_groups.items():
try:
receivers = ", ".join(group_config['members'])
receivers = ', '.join(group_config['members'])
self.info(f"Sending email to {group_name}: {receivers}",
metadata=metadata)
self.info(f'Sending email to {group_name}: {receivers}', metadata=metadata)
msg = MIMEMultipart()
msg.attach(MIMEText(group_config['html'], 'html'))
msg['From'] = self.sender_email
msg['To'] = receivers
msg['Subject'] = f"SIENTIA™ {mail_type}"
msg['Subject'] = f'SIENTIA™ {mail_type}'
msg = self.handle_attachments(
[
{
"filename": f"{notification['trigger']}_{notification['notification_id']}.txt",
"attachment_content": notification['attachment_content']
'filename': f'{notification["trigger"]}_{notification["notification_id"]}.txt',
'attachment_content': notification['attachment_content'],
}
for notification in group_config['notifications']
if notification.get('attachment_content') is not None],
msg)
if notification.get('attachment_content') is not None
],
msg,
)
self.try_send_email(msg, receivers)
except Exception as e:
self.error(f"Failed to send email to {group_name}: {e}",
metadata=metadata)
self.error(f'Failed to send email to {group_name}: {e}', metadata=metadata)
traceback.print_exc()
group_config['status'] = 'failed'
else:
@@ -244,13 +235,11 @@ class Email(BaseActivity):
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
email_group=group_name
email_group=group_name,
).inc()
self.info(f"Email sent to {group_name}: {receivers}",
metadata=metadata)
self.info(f'Email sent to {group_name}: {receivers}', metadata=metadata)
self.info(f"Email sent for {mail_type} mail type.",
metadata=metadata)
self.info(f'Email sent for {mail_type} mail type.', metadata=metadata)
return receiver_groups

354
orchestrator/activities/formatters.py
View File

@@ -2,19 +2,25 @@ from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import json
from typing import Any
from logging import Logger
from math import ceil
from typing import Any
from pandas import DataFrame
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 orchestrator.utils.orchestrator_functions import (
scouter, predictions_batch, gather_read_tags, build_tag_config, minimal_retrain
)
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from math import ceil
topic_separator = "\n ========== \n"
from orchestrator.utils.orchestrator_functions import (
build_tag_config,
gather_read_tags,
minimal_retrain,
predictions_batch,
scouter,
)
topic_separator = '\n ========== \n'
class Formatters(BaseActivity):
@@ -41,16 +47,18 @@ class Formatters(BaseActivity):
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self,
scouter_namespace: str,
laborious_namespace: str,
logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
scouter_namespace: str,
laborious_namespace: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace
BaseActivity.__init__(self, logger=logger,
notification_handler=notification_handler)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
@activity.defn(name="process_schedules")
@activity.defn(name='process_schedules')
async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Process pipeline configurations into Temporal-compatible schedule configurations.
@@ -61,7 +69,7 @@ class Formatters(BaseActivity):
Pipeline Types Supported:
- scouter: Data collection workflows with OPC tag configurations
- predictions_batch: ML prediction workflows with OPC write configurations
- predictions_batch: ML prediction workflows with OPC write configurations
- minimal_retrain: Model retraining workflows with SQL query configurations
Args:
@@ -73,46 +81,43 @@ class Formatters(BaseActivity):
- dict[str, Any]: The schedule config dictionary
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Processing schedules...", metadata=metadata)
self.info('Processing schedules...', metadata=metadata)
pipelines = input_data['pipelines']
schedule_config = {
self.scouter_namespace: {},
self.laborious_namespace: {}
}
schedule_config = {self.scouter_namespace: {}, self.laborious_namespace: {}}
for pipeline in pipelines:
if pipeline['workflow_type'] == 'scouter':
schedule_config[self.scouter_namespace][pipeline['schedule_name']] = {
**scouter(pipeline),
"updated_at": pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ))
'updated_at': pipeline.get(
'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
}
elif pipeline['workflow_type'] == 'predictions_batch':
schedule_config[self.laborious_namespace][pipeline['schedule_name']
] = {
schedule_config[self.laborious_namespace][pipeline['schedule_name']] = {
**predictions_batch(pipeline),
"updated_at": pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ))
'updated_at': pipeline.get(
'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
}
elif pipeline['workflow_type'] == 'minimal_retrain':
schedule_config[self.laborious_namespace][pipeline['schedule_name']
] = {
schedule_config[self.laborious_namespace][pipeline['schedule_name']] = {
**minimal_retrain(pipeline),
"updated_at": pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ))
'updated_at': pipeline.get(
'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
}
self.info("Processed schedules", metadata=metadata)
self.debug(json.dumps(
schedule_config, indent=4, sort_keys=True), metadata=metadata)
self.info('Processed schedules', metadata=metadata)
self.debug(json.dumps(schedule_config, indent=4, sort_keys=True), metadata=metadata)
return schedule_config
@activity.defn(name="process_slots")
@activity.defn(name='process_slots')
async def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Extracts all read tags from input pipelines, divides them into slots and
@@ -130,9 +135,9 @@ class Formatters(BaseActivity):
- dict[str, Any]: The slot config dictionary
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Processing slots...", metadata=metadata)
self.info('Processing slots...', metadata=metadata)
pipelines = input_data['pipelines']
opc_servers_list = input_data['opc_servers']
@@ -154,43 +159,42 @@ class Formatters(BaseActivity):
last_index = 0
for i in range(1, number_of_slots):
slot_config[f"{i}"] = {}
for tag in tags[last_index:last_index + tags_per_slot]:
slot_config[f'{i}'] = {}
for tag in tags[last_index : last_index + tags_per_slot]:
try:
slot_config = build_tag_config(
tag, slot_config.copy(), opc_servers, i)
slot_config = build_tag_config(tag, slot_config.copy(), opc_servers, i)
except ValueError as e:
self.send_notification(
metadata=metadata,
notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR",
notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR',
message=str(e),
block="orchestrator",
level=NotificationLevel.ERROR
block='orchestrator',
level=NotificationLevel.ERROR,
)
last_index += tags_per_slot
slot_config[f"{number_of_slots}"] = {}
slot_config[f'{number_of_slots}'] = {}
for tag in tags[last_index:]:
try:
slot_config = build_tag_config(
tag, slot_config.copy(), opc_servers, number_of_slots)
tag, slot_config.copy(), opc_servers, number_of_slots
)
except ValueError as e:
self.send_notification(
metadata=metadata,
notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR",
notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR',
message=str(e),
block="orchestrator",
level=NotificationLevel.ERROR
block='orchestrator',
level=NotificationLevel.ERROR,
)
self.info("Processed slots", metadata=metadata)
self.debug(json.dumps(
slot_config, indent=4, sort_keys=True), metadata=metadata)
self.info('Processed slots', metadata=metadata)
self.debug(json.dumps(slot_config, indent=4, sort_keys=True), metadata=metadata)
return slot_config
@activity.defn(name="format_schedule_config")
@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.
@@ -202,9 +206,9 @@ class Formatters(BaseActivity):
Returns:
dict[str, Any]: The formatted schedule config.
"""
metadata = input_data["metadata"]
metadata = input_data['metadata']
self.info("Formatting schedule config...", metadata=metadata)
self.info('Formatting schedule config...', metadata=metadata)
schedule_config = input_data['schedule_config']
@@ -219,16 +223,20 @@ class Formatters(BaseActivity):
config[namespace][schedule_name] = updated_at
self.info("Formatted schedule config", metadata=metadata)
self.debug(json.dumps(
config, indent=4, sort_keys=True), metadata=metadata)
self.info('Formatted schedule config', metadata=metadata)
self.debug(json.dumps(config, indent=4, sort_keys=True), metadata=metadata)
return config
def compare_config_timestamps(self,
schedules: dict[str, Any], current_schedules: dict[str, Any],
to_update: dict[str, Any], to_create: dict[str, Any],
namespace: str, metadata: dict[str, Any]):
def compare_config_timestamps(
self,
schedules: dict[str, Any],
current_schedules: dict[str, Any],
to_update: dict[str, Any],
to_create: dict[str, Any],
namespace: str,
metadata: dict[str, Any],
):
"""
Compares the timestamps of the schedule and the current schedule to determine
which schedules need to be updated or created.
@@ -243,23 +251,22 @@ class Formatters(BaseActivity):
"""
for schedule_name, schedule in schedules.items():
if schedule_name in current_schedules:
update_timestamp = schedule.get(
'updated_at', now())
update_timestamp = schedule.get('updated_at', now())
old_timestamp = current_schedules[schedule_name]
self.debug(
f"Comparing schedule {schedule_name}:{update_timestamp} vs {old_timestamp}", metadata=metadata)
f'Comparing schedule {schedule_name}:{update_timestamp} vs {old_timestamp}',
metadata=metadata,
)
if update_timestamp > old_timestamp:
to_update[namespace][schedule_name] = schedule
else:
to_create[namespace][schedule_name] = schedule
@activity.defn(name="create_schedule_config")
async def create_schedule_config(self,
input_data: dict[str, Any]) -> dict[str, Any]:
@activity.defn(name='create_schedule_config')
async def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Creates a schedule config dictionary based on the input data.
Checks the existing schedule config and updates it with the new schedule config,
@@ -276,51 +283,37 @@ class Formatters(BaseActivity):
- dict[str, Any]: The schedule config dictionary
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Creating schedule config...", metadata=metadata)
self.info('Creating schedule config...', metadata=metadata)
current_schedule_config = input_data['current_schedule_config']
schedule_config = input_data['schedule_config']
to_update = {
self.scouter_namespace: {},
self.laborious_namespace: {}
}
to_create = {
self.scouter_namespace: {},
self.laborious_namespace: {}
}
to_delete = {
self.scouter_namespace: [],
self.laborious_namespace: []
}
to_update = {self.scouter_namespace: {}, self.laborious_namespace: {}}
to_create = {self.scouter_namespace: {}, self.laborious_namespace: {}}
to_delete = {self.scouter_namespace: [], self.laborious_namespace: []}
for namespace, schedules in schedule_config.items():
current_schedules = current_schedule_config.get(namespace, {})
self.compare_config_timestamps(
schedules, current_schedules, to_update, to_create, namespace, metadata)
schedules, current_schedules, to_update, to_create, namespace, metadata
)
for namespace, schedules in current_schedule_config.items():
for schedule_name in schedules:
if schedule_name not in schedule_config[namespace]:
to_delete[namespace].append(schedule_name)
output = {
"to_update": to_update,
"to_create": to_create,
"to_delete": to_delete
}
output = {'to_update': to_update, 'to_create': to_create, 'to_delete': to_delete}
self.info("Created schedule config", metadata=metadata)
self.debug(json.dumps(
output, indent=4, sort_keys=True), metadata=metadata)
self.info('Created schedule config', metadata=metadata)
self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
return output
@activity.defn(name="create_slot_config")
async def create_slot_config(self,
input_data: dict[str, Any]) -> dict[str, Any]:
@activity.defn(name='create_slot_config')
async def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Creates a slot config dictionary based on the input data.
Checks the existing slot config and updates it with the new slot config,
@@ -337,9 +330,9 @@ class Formatters(BaseActivity):
- dict[str, Any]: The slot config dictionary
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Creating slot config...", metadata=metadata)
self.info('Creating slot config...', metadata=metadata)
current_slot_config = input_data['current_slot_config']
slot_config = input_data['slot_config']
@@ -349,21 +342,18 @@ class Formatters(BaseActivity):
number_of_slots = len(slot_config)
if number_of_current_slots > number_of_slots:
to_delete = [str(i) for i in range(
number_of_slots + 1, number_of_current_slots + 1)]
to_delete = [str(i) for i in range(number_of_slots + 1, number_of_current_slots + 1)]
output = {
"to_delete": to_delete,
"to_insert": slot_config
}
output = {'to_delete': to_delete, 'to_insert': slot_config}
self.info("Created slot config", metadata=metadata)
self.debug(json.dumps(
output, indent=4, sort_keys=True), metadata=metadata)
self.info('Created slot config', metadata=metadata)
self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
return output
def send_success_report(self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str = None) -> None:
def send_success_report(
self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str = None
) -> None:
"""
Sends a success notification report.
@@ -377,13 +367,14 @@ class Formatters(BaseActivity):
metadata=metadata,
notification_id=notification_id,
message=message,
block="report_orchestration",
block='report_orchestration',
level=NotificationLevel.INFO,
attachment_content=json.dumps(attachment, indent=4, sort_keys=True)
attachment_content=json.dumps(attachment, indent=4, sort_keys=True),
)
def send_error_report(self, metadata: dict[str, Any], message: str, notification_id: str,
attachment: str) -> None:
def send_error_report(
self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str
) -> None:
"""
Sends an error notification report.
@@ -397,9 +388,9 @@ class Formatters(BaseActivity):
metadata=metadata,
notification_id=notification_id,
message=message,
block="report_orchestration",
block='report_orchestration',
level=NotificationLevel.ERROR,
attachment_content=attachment
attachment_content=attachment,
)
def parse_report_schedule(self, input_data: dict[str, Any]) -> tuple[list[str], dict[str, Any]]:
@@ -408,7 +399,7 @@ class Formatters(BaseActivity):
Args:
input_data (dict[str, Any]): The input data containing schedule reports.
Each item should have 'namespace', 'schedule_name', 'success', 'message',
Each item should have 'namespace', 'schedule_name', 'success', 'message',
and optionally 'attachment' fields.
Returns:
@@ -416,14 +407,20 @@ class Formatters(BaseActivity):
- List of successful schedule keys in format "namespace/schedule_name"
- Dictionary of error keys mapped to their error details
"""
success_keys = [f"{value['namespace']}/{value['schedule_name']}"
for value in input_data if value['success']]
success_keys = [
f'{value["namespace"]}/{value["schedule_name"]}'
for value in input_data
if value['success']
]
error_keys = {f"{value['namespace']}/{value['schedule_name']}": {
'message': value['message'],
'attachment': value.get('attachment', None)
error_keys = {
f'{value["namespace"]}/{value["schedule_name"]}': {
'message': value['message'],
'attachment': value.get('attachment', None),
}
for value in input_data
if not value['success']
}
for value in input_data if not value['success']}
return success_keys, error_keys
@@ -440,15 +437,20 @@ class Formatters(BaseActivity):
- List of successful keys
- List of error keys
"""
success_keys = [key for key, value
in input_data.items() if value['success']]
success_keys = [key for key, value in input_data.items() if value['success']]
error_keys = [key for key, value
in input_data.items() if not value['success']]
error_keys = [key for key, value in input_data.items() if not value['success']]
return success_keys, error_keys
def manage_and_send_report(self, metadata: dict[str, Any], success_keys: list[str], error_keys: dict[str, Any], schedule_type: str, schedule_data: dict[str, Any]):
def manage_and_send_report(
self,
metadata: dict[str, Any],
success_keys: list[str],
error_keys: dict[str, Any],
schedule_type: str,
schedule_data: dict[str, Any],
):
"""
Manages and sends success and error reports based on the provided keys and data.
@@ -462,30 +464,28 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
metadata=metadata,
message=f"Successfully {schedule_type}: \n {', '.join(success_keys)}",
message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}',
notification_id=schedule_data['id'],
attachment=schedule_data['items']
attachment=schedule_data['items'],
)
if len(error_keys) > 0:
attachment = []
for key, value in error_keys.items():
if value['attachment'] is not None:
attachment.append(
f"{key}:\n{value['message']}\n{value['attachment']}")
attachment.append(f'{key}:\n{value["message"]}\n{value["attachment"]}')
else:
attachment.append(f"{key}:\n{value['message']}")
attachment.append(f'{key}:\n{value["message"]}')
self.send_error_report(
metadata=metadata,
message=f"Fails on {schedule_type}: \n {', '.join(error_keys)}",
notification_id=f"{schedule_data['id']}_ERROR",
attachment=topic_separator.join(attachment)
message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}',
notification_id=f'{schedule_data["id"]}_ERROR',
attachment=topic_separator.join(attachment),
)
@activity.defn(name="report_schedule_orchestration")
async def report_schedule_orchestration(self,
input_data: dict[str, Any]) -> None:
@activity.defn(name='report_schedule_orchestration')
async def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None:
"""
Reports the orchestration result to the notification handler.
@@ -496,9 +496,9 @@ class Formatters(BaseActivity):
- updated_schedules (dict[str, Any]): The updated schedules.
- deleted_schedules (list[str]): The deleted schedules.
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Reporting orchestration...", metadata=metadata)
self.info('Reporting orchestration...', metadata=metadata)
created_schedules = input_data['created_schedules']
updated_schedules = input_data['updated_schedules']
@@ -507,34 +507,32 @@ class Formatters(BaseActivity):
schedules_report = {
'created schedules': {
'items': created_schedules,
'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES'
'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES',
},
'updated schedules': {
'items': updated_schedules,
'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES'
'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES',
},
'deleted schedules': {
'items': deleted_schedules,
'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES'
}
'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES',
},
}
for schedule_type, schedule_data in schedules_report.items():
if len(schedule_data['items']) > 0:
success_keys, error_keys = self.parse_report_schedule(
schedule_data['items'])
success_keys, error_keys = self.parse_report_schedule(schedule_data['items'])
self.manage_and_send_report(
metadata=metadata,
success_keys=success_keys,
error_keys=error_keys,
schedule_type=schedule_type,
schedule_data=schedule_data
schedule_data=schedule_data,
)
@activity.defn(name="report_slot_orchestration")
async def report_slot_orchestration(self,
input_data: dict[str, Any]) -> None:
@activity.defn(name='report_slot_orchestration')
async def report_slot_orchestration(self, input_data: dict[str, Any]) -> None:
"""
Reports the orchestration result to the notification handler.
@@ -545,9 +543,9 @@ class Formatters(BaseActivity):
- deleted_slots (list[str]): The deleted slots.
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Reporting orchestration...", metadata=metadata)
self.info('Reporting orchestration...', metadata=metadata)
inserted_slots = input_data['inserted_slots']
deleted_slots = input_data['deleted_slots']
@@ -558,16 +556,16 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
metadata=metadata,
message=f"Inserted slots: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS"
message=f'Inserted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
)
if len(error_keys) > 0:
self.send_error_report(
metadata=metadata,
message=f"Failed to insert slots: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS",
attachment=inserted_slots
message=f'Failed to insert slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
attachment=inserted_slots,
)
if len(deleted_slots) > 0:
@@ -576,19 +574,19 @@ class Formatters(BaseActivity):
if len(success_keys) > 0:
self.send_success_report(
metadata=metadata,
message=f"Deleted slots: \n {', '.join(success_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS"
message=f'Deleted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
)
if len(error_keys) > 0:
self.send_error_report(
metadata=metadata,
message=f"Failed to delete slots: \n {', '.join(error_keys)}",
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS",
attachment=deleted_slots
message=f'Failed to delete slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
attachment=deleted_slots,
)
@activity.defn(name="format_log_report")
@activity.defn(name='format_log_report')
async def format_log_report(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Formats the receiver_groups status to a dataframe to be stored in the database.
@@ -602,23 +600,21 @@ class Formatters(BaseActivity):
Returns:
dict[str, Any]: The formatted log report as a dictionary representation of a DataFrame.
"""
metadata = input_data["metadata"]
mail_type = input_data["mail_type"]
metadata = input_data['metadata']
mail_type = input_data['mail_type']
self.info("Formatting log report...", metadata=metadata)
self.info('Formatting log report...', metadata=metadata)
receiver_groups = input_data['receiver_groups']
data = {}
for group_name, group_config in receiver_groups.items():
for notification in group_config['notifications']:
notification_id = notification['notification_id']
trigger = notification['trigger']
key = f"{notification_id}:{trigger}"
key = f'{notification_id}:{trigger}'
if key not in data:
data[key] = {
@@ -634,7 +630,7 @@ class Formatters(BaseActivity):
'project': notification['project'],
'model_name': notification['model_name'],
'model_id': notification['model_id'],
'mail_type': mail_type
'mail_type': mail_type,
}
else:
if group_name not in data[key]['groups']:
@@ -642,7 +638,7 @@ class Formatters(BaseActivity):
return DataFrame(list(data.values())).to_dict()
@activity.defn(name="filter_notification_reports")
@activity.defn(name='filter_notification_reports')
async def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Filter notifications for comprehensive scheduled reports.
@@ -678,16 +674,13 @@ class Formatters(BaseActivity):
notification_package = input_data['notification_package']
sending_configs = input_data['sending_configs']
self.info("Filtering notification reports...", metadata=metadata)
self.info('Filtering notification reports...', metadata=metadata)
receiver_groups = {}
for receiver_group in sending_configs:
group_name = receiver_group['group_name']
receiver_groups[group_name] = {
**receiver_group,
"notifications": []
}
receiver_groups[group_name] = {**receiver_group, 'notifications': []}
receiver_groups[group_name]['notifications'] = []
already_added_keys = []
@@ -695,15 +688,18 @@ class Formatters(BaseActivity):
ignore_list = receiver_group.get('ignore', [])
for notification in notification_package:
alert_type = "reports"
alert_type = 'reports'
notification_id = notification['notification_id']
key = f"{notification['trigger']}:{notification_id}"
key = f'{notification["trigger"]}:{notification_id}'
# Check if this group must be notified
if alert_type in receiver_group['contents'] and notification_id not in ignore_list and key not in already_added_keys:
receiver_groups[group_name]["notifications"].append(
notification)
if (
alert_type in receiver_group['contents']
and notification_id not in ignore_list
and key not in already_added_keys
):
receiver_groups[group_name]['notifications'].append(notification)
already_added_keys.append(key)
# Remove groups with no notifications

327
orchestrator/activities/mongo_db.py
View File

@@ -1,16 +1,18 @@
from temporalio import workflow, activity
from datetime import UTC
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
import traceback
from datetime import datetime
from logging import Logger
from typing import Any
from pymongo import MongoClient
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from datetime import datetime, timezone
def clear_mongo_id(docs: list) -> list:
@@ -28,10 +30,10 @@ def clear_mongo_id(docs: list) -> list:
clear_mongo_id(doc)
elif isinstance(doc, dict):
if "_id" in doc:
del doc["_id"]
if '_id' in doc:
del doc['_id']
for key, value in doc.items():
for _key, value in doc.items():
if isinstance(value, list):
clear_mongo_id(value)
elif isinstance(value, dict):
@@ -57,14 +59,18 @@ class MongoDB(BaseActivity):
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self, connection_string: str, database_name: str, ttl_index_seconds: int,
logger: Logger,
notification_handler: NotificationHandler):
def __init__(
self,
connection_string: str,
database_name: str,
ttl_index_seconds: int,
logger: Logger,
notification_handler: NotificationHandler,
):
self.connection_string = connection_string
self.database_name = database_name
self.client = MongoClient(
self.connection_string, serverSelectionTimeoutMS=5000)
self.client = MongoClient(self.connection_string, serverSelectionTimeoutMS=5000)
self.client.server_info() # Trigger an exception if connection fails
self.database = self.client[self.database_name]
@@ -72,11 +78,9 @@ class MongoDB(BaseActivity):
self.ttl_index_seconds = ttl_index_seconds
# Initialize MongoDB client here (omitted for brevity)
logger.info("MongoDB connection initialized")
logger.info('MongoDB connection initialized')
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
def shutdown(self):
"""
@@ -84,11 +88,11 @@ class MongoDB(BaseActivity):
"""
try:
if self.client:
self.logger.info("Closing MongoDB connection...")
self.logger.info('Closing MongoDB connection...')
self.client.close()
self.logger.info("MongoDB connection closed successfully")
self.logger.info('MongoDB connection closed successfully')
except Exception as e:
self.logger.error(f"Failed to close MongoDB connection: {e}")
self.logger.error(f'Failed to close MongoDB connection: {e}')
def __del__(self):
"""
@@ -109,13 +113,15 @@ class MongoDB(BaseActivity):
"""
collection = self.database[collection_name]
documents = list(collection.find(filters, {"_id": 0}))
documents = list(collection.find(filters, {'_id': 0}))
documents = clear_mongo_id(documents)
return documents
@activity.defn(name="find_documents_in_mongodb",)
@activity.defn(
name='find_documents_in_mongodb',
)
async def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Find documents in a MongoDB collection based on the provided query parameters.
@@ -128,33 +134,39 @@ class MongoDB(BaseActivity):
list[dict]: List of documents matching the query.
"""
query = input_data.get("query", {})
metadata = input_data.get("metadata", {})
timestamp_fields = input_data.get("timestamp_fields", [])
query = input_data.get('query', {})
metadata = input_data.get('metadata', {})
timestamp_fields = input_data.get('timestamp_fields', [])
collection_name = query.get("collection")
collection_name = query.get('collection')
if not collection_name:
raise ValueError("Collection name must be provided in the query.")
raise ValueError('Collection name must be provided in the query.')
filters = query.get("filters", {})
filters = query.get('filters', {})
self.info(
f"Loading documents from collection '{collection_name}' with filters: {filters}", metadata=metadata)
f"Loading documents from collection '{collection_name}' with filters: {filters}",
metadata=metadata,
)
try:
documents = self.find(collection_name, filters)
self.info(
f"Loaded {len(documents)} documents from collection '{collection_name}'", metadata=metadata)
f"Loaded {len(documents)} documents from collection '{collection_name}'",
metadata=metadata,
)
for document in documents:
for timestamp_field in timestamp_fields:
if timestamp_field in document:
document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime(
DATETIME_FORMAT_MS_WITH_TZ)
document[timestamp_field] = (
document[timestamp_field]
.replace(tzinfo=UTC)
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
)
self.debug(
f"Documents loaded: {documents}", metadata=metadata)
self.debug(f'Documents loaded: {documents}', metadata=metadata)
return documents
@@ -162,19 +174,20 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_QUERY_ERROR",
message=f"Failed to execute MongoDB query: {e}",
block="load_query_from_mongodb",
notification_id='MONGODB_QUERY_ERROR',
message=f'Failed to execute MongoDB query: {e}',
block='load_query_from_mongodb',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name="aggregate_documents_in_mongodb")
async def aggregate_documents_in_mongodb(self,
input_data: dict[str, Any]) -> list[dict[str, Any]]:
@activity.defn(name='aggregate_documents_in_mongodb')
async def aggregate_documents_in_mongodb(
self, input_data: dict[str, Any]
) -> list[dict[str, Any]]:
"""
Aggregate documents in a MongoDB collection based on the provided aggregation pipeline.
@@ -186,40 +199,45 @@ class MongoDB(BaseActivity):
list[dict]: List of aggregated documents.
"""
query = input_data.get("query", {})
metadata = input_data.get("metadata", {})
timestamp_fields = input_data.get("timestamp_fields", [])
query = input_data.get('query', {})
metadata = input_data.get('metadata', {})
timestamp_fields = input_data.get('timestamp_fields', [])
collection_name = query.get("collection")
collection_name = query.get('collection')
if not collection_name:
raise ValueError("Collection name must be provided in the query.")
aggregation = query.get("aggregation")
raise ValueError('Collection name must be provided in the query.')
aggregation = query.get('aggregation')
if not aggregation:
raise ValueError("Aggregation must be provided.")
aggregation.append({"$project": {"_id": 0}})
raise ValueError('Aggregation must be provided.')
aggregation.append({'$project': {'_id': 0}})
self.info(
f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}", metadata=metadata)
f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}",
metadata=metadata,
)
try:
collection = self.database[collection_name]
aggregated_documents = list(
collection.aggregate(aggregation))
aggregated_documents = list(collection.aggregate(aggregation))
aggregated_documents = clear_mongo_id(aggregated_documents)
self.info(
f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'", metadata=metadata)
f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'",
metadata=metadata,
)
for document in aggregated_documents:
for timestamp_field in timestamp_fields:
if timestamp_field in document:
document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime(
DATETIME_FORMAT_MS_WITH_TZ)
document[timestamp_field] = (
document[timestamp_field]
.replace(tzinfo=UTC)
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
)
self.debug(
f"Aggregation result: {aggregated_documents}", metadata=metadata)
self.debug(f'Aggregation result: {aggregated_documents}', metadata=metadata)
return aggregated_documents
@@ -227,83 +245,85 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_AGGREGATION_ERROR",
message=f"Failed to execute MongoDB aggregation: {e}",
block="aggregate_documents_in_mongodb",
notification_id='MONGODB_AGGREGATION_ERROR',
message=f'Failed to execute MongoDB aggregation: {e}',
block='aggregate_documents_in_mongodb',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name="update_pipelines_timestamps")
@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", [])
metadata = input_data.get("metadata", {})
updated_pipelines = input_data.get('updated_pipelines', [])
metadata = input_data.get('metadata', {})
date_now = now()
collection = self.database["orchestrated_schedules"]
collection = self.database['orchestrated_schedules']
self.info("Updating pipelines timestamps...", metadata=metadata)
self.info('Updating pipelines timestamps...', metadata=metadata)
success_count = 0
argument = [
{"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"]}
for pipeline in updated_pipelines if pipeline["success"]
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
for pipeline in updated_pipelines
if pipeline['success']
]
data_filter = {"$or": argument} if argument else {}
data_filter = {'$or': argument} if argument else {}
try:
collection.update_many(
data_filter,
{"$set": {"updated_at": date_now}}
)
collection.update_many(data_filter, {'$set': {'updated_at': date_now}})
success_count += 1
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_UPDATE_PIPELINES_ERROR",
message=f"Failed to update pipelines timestamps: {e}",
block="update_pipelines_timestamps",
notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message=f'Failed to update pipelines timestamps: {e}',
block='update_pipelines_timestamps',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(
f"Updated {success_count} of {len(updated_pipelines)} pipelines timestamps", metadata=metadata)
f'Updated {success_count} of {len(updated_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="create_pipelines_timestamps")
@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", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
created_pipelines = input_data.get('created_pipelines', [])
metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info("Creating pipelines timestamps...", metadata=metadata)
self.info('Creating pipelines timestamps...', metadata=metadata)
success_count = 0
date_now = now()
argument = [
{"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"],
"updated_at": date_now}
for pipeline in created_pipelines if pipeline["success"]
{
'schedule_name': pipeline['schedule_name'],
'namespace': pipeline['namespace'],
'updated_at': date_now,
}
for pipeline in created_pipelines
if pipeline['success']
]
data_filter = argument if argument else {}
@@ -315,39 +335,41 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_CREATE_PIPELINES_ERROR",
message=f"Failed to create pipelines timestamps: {e}",
block="create_pipelines_timestamps",
notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message=f'Failed to create pipelines timestamps: {e}',
block='create_pipelines_timestamps',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(
f"Created {success_count} of {len(created_pipelines)} pipelines timestamps", metadata=metadata)
f'Created {success_count} of {len(created_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="delete_pipelines_timestamps")
@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", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
deleted_pipelines = input_data.get('deleted_pipelines', [])
metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info("Deleting pipelines timestamps...", metadata=metadata)
self.info('Deleting pipelines timestamps...', metadata=metadata)
success_count = 0
argument = [
{"schedule_name": pipeline["schedule_name"],
"namespace": pipeline["namespace"]}
for pipeline in deleted_pipelines if pipeline["success"]
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
for pipeline in deleted_pipelines
if pipeline['success']
]
data_filter = {"$or": argument} if argument else {}
data_filter = {'$or': argument} if argument else {}
try:
collection.delete_many(data_filter)
@@ -356,19 +378,21 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_DELETE_PIPELINES_ERROR",
message=f"Failed to delete pipelines timestamps: {e}",
block="delete_pipelines_timestamps",
notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message=f'Failed to delete pipelines timestamps: {e}',
block='delete_pipelines_timestamps',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(
f"Deleted {success_count} of {len(deleted_pipelines)} pipelines timestamps", metadata=metadata)
f'Deleted {success_count} of {len(deleted_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="create_collection_with_ttl_index")
@activity.defn(name='create_collection_with_ttl_index')
async def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
"""
Create a collection with a TTL index.
@@ -376,20 +400,21 @@ class MongoDB(BaseActivity):
- collection_name (str): The name of the collection to create.
- ttl_index (str): The name of the TTL index to create.
"""
pipelines = input_data.get("pipelines", {})
metadata = input_data.get("metadata", {})
pipelines = input_data.get('pipelines', {})
metadata = input_data.get('metadata', {})
self.info(
f"Creating collection with TTL index for pipelines: {list(pipelines.keys())}",
metadata=metadata)
f'Creating collection with TTL index for pipelines: {list(pipelines.keys())}',
metadata=metadata,
)
collection_names = self.database.list_collection_names()
created_collections = []
created_indexes = []
for pipeline_name, pipeline_config in pipelines.items():
collection = pipeline_config["topic"]
for _pipeline_name, pipeline_config in pipelines.items():
collection = pipeline_config['topic']
try:
# Check if collection exists
@@ -402,16 +427,17 @@ class MongoDB(BaseActivity):
existing_indexes = collection.list_indexes()
ttl_index_exists = False
for index in existing_indexes:
if "inserted_at" in index["key"] and index.get("expireAfterSeconds") is not None:
if (
'inserted_at' in index['key']
and index.get('expireAfterSeconds') is not None
):
ttl_index_exists = True
break
# Create TTL index if it doesn't exist
if not ttl_index_exists:
collection.create_index(
"inserted_at",
expireAfterSeconds=self.ttl_index_seconds,
background=True
'inserted_at', expireAfterSeconds=self.ttl_index_seconds, background=True
)
created_indexes.append(collection)
@@ -419,30 +445,24 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGODB_CREATE_COLLECTION_ERROR",
message=f"Failed to create collection {collection} with TTL index: {e}",
block="create_collection_with_ttl_index",
notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message=f'Failed to create collection {collection} with TTL index: {e}',
block='create_collection_with_ttl_index',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(
f"Created {len(created_collections)} collections and {len(created_indexes)} indexes",
metadata=metadata
f'Created {len(created_collections)} collections and {len(created_indexes)} indexes',
metadata=metadata,
)
self.debug(
f"Created collections: {created_collections}",
metadata=metadata
)
self.debug(
f"Created indexes: {created_indexes}",
metadata=metadata
)
self.debug(f'Created collections: {created_collections}', metadata=metadata)
self.debug(f'Created indexes: {created_indexes}', metadata=metadata)
@activity.defn(name="load_latest_data")
@activity.defn(name='load_latest_data')
async def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Load the latest data from MongoDB collection since a specified timestamp.
@@ -470,58 +490,43 @@ class MongoDB(BaseActivity):
last_data_timestamp = input_data['last_data_timestamp']
base_data_filter = input_data['base_data_filter']
self.debug(
f"Loading data from MongoDB: {input_data}",
metadata=metadata
)
self.debug(f'Loading data from MongoDB: {input_data}', metadata=metadata)
try:
if last_data_timestamp is None:
data_filter = base_data_filter
else:
data_filter = {
**base_data_filter,
"timestamp": {
"$gt": datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ)
}
'timestamp': {
'$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ)
},
}
self.debug(
f"Data filter: {data_filter}",
metadata=metadata
)
self.debug(f'Data filter: {data_filter}', metadata=metadata)
data = self.find(collection_name, data_filter)
self.debug(
f"Collected: {data}",
metadata=metadata
)
self.debug(f'Collected: {data}', metadata=metadata)
for item in data:
item['timestamp'] = item['timestamp'].replace(tzinfo=timezone.utc).strftime(
DATETIME_FORMAT_MS_WITH_TZ)
item['timestamp'] = (
item['timestamp'].replace(tzinfo=UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ)
)
self.info(
f"Loaded {len(data)} documents from MongoDB",
metadata=metadata
)
self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata)
self.debug(
f"Loaded data: {data}",
metadata=metadata
)
self.debug(f'Loaded data: {data}', metadata=metadata)
return data
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="MONGO_LOAD_ERROR",
message=f"Error loading data from MongoDB: {e}",
block="load_latest_data",
notification_id='MONGO_LOAD_ERROR',
message=f'Error loading data from MongoDB: {e}',
block='load_latest_data',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
raise e

206
orchestrator/activities/slot_manager.py
View File

@@ -1,17 +1,17 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
import traceback
import json
import traceback
from datetime import datetime, timedelta
from logging import Logger
from sientia_do.temporal.activities.redis_base import Redis
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.redis_base import Redis
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from pandas import DataFrame
from datetime import datetime, timedelta
class SlotManager(Redis):
@@ -40,14 +40,18 @@ class SlotManager(Redis):
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self, host: str, port: int,
username: str, password: str,
logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
host: str,
port: int,
username: str,
password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
Redis.__init__(self, host, port, username, password, logger, notification_handler)
Redis.__init__(self, host, port, username,
password, logger, notification_handler)
@activity.defn(name="load_opc_slots")
@activity.defn(name='load_opc_slots')
async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Load all OPC slots from Redis for current system state assessment.
@@ -69,17 +73,16 @@ class SlotManager(Redis):
Exception: If Redis connection fails or data retrieval errors occur
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Loading OPC slots...", metadata=metadata)
self.info('Loading OPC slots...', metadata=metadata)
opc_slots = {}
try:
slot_keys = self.redis_client.keys('slot:opc_tags:*')
slot_keys = self.redis_client.keys("slot:opc_tags:*")
self.debug(f"Slot keys: {slot_keys}", metadata=metadata)
self.debug(f'Slot keys: {slot_keys}', metadata=metadata)
if slot_keys:
if isinstance(slot_keys[0], bytes):
@@ -93,20 +96,20 @@ class SlotManager(Redis):
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",
notification_id='REDIS_GET_ERROR',
message=f'Failed to load OPC slots: {e}',
block='load_opc_slots',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(f"Loaded {len(opc_slots)} OPC slots", metadata=metadata)
self.info(f'Loaded {len(opc_slots)} OPC slots', metadata=metadata)
return opc_slots
@activity.defn(name="load_active_ingestors")
@activity.defn(name='load_active_ingestors')
async def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
"""
Load all active ingestors from Redis
@@ -115,19 +118,16 @@ class SlotManager(Redis):
list[str]: A list of active ingestors
"""
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Loading active ingestors...", metadata=metadata)
self.info('Loading active ingestors...', metadata=metadata)
try:
active_ingestors = self.redis_client.keys('heartbeat:ingestor:*')
active_ingestors = self.redis_client.keys("heartbeat:ingestor:*")
self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata)
self.info(
f"Loaded {len(active_ingestors)} active ingestors", metadata=metadata)
self.debug(
f"Active ingestors: \n {active_ingestors}", metadata=metadata)
self.debug(f'Active ingestors: \n {active_ingestors}', metadata=metadata)
ingestors = []
@@ -142,16 +142,16 @@ class SlotManager(Redis):
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",
notification_id='REDIS_GET_ERROR',
message=f'Failed to load active ingestors: {e}',
block='load_active_ingestors',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name="update_slots")
@activity.defn(name='update_slots')
async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Update OPC slots in Redis
@@ -166,8 +166,8 @@ class SlotManager(Redis):
"""
to_insert = input_data['to_insert']
metadata = input_data.get("metadata", {})
self.info("Updating OPC slots...", metadata=metadata)
metadata = input_data.get('metadata', {})
self.info('Updating OPC slots...', metadata=metadata)
report = {}
@@ -175,30 +175,20 @@ class SlotManager(Redis):
for slot in to_insert:
try:
self.set(f"slot:opc_tags:{slot}",
to_insert[slot], ttl=None)
report[slot] = {
"success": True,
"message": "Slot updated successfully"
}
self.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
report[slot] = {'success': True, 'message': 'Slot updated successfully'}
success_count += 1
except Exception as e:
self.error(
f"Failed to update slot {slot}: {str(e)}", metadata=metadata)
report[slot] = {
"success": False,
"message": str(e)
}
self.error(f'Failed to update slot {slot}: {str(e)}', metadata=metadata)
report[slot] = {'success': False, 'message': str(e)}
self.info(
f"Updated {success_count} of {len(to_insert)} OPC slots", metadata=metadata)
self.info(f'Updated {success_count} of {len(to_insert)} OPC slots', metadata=metadata)
self.debug(
f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
return report
@activity.defn(name="delete_slots")
@activity.defn(name='delete_slots')
async def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Delete OPC slots from Redis
@@ -213,9 +203,9 @@ class SlotManager(Redis):
"""
to_delete = input_data['to_delete']
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
self.info("Deleting OPC slots...", metadata=metadata)
self.info('Deleting OPC slots...', metadata=metadata)
report = {}
@@ -223,29 +213,20 @@ class SlotManager(Redis):
for slot in to_delete:
try:
self.redis_client.delete(f"slot:opc_tags:{slot}")
report[slot] = {
"success": True,
"message": "Slot deleted successfully"
}
self.redis_client.delete(f'slot:opc_tags:{slot}')
report[slot] = {'success': True, 'message': 'Slot deleted successfully'}
success_count += 1
except Exception as e:
self.error(
f"Failed to delete slot {slot}: {str(e)}", metadata=metadata)
report[slot] = {
"success": False,
"message": str(e)
}
self.error(f'Failed to delete slot {slot}: {str(e)}', metadata=metadata)
report[slot] = {'success': False, 'message': str(e)}
self.info(
f"Deleted {success_count} of {len(to_delete)} OPC slots", metadata=metadata)
self.info(f'Deleted {success_count} of {len(to_delete)} OPC slots', metadata=metadata)
self.debug(
f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
return report
@activity.defn(name="get_last_data_timestamp")
@activity.defn(name='get_last_data_timestamp')
async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
"""
Gets the last data timestamp from redis.
@@ -259,32 +240,29 @@ class SlotManager(Redis):
str | None: The last data timestamp as a string, or None if no timestamp exists.
"""
metadata = input_data['metadata']
key = f"notification_last_timestamp:{input_data['mail_type']}"
key = f'notification_last_timestamp:{input_data["mail_type"]}'
try:
data_hold = self.get(key)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_GET_ERROR",
message=f"Error getting last data timestamp: {e}",
block="get_last_data_timestamp",
notification_id='REDIS_GET_ERROR',
message=f'Error getting last data timestamp: {e}',
block='get_last_data_timestamp',
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
attachment_content=traceback.format_exc(),
)
raise e
self.debug(
f"Last collected timestamp: {data_hold}",
metadata=metadata
)
self.debug(f'Last collected timestamp: {data_hold}', metadata=metadata)
if not data_hold:
return None
return data_hold
@activity.defn(name="put_last_data_timestamp")
@activity.defn(name='put_last_data_timestamp')
async def put_last_data_timestamp(self, input_data: dict[str, Any]):
"""
Puts the last data timestamp into redis.
@@ -299,39 +277,34 @@ class SlotManager(Redis):
str | None: The last data timestamp that was stored, or None if no data exists.
"""
metadata = input_data['metadata']
key = f"notification_last_timestamp:{input_data['mail_type']}"
key = f'notification_last_timestamp:{input_data["mail_type"]}'
data = DataFrame(input_data['data'])
if data.empty:
self.warning("No data to insert",
metadata=metadata
)
self.warning('No data to insert', metadata=metadata)
return None
last_data_timestamp = data['timestamp'].max()
self.debug(
f"Last collected timestamp to insert: {last_data_timestamp}",
metadata=metadata
)
self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
try:
self.set(key, last_data_timestamp, ttl=60*60*5)
self.set(key, last_data_timestamp, ttl=60 * 60 * 5)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_SET_ERROR",
message=f"Error setting last data timestamp: {e}",
block="put_last_data_timestamp",
notification_id='REDIS_SET_ERROR',
message=f'Error setting last data timestamp: {e}',
block='put_last_data_timestamp',
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
attachment_content=traceback.format_exc(),
)
raise e
return last_data_timestamp
@activity.defn(name="filter_notification_alerts")
@activity.defn(name='filter_notification_alerts')
async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Filter notification alerts with intelligent TTL-based duplicate prevention.
@@ -369,16 +342,13 @@ class SlotManager(Redis):
sending_configs = input_data['sending_configs']
notification_ttl = input_data['notification_ttl']
self.info("Filtering notification alerts...", metadata=metadata)
self.info('Filtering notification alerts...', metadata=metadata)
receiver_groups = {}
for receiver_group in sending_configs:
group_name = receiver_group['group_name']
receiver_groups[group_name] = {
**receiver_group,
"notifications": []
}
receiver_groups[group_name] = {**receiver_group, 'notifications': []}
receiver_groups[group_name]['notifications'] = []
already_added_keys = []
@@ -386,28 +356,30 @@ class SlotManager(Redis):
ignore_list = receiver_group.get('ignore', [])
for notification in notification_package:
alert_type = "do_nothing"
alert_type = 'do_nothing'
notification_id = notification['notification_id']
# Check if notification was recently sent
key = f"{notification['trigger']}:{notification_id}"
key = f'{notification["trigger"]}:{notification_id}'
last_sent = self.get(key)
if last_sent is None:
alert_type = "core_alerts"
alert_type = 'core_alerts'
else:
last_sent = datetime.strptime(
last_sent, DATETIME_FORMAT_MS_WITH_TZ)
last_sent = datetime.strptime(last_sent, DATETIME_FORMAT_MS_WITH_TZ)
# Check if "notification_ttl" seconds has passed since last sent
if (now() - last_sent) > timedelta(seconds=notification_ttl):
alert_type = "persistent_alerts"
alert_type = 'persistent_alerts'
# Check if this group must be notified
if alert_type in receiver_group['contents'] and notification_id not in ignore_list and key not in already_added_keys:
receiver_groups[group_name]["notifications"].append(
notification)
if (
alert_type in receiver_group['contents']
and notification_id not in ignore_list
and key not in already_added_keys
):
receiver_groups[group_name]['notifications'].append(notification)
already_added_keys.append(key)
# Remove groups with no notifications
@@ -419,7 +391,7 @@ class SlotManager(Redis):
return receiver_groups
@activity.defn(name="store_notification_cache")
@activity.defn(name='store_notification_cache')
async def store_notification_cache(self, input_data: dict[str, Any]) -> None:
"""
Store notification cache in Redis to track recently sent notifications.
@@ -434,14 +406,14 @@ class SlotManager(Redis):
log_report = DataFrame(input_data['log_report'])
sent_ttl = input_data['sent_ttl']
self.info("Storing notification cache...", metadata=metadata)
self.info('Storing notification cache...', metadata=metadata)
date_now = now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
for index, row in log_report.iterrows():
for _index, row in log_report.iterrows():
status = row['status']
if status == 'sent':
key = f"{row['schedule']}:{row['notification_id']}"
key = f'{row["schedule"]}:{row["notification_id"]}'
self.set(key, date_now, ttl=sent_ttl)
self.info("Notification cache stored...", metadata=metadata)
self.info('Notification cache stored...', metadata=metadata)

308
orchestrator/activities/temporal_manager.py
View File

@@ -1,20 +1,26 @@
from temporalio import activity, workflow
from temporalio.client import (
Client, Schedule, ScheduleActionStartWorkflow, ScheduleIntervalSpec, ScheduleSpec, ScheduleUpdate, ScheduleUpdateInput)
Client,
Schedule,
ScheduleActionStartWorkflow,
ScheduleIntervalSpec,
ScheduleSpec,
ScheduleUpdate,
ScheduleUpdateInput,
)
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
import json
from asyncio import sleep
import traceback
from datetime import timedelta
from logging import Logger
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
from orchestrator.utils.converters import parse_frequency
@@ -30,69 +36,65 @@ class TemporalManager(BaseActivity):
Args:
host (str): Temporal server host address
scouter_namespace (str): Scouter workflow namespace
laborious_namespace (str): Laborious workflow namespace
laborious_namespace (str): Laborious workflow namespace
logger (Logger): Application logger instance
notification_handler (NotificationHandler): Notification management handler
"""
def __init__(self, host: str, scouter_namespace: str, laborious_namespace: str,
task_timeout: int, run_timeout: int, execution_timeout: int, logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
host: str,
scouter_namespace: str,
laborious_namespace: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.temporal_host = host
self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace
self.task_timeout = task_timeout
self.run_timeout = run_timeout
self.execution_timeout = execution_timeout
self.temporal_clients = {}
self.model_id_id_key = SearchAttributeKey.for_keyword("model_id")
self.model_name_id_key = SearchAttributeKey.for_keyword("model_name")
self.orchestrated_id_key = SearchAttributeKey.for_keyword(
"orchestrated")
self.model_id_id_key = SearchAttributeKey.for_keyword('model_id')
self.model_name_id_key = SearchAttributeKey.for_keyword('model_name')
self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated')
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
async def connect_to_temporal(self):
"""
Connect to Temporal server namespaces for scouter and laborious workflows.
Creates client connections to both namespaces and stores them for later use.
"""
self.logger.info(
f"Connecting to Temporal side namespaces at {self.temporal_host}")
self.logger.info(f"Scouter namespace: {self.scouter_namespace}")
self.logger.info(f'Connecting to Temporal side namespaces at {self.temporal_host}')
self.logger.info(f'Scouter namespace: {self.scouter_namespace}')
scouter_client = await Client.connect(
target_host=self.temporal_host,
namespace=self.scouter_namespace
target_host=self.temporal_host, namespace=self.scouter_namespace
)
self.logger.info(f"Laborious namespace: {self.laborious_namespace}")
self.logger.info(f'Laborious namespace: {self.laborious_namespace}')
laborious_client = await Client.connect(
target_host=self.temporal_host,
namespace=self.laborious_namespace
target_host=self.temporal_host, namespace=self.laborious_namespace
)
self.temporal_clients = {
self.scouter_namespace: scouter_client,
self.laborious_namespace: laborious_client
self.laborious_namespace: laborious_client,
}
@activity.defn(name="normalize_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.
"""
metadata = input_data["metadata"]
metadata = input_data['metadata']
remove_count = 0
self.info("Getting orchestrated schedules...", metadata=metadata)
self.info('Getting orchestrated schedules...', metadata=metadata)
orchestrated_schedules = input_data.get('orchestrated_schedules', {})
@@ -100,20 +102,20 @@ class TemporalManager(BaseActivity):
try:
schedules = orchestrated_schedules.get(namespace, {})
self.info(
f"Getting orchestrated schedules for {namespace}", metadata=metadata)
self.info(f'Getting orchestrated schedules for {namespace}', metadata=metadata)
async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]:
search_attrs = getattr(schedule, 'search_attributes', {})
if search_attrs.get('orchestrated', ['false']) == ['true']:
schedule_id = schedule.id
if schedule_id not in schedules:
self.info(
f"Schedule {schedule_id} not found in mongo db, cleaning up", metadata=metadata)
f'Schedule {schedule_id} not found in mongo db, cleaning up',
metadata=metadata,
)
handle = client.get_schedule_handle(
schedule_id)
handle = client.get_schedule_handle(schedule_id)
await handle.delete()
@@ -123,19 +125,18 @@ class TemporalManager(BaseActivity):
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",
notification_id='TEMPORAL_NORMALIZE_SCHEDULES_ERROR',
message=f'Failed to normalize schedules: {e}',
block='normalize_schedules',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
self.info(
f"Removed {remove_count} schedules", metadata=metadata)
self.info(f'Removed {remove_count} schedules', metadata=metadata)
@activity.defn(name="create_schedules")
@activity.defn(name='create_schedules')
async def create_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Create schedules in Temporal
@@ -150,47 +151,41 @@ class TemporalManager(BaseActivity):
"""
schedules_to_create = input_data['schedules']
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
report = []
success_count = 0
self.info("Creating schedules...", metadata=metadata)
self.info('Creating schedules...', metadata=metadata)
for namespace, schedules in schedules_to_create.items():
client = self.temporal_clients.get(namespace)
if not client:
raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}")
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name, schedule in schedules.items():
search_attributes = TypedSearchAttributes([
SearchAttributePair(
key=self.model_id_id_key,
value=schedule['model_id']
),
SearchAttributePair(
key=self.model_name_id_key,
value=schedule['model_name']
),
SearchAttributePair(
key=self.orchestrated_id_key,
value="true"
)
])
search_attributes = TypedSearchAttributes(
[
SearchAttributePair(key=self.model_id_id_key, value=schedule['model_id']),
SearchAttributePair(
key=self.model_name_id_key, value=schedule['model_name']
),
SearchAttributePair(key=self.orchestrated_id_key, value='true'),
]
)
workflow_type = schedule['workflow_type']
try:
execution_timeout_seconds = schedule.get(
'execution_timeout_seconds', 300)
task_timeout_seconds = schedule.get(
'task_timeout_seconds', 300)
execution_timeout_seconds = schedule.get('execution_timeout_seconds', 300)
task_timeout_seconds = schedule.get('task_timeout_seconds', 300)
self.debug(f'Creating schedule {schedule_name}:', metadata=metadata)
self.debug(
f"Creating schedule {schedule_name}:", metadata=metadata)
self.debug(
f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata)
f'{json.dumps(schedule, indent=4, sort_keys=True)}', metadata=metadata
)
await client.create_schedule(
schedule_name,
@@ -199,54 +194,57 @@ class TemporalManager(BaseActivity):
workflow_type,
schedule,
id=schedule_name,
task_queue=f"{workflow_type}-queue",
execution_timeout=timedelta(
seconds=execution_timeout_seconds),
run_timeout=timedelta(
seconds=execution_timeout_seconds),
task_timeout=timedelta(
seconds=task_timeout_seconds),
typed_search_attributes=search_attributes
task_queue=f'{workflow_type}-queue',
execution_timeout=timedelta(seconds=execution_timeout_seconds),
run_timeout=timedelta(seconds=execution_timeout_seconds),
task_timeout=timedelta(seconds=task_timeout_seconds),
typed_search_attributes=search_attributes,
),
spec=ScheduleSpec(
intervals=[
ScheduleIntervalSpec(
every=timedelta(seconds=parse_frequency(
schedule.get('frequency', '1m')))
every=timedelta(
seconds=parse_frequency(schedule.get('frequency', '1m'))
)
)
]
)
),
),
search_attributes=search_attributes
search_attributes=search_attributes,
)
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True,
"message": "Schedule created successfully"
})
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': True,
'message': 'Schedule created successfully',
}
)
success_count += 1
except Exception as e:
self.error(
f"Failed to create schedule {schedule_name}: {str(e)}", metadata=metadata)
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False,
"message": str(e)
})
f'Failed to create schedule {schedule_name}: {str(e)}', metadata=metadata
)
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': False,
'message': str(e),
}
)
self.info(
f"Created {success_count} of {len(schedules_to_create)} schedules", metadata=metadata)
f'Created {success_count} of {len(schedules_to_create)} schedules', metadata=metadata
)
self.debug(
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
return report
@activity.defn(name="update_schedules")
@activity.defn(name='update_schedules')
async def update_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Update schedules in Temporal
@@ -261,27 +259,27 @@ class TemporalManager(BaseActivity):
"""
schedules_to_update = input_data['schedules']
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
report = []
success_count = 0
self.info("Updating schedules...", metadata=metadata)
self.info('Updating schedules...', metadata=metadata)
for namespace, schedules in schedules_to_update.items():
client = self.temporal_clients.get(namespace)
if not client:
raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}")
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name, schedule in schedules.items():
try:
handler = client.get_schedule_handle(
schedule_name)
handler = client.get_schedule_handle(schedule_name)
if not handler:
raise ValueError(f"Schedule {schedule_name} not found")
raise ValueError(f'Schedule {schedule_name} not found')
# fmt: off
async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR
@@ -310,33 +308,38 @@ class TemporalManager(BaseActivity):
del update_schedule
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True,
"message": "Schedule updated successfully"
})
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': True,
'message': 'Schedule updated successfully',
}
)
success_count += 1
except Exception as e:
self.error(
f"Failed to update schedule {schedule_name}: {str(e)}", metadata=metadata)
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False,
"message": str(e)
})
f'Failed to update schedule {schedule_name}: {str(e)}', metadata=metadata
)
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': False,
'message': str(e),
}
)
self.info(
f"Updated {success_count} of {len(schedules_to_update)} schedules", metadata=metadata)
f'Updated {success_count} of {len(schedules_to_update)} schedules', metadata=metadata
)
self.debug(
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
return report
@activity.defn(name="delete_schedules")
@activity.defn(name='delete_schedules')
async def delete_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Delete schedules in Temporal
@@ -351,54 +354,59 @@ class TemporalManager(BaseActivity):
"""
schedules_to_delete = input_data['schedules']
metadata = input_data.get("metadata", {})
metadata = input_data.get('metadata', {})
report = []
success_count = 0
self.info("Deleting schedules...", metadata=metadata)
self.info('Deleting schedules...', metadata=metadata)
for namespace, schedules in schedules_to_delete.items():
client = self.temporal_clients.get(namespace)
if not client:
raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}")
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name in schedules:
try:
handler = client.get_schedule_handle(
schedule_name)
handler = client.get_schedule_handle(schedule_name)
if not handler:
raise ValueError(f"Schedule {schedule_name} not found")
raise ValueError(f'Schedule {schedule_name} not found')
await handler.delete()
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": True,
"message": "Schedule deleted successfully"
})
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': True,
'message': 'Schedule deleted successfully',
}
)
success_count += 1
except Exception as e:
trace = traceback.format_exc()
self.error(
f"Failed to delete schedule {schedule_name}: {str(e)}", metadata=metadata)
report.append({
"namespace": namespace,
"schedule_name": schedule_name,
"success": False,
"message": str(e),
"attachment": trace
})
f'Failed to delete schedule {schedule_name}: {str(e)}', metadata=metadata
)
report.append(
{
'namespace': namespace,
'schedule_name': schedule_name,
'success': False,
'message': str(e),
'attachment': trace,
}
)
self.info(
f"Deleted {success_count} of {len(schedules_to_delete)} schedules", metadata=metadata)
f'Deleted {success_count} of {len(schedules_to_delete)} schedules', metadata=metadata
)
self.debug(
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
return report