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

SIENTIAPDE-1172

Browse Source 
feat: enhance email and slot manager activities with new filtering and caching functionalities

- Updated Email class to improve attachment handling and logging.
- Refactored Formatters class to process notifications more effectively by grouping them based on their configurations.
- Introduced new methods in SlotManager for filtering notification alerts and storing notification cache.
- Enhanced Alerts workflow to integrate new filtering and caching activities.
- Added tests for new functionalities in SlotManager and Formatters classes to ensure reliability.
This commit is contained in:
vitor-aignosi
2025-07-28 15:07:09 -03:00
parent f2cb0f057d
commit ab657d0fda
10 changed files with 786 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
orchestrator/activities/email.py
View File

@@ -1,9 +1,8 @@
from email import encoders
from email.mime.base import MIMEBase
import traceback
from temporalio import workflow, activity
with workflow.unsafe.imports_passed_through():
import traceback
import smtplib
from typing import Any
from sientia_do.temporal.activities.base import BaseActivity
@@ -12,6 +11,8 @@ with workflow.unsafe.imports_passed_through():
from orchestrator.utils.email_builder import EmailBuilder
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
class Email(BaseActivity):
@@ -24,7 +25,6 @@ class Email(BaseActivity):
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)
@@ -68,7 +68,7 @@ class Email(BaseActivity):
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.
the keys 'filename' and 'content' 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.
@@ -124,7 +124,11 @@ class Email(BaseActivity):
msg['Subject'] = f"SIENTIA™ {mail_type}"
msg = self.handle_attachments(
[notification['attachment_content']
[
{
"filename": f"{notification['trigger']}_{notification['notification_id']}.txt",
"content": notification['attachment_content']
}
for notification in group_config['notifications']
if notification['attachment_content']],
msg)

44
orchestrator/activities/formatters.py
View File

@@ -508,28 +508,30 @@ class Formatters(BaseActivity):
for group_name, group_config in receiver_groups.items():
notification_id = group_config['notification_id']
trigger = group_config['trigger']
for notification in group_config['notifications']:
key = f"{notification_id}:{trigger}"
notification_id = notification['notification_id']
trigger = notification['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)
key = f"{notification_id}:{trigger}"
if key not in data:
data[key] = {
'status': group_config['status'],
'timestamp': notification['timestamp'],
'groups': [group_name],
'message': notification['message'],
'level': notification['level'],
'notification_id': notification_id,
'block': notification['block'],
'schedule': trigger,
'pipeline': notification['pipeline'],
'project': notification['project'],
'model_name': notification['model_name'],
'model_id': notification['model_id'],
'mail_type': mail_type
}
else:
data[key]['groups'].append(group_name)
return DataFrame(list(data.values())).to_dict()

76
orchestrator/activities/slot_manager.py
View File

@@ -1,4 +1,4 @@
from pandas import DataFrame
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
@@ -9,6 +9,9 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.redis_base import Redis
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT
from pandas import DataFrame
from datetime import datetime, timedelta
class SlotManager(Redis):
@@ -263,3 +266,74 @@ class SlotManager(Redis):
raise e
return last_data_timestamp
@activity.defn(name="filter_notification_alerts")
async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Filter notification alerts
"""
metadata = input_data['metadata']
notification_package = input_data['notification_package']
sending_configs = input_data['sending_configs']
notification_ttl = input_data['notification_ttl']
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]['notifications'] = []
ignore_list = receiver_group.get('ignore', [])
for notification in notification_package:
alert_type = "do_nothing"
notification_id = notification['notification_id']
# Check if notification was recently sent
key = f"{notification['trigger']}:{notification_id}"
last_sent = self.get(key)
if last_sent is None:
alert_type = "core_alerts"
else:
last_sent = datetime.strptime(
last_sent, DEFAULT_DATE_FORMAT)
# Check if "notification_ttl" seconds has passed since last sent
if (datetime.now() - last_sent) > timedelta(seconds=notification_ttl):
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:
receiver_groups[group_name]["notifications"].append(
notification)
return receiver_groups
@activity.defn(name="store_notification_cache")
async def store_notification_cache(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Store notification cache
"""
metadata = input_data['metadata']
log_report = DataFrame(input_data['log_report'])
sent_ttl = input_data['sent_ttl']
self.info("Storing notification cache...", metadata=metadata)
now = datetime.now().strftime(DEFAULT_DATE_FORMAT)
for index, row in log_report.iterrows():
status = row['status']
if status == 'sent':
key = f"{row['schedule']}:{row['notification_id']}"
self.set(key, now, ttl=sent_ttl)
return log_report

70
orchestrator/workflows/alerts.py
View File

@@ -9,14 +9,80 @@ with workflow.unsafe.imports_passed_through():
class Alerts:
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Workflow to send alerts to the users
Args:
input_data (dict[str, Any]): Input data. It contains the following keys:
- schedule_name: str - Name of the schedule
- notification_ttl: int - Period before consider some notification persistent
- sent_ttl: int - Time to live for the sent notification
Returns:
None
Raises:
Exception: If the workflow fails
"""
metadata = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'workflow_name': 'alerts',
'model_name': '-',
'model_id': '-'
}
}
mail_type = "Alerts"
input_data['metadata'] = metadata
input_data['base_data_filter'] = {
'level': 'ERROR'
}
# Call subworkflow "load_notification_package" passing the static filters
# (level = "ERROR" and timestamp > last timestamp)
package = await workflow.execute_child_workflow(
'load_notification_package',
input_data
)
if not package['notification_package'] or not package['sending_configs']:
return
# Filter notification package by groups custom configs, levels and
# timestamp cached
receiver_groups = await workflow.execute_local_activity_method(
Activities.filter_notification_alerts,
{
**metadata,
'notification_package': package['notification_package'],
'sending_configs': package['sending_configs'],
'notification_ttl': input_data['notification_ttl']
}
)
# Call subworkflow "process_notifications" passing the notification package
log_report = await workflow.execute_child_workflow(
'process_notifications',
{
'metadata': metadata,
'mail_type': mail_type,
'notification_package': receiver_groups,
'schema': 'sientia_data',
'table_name': 'log_report'
}
)
# Store the notification_id sendings to avoid sending them again
pass
await workflow.execute_activity_method(
Activities.store_notification_cache,
{
**metadata,
'log_report': log_report,
'sent_ttl': input_data['sent_ttl']
}
)

11
orchestrator/workflows/subworkflows/load_notification_package.py
View File

@@ -62,14 +62,19 @@ class LoadNotificationPackage:
**metadata,
'collection_name': 'notification_queue',
'last_data_timestamp': last_timestamp,
'base_data_filter': {
'level': 'ERROR'
}
'base_data_filter': input_data['base_data_filter']
},
start_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy
)
if not notification_package:
return {
'last_timestamp': last_timestamp,
'notification_package': [],
'sending_configs': await sending_configs_handler
}
# Put last collected timestamp in redis "notification_last_timestamp"
await workflow.start_activity_method(

277
tests/orchestrator/activities/test_email.py Normal file
View File

@@ -0,0 +1,277 @@
from unittest.mock import MagicMock, call, patch
from pytest import mark, fixture
from orchestrator.activities.email import Email
@fixture
@patch('orchestrator.activities.email.EmailBuilder')
@patch('orchestrator.activities.email.smtplib')
def email(smtplib, email_builder):
email = Email(
sender_email="test@test.com",
sender_password="test",
smpt_server="test",
port=587,
logger=MagicMock(),
notification_handler=MagicMock()
)
email_builder.send_notification = MagicMock()
return email
@patch('orchestrator.activities.email.EmailBuilder')
@patch('orchestrator.activities.email.smtplib')
def test___init___with_password(smtplib, email_builder):
email = Email(
sender_email="test@test.com",
sender_password="test",
smpt_server="test",
port=587,
logger=MagicMock(),
notification_handler=MagicMock()
)
assert email.sender_email == "test@test.com"
assert email.sender_password == "test"
assert email.port == 587
smtplib.SMTP_SSL.assert_called_once_with("test", 587)
smtplib.SMTP_SSL.return_value.login.assert_called_once_with(
"test@test.com", "test")
assert email.server == smtplib.SMTP_SSL.return_value
@patch('orchestrator.activities.email.EmailBuilder')
@patch('orchestrator.activities.email.smtplib')
def test___init___without_password(smtplib, email_builder):
email = Email(
sender_email="test@test.com",
sender_password=None,
smpt_server="test",
port=587,
logger=MagicMock(),
notification_handler=MagicMock()
)
assert email.sender_email == "test@test.com"
assert email.sender_password is None
assert email.port == 587
smtplib.SMTP.assert_called_once_with("test", 587)
assert email.server == smtplib.SMTP.return_value
metadata = {
"metadata": {
"schedule_name": "test",
"model_name": "test",
"model_id": "test",
"workflow_name": "test",
}
}
@mark.asyncio
async def test_build_email_html(email):
email.email_builder.build_email = MagicMock(
return_value="test"
)
input_data = {
**metadata,
"receiver_groups": {
"group_1": {
"notifications": [
{
"type": "test",
"subject": "test",
"body": "test"
},
{
"type": "test",
"subject": "test",
"body": "test"
}
]
}
},
"mail_type": "test"
}
response = await email.build_email_html(input_data)
assert response == {
"group_1": {
"notifications": [
{
"type": "test",
"subject": "test",
"body": "test",
},
{
"type": "test",
"subject": "test",
"body": "test",
}
],
"html": "test"
}
}
email.email_builder.build_email.assert_called_once_with(
input_data['receiver_groups']['group_1']['notifications'],
input_data['mail_type']
)
@patch('orchestrator.activities.email.MIMEBase')
@patch('orchestrator.activities.email.encoders')
def test_handle_attachments_success(encoders, mime_base, email):
message = MagicMock()
attachments = [
{
"filename": "file_1",
"attachment_content": "test_content_1"
},
{
"filename": "file_2",
"attachment_content": "test_content_2"
},
{
"filename": "file_3",
"attachment_content": "test_content_3"
}
]
response = email.handle_attachments(attachments, message)
assert response == message
mime_base.assert_called_with('application', 'octet-stream')
assert mime_base.return_value.set_payload.call_count == 3
mime_base.return_value.set_payload.assert_has_calls(
[
call("test_content_1".encode('utf-8')),
call("test_content_2".encode('utf-8')),
call("test_content_3".encode('utf-8'))
]
)
encoders.encode_base64.assert_called_with(mime_base.return_value)
assert encoders.encode_base64.call_count == 3
mime_base.return_value.add_header.assert_has_calls(
[
call('Content-Disposition', 'attachment; filename="file_1"'),
call('Content-Disposition', 'attachment; filename="file_2"'),
call('Content-Disposition', 'attachment; filename="file_3"')
]
)
message.attach.assert_called_with(mime_base.return_value)
assert message.attach.call_count == 3
@patch('orchestrator.activities.email.MIMEBase')
def test_handle_attachments_failure(mime_base, email):
message = MagicMock()
mime_base.side_effect = Exception("test")
attachments = [
{
"filename": "file_1",
"attachment_content": "test_content_1"
}
]
email.handle_attachments(attachments, message)
assert message.attach.call_count == 0
@mark.asyncio
@patch('orchestrator.activities.email.MIMEText')
@patch('orchestrator.activities.email.MIMEMultipart')
async def test_send_email(mimemultipart, mimetext, email):
side_effect_1 = MagicMock()
side_effect_2 = MagicMock()
mimemultipart.side_effect = [
side_effect_1,
side_effect_2
]
email.email_builder.send_notification = MagicMock()
email.server.sendmail = MagicMock(
side_effect=[
None,
Exception("test")
]
)
input_data = {
**metadata,
"receiver_groups": {
"group_1": {
"members": ["test@test.com", "test2@test.com"],
"notifications": [
{
"attachment_content": "test_content_1",
"trigger": "test_trigger",
"notification_id": "test_notification_id"
}
],
"html": "test_html1"
},
"group_2": {
"members": ["test3@test.com", "test4@test.com"],
"notifications": [],
"html": "test_html2"
}
},
"mail_type": "test_TYPE"
}
response = await email.send_email(input_data)
assert response['group_1']['status'] == 'sent'
assert response['group_2']['status'] == 'failed'
assert mimemultipart.call_count == 2
mimetext.assert_has_calls(
[
call("test_html1", "html"),
call("test_html2", "html")
]
)
side_effect_1.__setitem__.assert_has_calls(
[
call('From', 'test@test.com'),
call('To', 'test@test.com, test2@test.com'),
call('Subject', 'SIENTIA™ test_TYPE')
]
)
side_effect_2.__setitem__.assert_has_calls(
[
call('From', 'test@test.com'),
call('To', 'test3@test.com, test4@test.com'),
call('Subject', 'SIENTIA™ test_TYPE')
]
)
email.server.sendmail.assert_has_calls(
[
call('test@test.com', 'test@test.com, test2@test.com',
side_effect_1.as_string.return_value),
call('test@test.com', 'test3@test.com, test4@test.com',
side_effect_2.as_string.return_value)
]
)

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

@@ -1,5 +1,6 @@
from unittest.mock import MagicMock, patch, call, ANY
import json
from pandas import DataFrame
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.formatters import Formatters
@@ -696,3 +697,71 @@ async def test_report_slot_orchestration(formatters):
attachment=input_data['deleted_slots']
)
])
@mark.asyncio
async def test_format_log_report(formatters):
input_data = {
**metadata,
"receiver_groups": {
"test_receiver_group": {
"notifications": [
{
"notification_id": "test_notification_id",
"trigger": "test_trigger",
"timestamp": "2021-01-01",
"message": "test_message",
"level": "test_level",
"block": "test_block",
"pipeline": "test_pipeline",
"project": "test_project",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
],
"status": "sent",
},
"test_receiver_group2": {
"notifications": [
{
"notification_id": "test_notification_id",
"trigger": "test_trigger",
"timestamp": "2021-01-01",
"message": "test_message",
"level": "test_level",
"block": "test_block",
"pipeline": "test_pipeline",
"project": "test_project",
"model_name": "test_model_name",
"model_id": "test_model_id"
}
],
"status": "sent",
}
},
"mail_type": "test_mail_type"
}
result = await formatters.format_log_report(input_data)
expected_result = DataFrame(
[
{
"status": "sent",
"timestamp": "2021-01-01",
"groups": ["test_receiver_group", "test_receiver_group2"],
"message": "test_message",
"level": "test_level",
"notification_id": "test_notification_id",
"block": "test_block",
"schedule": "test_trigger",
"pipeline": "test_pipeline",
"project": "test_project",
"model_name": "test_model_name",
"model_id": "test_model_id",
"mail_type": "test_mail_type"
}
]
)
assert DataFrame(result).equals(expected_result)

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

@@ -1,8 +1,10 @@
from unittest.mock import MagicMock, patch, call, ANY
from datetime import datetime, timedelta
from pandas import DataFrame
from pytest import mark, fixture
from orchestrator.activities.slot_manager import SlotManager
from sientia_do.notifications.models import NotificationLevel
from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT
metadata = {
"metadata": {
@@ -359,3 +361,98 @@ async def test_put_last_data_timestamp_error(slot_manager):
else:
assert False, "Expected exception"
@mark.asyncio
async def test_filter_notification_alerts(slot_manager):
slot_manager.get = MagicMock(side_effect=[
None,
(datetime.now() - timedelta(seconds=600)
).strftime(DEFAULT_DATE_FORMAT),
datetime.now().strftime(DEFAULT_DATE_FORMAT),
None,
(datetime.now() - timedelta(seconds=600)
).strftime(DEFAULT_DATE_FORMAT),
datetime.now().strftime(DEFAULT_DATE_FORMAT)])
input_data = {
**metadata,
'notification_package': [
{
'trigger': 'test_trigger_1',
'notification_id': 'test_notification_id_1'
},
{
'trigger': 'test_trigger_2',
'notification_id': 'test_notification_id_2'
},
{
'trigger': 'test_trigger_3',
'notification_id': 'test_notification_id_3'
}
],
'sending_configs': [
{
'group_name': 'test_group_1',
'contents': ['core_alerts', 'persistent_alerts']
},
{
'group_name': 'test_group_2',
'contents': ['core_alerts']
}
],
'notification_ttl': 300
}
response = await slot_manager.filter_notification_alerts(input_data)
assert response == {
'test_group_1': {
'group_name': 'test_group_1',
'contents': ['core_alerts', 'persistent_alerts'],
'notifications': [
{
'trigger': 'test_trigger_1',
'notification_id': 'test_notification_id_1'
},
{
'trigger': 'test_trigger_2',
'notification_id': 'test_notification_id_2'
}
]
},
'test_group_2': {
'group_name': 'test_group_2',
'contents': ['core_alerts'],
'notifications': [
{
'trigger': 'test_trigger_1',
'notification_id': 'test_notification_id_1'
}
]
}
}
@mark.asyncio
async def test_store_notification_cache(slot_manager):
"""Test store_notification_cache"""
test_data = {
**metadata,
'log_report': DataFrame({
'status': ['sent', 'error'],
'schedule': ['test_schedule_1', 'test_schedule_2'],
'notification_id': ['test_notification_id_1', 'test_notification_id_2']
}).to_dict(),
'sent_ttl': 600
}
slot_manager.set = MagicMock()
await slot_manager.store_notification_cache(test_data)
slot_manager.set.assert_called_once_with(
"test_schedule_1:test_notification_id_1",
ANY,
ttl=600
)

79
tests/orchestrator/workflows/subworkflows/test_process_notifications.py Normal file
View File

@@ -0,0 +1,79 @@
from unittest.mock import AsyncMock, patch, ANY, call
from pytest import fixture, mark
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
from orchestrator.activities.activities import Activities
@fixture
def process_notifications():
return ProcessNotifications()
metadata = {
'metadata': {
'schedule_name': 'test-schedule-name',
'workflow_name': 'test-workflow',
'model_name': '-',
'model_id': '-',
}
}
@mark.asyncio
@patch("orchestrator.workflows.subworkflows.process_notifications.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, process_notifications):
input_data = {
'metadata': metadata,
'notification_package': ["content"],
'mail_type': 'test_mail_type',
'schema': 'test_schema',
'table_name': 'test_table_name',
}
response = await process_notifications.run(input_data)
assert response == workflow_mock.execute_local_activity_method.return_value
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.build_email_html,
{
**metadata,
'receiver_groups': input_data['notification_package']
},
schedule_to_close_timeout=ANY,
retry_policy=ANY
),
call(
Activities.format_log_report,
{
**metadata,
'receiver_groups': workflow_mock.execute_activity_method.return_value,
'mail_type': input_data['mail_type']
},
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.send_email,
{
**metadata,
'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
'mail_type': input_data['mail_type']
},
schedule_to_close_timeout=ANY,
retry_policy=ANY
),
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_local_activity_method.return_value
},
schedule_to_close_timeout=ANY,
retry_policy=ANY
)
])

80
tests/orchestrator/workflows/test_alerts.py Normal file
View File

@@ -0,0 +1,80 @@
from unittest.mock import AsyncMock, patch, ANY, call
from pytest import fixture, mark
from orchestrator.workflows.alerts import Alerts
from orchestrator.activities.activities import Activities
@fixture
def alerts():
return Alerts()
metadata = {
'metadata': {
'schedule_name': 'test-schedule-name',
'workflow_name': 'alerts',
'model_name': '-',
'model_id': '-',
}
}
@mark.asyncio
@patch("orchestrator.workflows.alerts.workflow", new_callable=AsyncMock)
async def test_run_full_flow(workflow_mock, alerts):
input_data = {
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await alerts.run(input_data)
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'load_notification_package',
{
**input_data,
'metadata': metadata,
'base_data_filter': {
'level': 'ERROR'
}
}
)
])
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'process_notifications',
{
'metadata': metadata,
'mail_type': 'Alerts',
'notification_package': workflow_mock.execute_local_activity_method.return_value,
'schema': 'sientia_data',
'table_name': 'log_report'
}
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.filter_notification_alerts,
{
**metadata,
'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'],
'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'],
'notification_ttl': input_data['notification_ttl']
}
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.store_notification_cache,
{
**metadata,
'log_report': workflow_mock.execute_child_workflow.return_value,
'sent_ttl': input_data['sent_ttl']
}
)
])