SIENTIAPDE-1172
feat: enhance orchestrator activities with new email and Postgres integrations - Added Email and Postgres classes to the Activities class for improved functionality. - Introduced new methods in MongoDB and SlotManager for loading and managing data. - Updated requirements.txt to include jinja2. - Added new formatting activity for log reports in Formatters class. - Enhanced test coverage for MongoDB and SlotManager activities.
This commit is contained in:
vitor-aignosi
@@ -1,6 +1,7 @@
|
||||
from temporalio import activity, workflow
|
||||
from temporalio.client import Client
|
||||
|
||||
from orchestrator.activities.email import Email
|
||||
from orchestrator.activities.mongo_db import MongoDB
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
@@ -10,17 +11,21 @@ with workflow.unsafe.imports_passed_through():
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
|
||||
|
||||
class Activities( # Couchbase,
|
||||
TemporalManager, SlotManager, Formatters, MongoDB):
|
||||
TemporalManager, SlotManager, Formatters, MongoDB, Email,
|
||||
Postgres):
|
||||
|
||||
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):
|
||||
|
||||
@@ -59,5 +64,24 @@ class Activities( # Couchbase,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
Email.__init__(self,
|
||||
sender_email=email_config['sender_email'],
|
||||
sender_password=email_config['sender_password'],
|
||||
smpt_server=email_config['smpt_server'],
|
||||
port=email_config['port'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
Postgres.__init__(self,
|
||||
host=postgres_config['host'],
|
||||
port=postgres_config['port'],
|
||||
user=postgres_config['username'],
|
||||
password=postgres_config['password'],
|
||||
dbname=postgres_config['database_name'],
|
||||
min_connections=postgres_config['min_connections'],
|
||||
max_connections=postgres_config['max_connections'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
def shutdown(self):
|
||||
MongoDB.shutdown(self)
|
||||
|
||||
149
orchestrator/activities/email.py
Normal file
149
orchestrator/activities/email.py
Normal file
@@ -0,0 +1,149 @@
|
||||
from email import encoders
|
||||
from email.mime.base import MIMEBase
|
||||
import traceback
|
||||
from temporalio import workflow, activity
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import smtplib
|
||||
from typing import Any
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
from sientia_do.temporal.utils.logger import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from orchestrator.utils.email_builder import EmailBuilder
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.text import MIMEText
|
||||
|
||||
|
||||
class Email(BaseActivity):
|
||||
def __init__(self, sender_email: str, sender_password: str,
|
||||
smpt_server: str, port: int,
|
||||
logger: Logger, notification_handler: NotificationHandler):
|
||||
|
||||
self.email_builder = EmailBuilder(logger=logger)
|
||||
|
||||
self.sender_email = sender_email
|
||||
self.sender_password = sender_password
|
||||
self.port = port
|
||||
self.logger = logger
|
||||
|
||||
if self.sender_password:
|
||||
self.server = smtplib.SMTP_SSL(smpt_server, port)
|
||||
self.server.login(self.sender_email, self.sender_password)
|
||||
else:
|
||||
self.server = smtplib.SMTP(smpt_server, port)
|
||||
|
||||
BaseActivity.__init__(self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
@activity.defn(name="build_email_html")
|
||||
async def build_email_html(self, input_data: dict[str, Any]) -> str:
|
||||
"""
|
||||
Builds the email html for each receiver group.
|
||||
input_data:
|
||||
- receiver_groups (dict): The receiver groups.
|
||||
- mail_type (str): The mail type.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
return receiver_groups
|
||||
|
||||
def handle_attachments(self, attachments: list[dict], msg: MIMEMultipart) -> MIMEMultipart:
|
||||
"""
|
||||
Attaches a list of attachments to an email message.
|
||||
Args:
|
||||
attachments (List[Dict]): A list of dictionaries where each dictionary contains
|
||||
the keys 'attachment_id', 'trigger', and 'timestamp' representing the attachment details.
|
||||
msg (MIMEMultipart): The email message object to which the attachments will be added.
|
||||
Returns:
|
||||
MIMEMultipart: The email message object with the attachments added.
|
||||
Raises:
|
||||
Exception: If an attachment cannot be added, an error is logged.
|
||||
"""
|
||||
|
||||
for attachment in attachments:
|
||||
att_name = attachment['filename']
|
||||
try:
|
||||
# Create the attachment as a MIMEBase object
|
||||
part = MIMEBase('application', 'octet-stream')
|
||||
part.set_payload(
|
||||
attachment['attachment_content'].encode('utf-8'))
|
||||
encoders.encode_base64(part)
|
||||
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}")
|
||||
|
||||
return msg
|
||||
|
||||
@activity.defn(name="send_email")
|
||||
async def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Sends an email to the receivers of each group.
|
||||
input_data:
|
||||
- receiver_groups (dict): The receiver groups.
|
||||
- mail_type (str): The mail type.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
receiver_groups = input_data['receiver_groups']
|
||||
mail_type = input_data['mail_type']
|
||||
|
||||
self.info(f"Sending email for {mail_type} mail type.",
|
||||
metadata=metadata)
|
||||
|
||||
for group_name, group_config in receiver_groups.items():
|
||||
|
||||
receivers = ", ".join(group_config['members'])
|
||||
|
||||
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 = self.handle_attachments(
|
||||
[notification['attachment_content']
|
||||
for notification in group_config['notifications']
|
||||
if notification['attachment_content']],
|
||||
msg)
|
||||
|
||||
try:
|
||||
self.server.sendmail(
|
||||
self.sender_email, receivers, msg.as_string())
|
||||
except Exception as e:
|
||||
self.error(f"Failed to send email to {group_name}: {e}",
|
||||
metadata=metadata)
|
||||
traceback.print_exc()
|
||||
group_config['status'] = 'failed'
|
||||
else:
|
||||
group_config['status'] = 'sent'
|
||||
|
||||
self.info(f"Email sent to {group_name}.",
|
||||
metadata=metadata)
|
||||
|
||||
self.info(f"Email sent for {mail_type} mail type.",
|
||||
metadata=metadata)
|
||||
|
||||
return receiver_groups
|
||||
@@ -1,4 +1,5 @@
|
||||
|
||||
from pandas import DataFrame
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from orchestrator.utils.orchestrator_functions import minimal_retrain
|
||||
@@ -488,3 +489,47 @@ class Formatters(BaseActivity):
|
||||
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS",
|
||||
attachment=deleted_slots
|
||||
)
|
||||
|
||||
@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.
|
||||
input_data:
|
||||
- receiver_groups (dict): The receiver groups.
|
||||
"""
|
||||
metadata = input_data["metadata"]
|
||||
mail_type = input_data["mail_type"]
|
||||
|
||||
self.info("Formatting log report...", metadata=metadata)
|
||||
|
||||
receiver_groups = input_data['receiver_groups']
|
||||
|
||||
data = {}
|
||||
|
||||
for group_name, group_config in receiver_groups.items():
|
||||
|
||||
notification_id = group_config['notification_id']
|
||||
trigger = group_config['trigger']
|
||||
|
||||
key = f"{notification_id}:{trigger}"
|
||||
|
||||
if key not in data:
|
||||
data[key] = {
|
||||
'status': group_config['status'],
|
||||
'timestamp': group_config['timestamp'],
|
||||
'groups': [group_name],
|
||||
'message': group_config['message'],
|
||||
'level': group_config['level'],
|
||||
'notification_id': notification_id,
|
||||
'block': group_config['block'],
|
||||
'schedule': trigger,
|
||||
'pipeline': group_config['pipeline'],
|
||||
'project': group_config['project'],
|
||||
'model_name': group_config['model_name'],
|
||||
'model_id': group_config['model_id'],
|
||||
'mail_type': mail_type
|
||||
}
|
||||
else:
|
||||
data[key]['groups'].append(group_name)
|
||||
|
||||
return DataFrame(list(data.values())).to_dict()
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from pandas import DataFrame
|
||||
from temporalio import workflow, activity
|
||||
|
||||
|
||||
@@ -80,6 +81,15 @@ class MongoDB(BaseActivity):
|
||||
"""
|
||||
self.shutdown()
|
||||
|
||||
def find(self, collection_name: str, filters: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
collection = self.database[collection_name]
|
||||
|
||||
documents = list(collection.find(filters, {"_id": 0}))
|
||||
|
||||
documents = clear_mongo_id(documents)
|
||||
|
||||
return documents
|
||||
|
||||
@activity.defn(name="find_documents_in_mongodb",)
|
||||
async def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
@@ -106,11 +116,7 @@ class MongoDB(BaseActivity):
|
||||
f"Loading documents from collection '{collection_name}' with filters: {filters}", metadata=metadata)
|
||||
|
||||
try:
|
||||
collection = self.database[collection_name]
|
||||
|
||||
documents = list(collection.find(filters, {"_id": 0}))
|
||||
|
||||
documents = clear_mongo_id(documents)
|
||||
documents = self.find(collection_name, filters)
|
||||
|
||||
self.info(
|
||||
f"Loaded {len(documents)} documents from collection '{collection_name}'", metadata=metadata)
|
||||
@@ -352,3 +358,76 @@ class MongoDB(BaseActivity):
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
@activity.defn(name="load_latest_data")
|
||||
async def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Loads the latest data from MongoDB.
|
||||
input_data:
|
||||
- metadata (dict): The metadata of the workflow.
|
||||
- collection_name (str): The name of the collection to load data from.
|
||||
- last_data_timestamp (str): The timestamp of the last data to load.
|
||||
- base_data_filter (dict): The base data filter to apply to the query.
|
||||
returns:
|
||||
- data (list[dict]): The data loaded from MongoDB.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
collection_name = input_data['collection_name']
|
||||
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
|
||||
)
|
||||
|
||||
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, DEFAULT_DATE_FORMAT)
|
||||
}
|
||||
}
|
||||
|
||||
self.debug(
|
||||
f"Data filter: {data_filter}",
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
data = self.find(collection_name, data_filter)
|
||||
|
||||
self.debug(
|
||||
f"Collected: {data}",
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
for item in data:
|
||||
item['timestamp'] = item['timestamp'].strftime(
|
||||
DEFAULT_DATE_FORMAT)
|
||||
|
||||
self.info(
|
||||
f"Loaded {len(data)} documents from MongoDB",
|
||||
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",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace
|
||||
)
|
||||
raise e
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from pandas import DataFrame
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
@@ -193,3 +194,72 @@ class SlotManager(Redis):
|
||||
f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
|
||||
|
||||
return report
|
||||
|
||||
@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.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
key = "notification_last_timestamp"
|
||||
|
||||
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",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc()
|
||||
)
|
||||
raise e
|
||||
|
||||
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")
|
||||
async def put_last_data_timestamp(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Puts the last data timestamp into redis.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
key = "notification_last_timestamp"
|
||||
|
||||
data = DataFrame(input_data['data'])
|
||||
|
||||
if data.empty:
|
||||
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
|
||||
)
|
||||
|
||||
try:
|
||||
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",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc()
|
||||
)
|
||||
raise e
|
||||
|
||||
return last_data_timestamp
|
||||
|
||||
Reference in New Issue
Block a user