SIENTIAPDE-1163

Update notification handler references to use CoreNotificationHandler and add MongoDB configuration to connectors. Update GITHUB_BRANCH in values.yaml for optimization tasks.

laborious/activities/activities.py

@@ -2,7 +2,7 @@ from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.activities.postgres import Postgres
-    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.temporal.utils.logger import Logger
     from laborious.activities.mlflow import MLFlow
     from laborious.activities.gates import Gates

laborious/activities/gates.py

@@ -3,7 +3,7 @@ from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
     import traceback
-    from sientia_do.notifications.handlers import NotificationHandler
+    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.utils.logger import Logger
@@ -64,9 +64,6 @@ class Gates(BaseActivity):
             tuple[str | None, int, str]: (policy, confidence, comments) based in priority
                 list and filter configuration and functions.
         """
-
-        self.debug(f"Input data: {input_data}")
-
         metadata = input_data['metadata']
 
         self.debug("Performing input gate...", metadata)
@@ -93,7 +90,8 @@ class Gates(BaseActivity):
                     filter_output.append(config['policy'])
             except Exception as e:
                 trace = traceback.format_exc()
-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                     notification_id=f"INTPUT_GATE_ERROR__{fil}",
                     message=f"Error in filter {fil}:{config}: \n {e}",
                     block="input_gate",
@@ -148,7 +146,8 @@ class Gates(BaseActivity):
                 if mlflow_response_filter_functions[fil](data, config):
                     filter_output.append(config['policy'])
                     comments.append(data['content']['message'])
-                    self.notification_handler.build_and_send_notification(
+                    self.send_notification(
+                        metadata=metadata,
                         notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
                         message=data['content']['message'],
                         block="mlflow_gate",
@@ -157,7 +156,8 @@ class Gates(BaseActivity):
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                     notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
                     message=f"Error in filter {fil}:{config}: \n {e}",
                     block="mlflow_gate",
@@ -211,7 +211,8 @@ class Gates(BaseActivity):
             try:
                 if mlflow_content_filter_functions[fil](data, config):
                     filter_output.append(config['policy'])
-                    self.notification_handler.build_and_send_notification(
+                    self.send_notification(
+                        metadata=metadata,
                         notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
                         message=f"Data not passed the content filter {fil}:{config}",
                         block="mlflow_gate",
@@ -220,7 +221,8 @@ class Gates(BaseActivity):
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                     notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
                     message=f"Error in filter {fil}:{config}: \n {e}",
                     block="mlflow_gate",
@@ -260,7 +262,7 @@ class Gates(BaseActivity):
         data['prediction_confidence'] = input_data['prediction_confidence']
         data['prediction_status'] = 'Good'
         data['comments'] = ""
-        data.sort_values(by='timestamp', inplace=True)
+        data = data.sort_values(by='timestamp')
 
         return data.to_dict()
 

laborious/activities/mlflow.py

@@ -5,7 +5,7 @@ from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.activities.base import BaseActivity
-    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.temporal.utils.logger import Logger
     from laborious.utils.repository.model_repository import MLFlowRepository
     from typing import Any

laborious/activities/opc.py

@@ -2,7 +2,7 @@ from temporalio import activity, workflow
 
 
 with workflow.unsafe.imports_passed_through():
-    from sientia_do.notifications.handlers import NotificationHandler
+    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.utils.logger import Logger
@@ -40,7 +40,7 @@ class OPC(BaseActivity):
         BaseActivity.__init__(self, logger, notification_handler)
 
     def write_data(self, server_id: str, tag: str, data: Any,
-                   data_type: str, tag_type: str) -> bool:
+                   data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool:
         """
         Write data to OPC server.
 
@@ -57,11 +57,11 @@ class OPC(BaseActivity):
 
         try:
             return self.opc_repository[server_id].write_data(
-                tag, data, data_type)
-            self.logger.debug(f"Wrote {tag_type} to {tag}")
+                tag, data, data_type, self.logger, metadata)
         except Exception as e:
             trace = traceback.format_exc()
-            self.notification_handler.build_and_send_notification(
+            self.send_notification(
+                metadata=metadata,
                 notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
                 message=f"Error writing data to OPC server: {e}",
                 block="write_opc_data",
@@ -109,7 +109,8 @@ class OPC(BaseActivity):
                         tag=tag,
                         data=data.head(1)['prediction'].values[0],
                         data_type=tag_config['data_type'],
-                        tag_type='prediction'
+                        tag_type='prediction',
+                        metadata=metadata
                     )
 
             if 'confidence_tags' in config:
@@ -119,7 +120,8 @@ class OPC(BaseActivity):
                         tag=tag,
                         data=data.head(1)['prediction_confidence'].values[0],
                         data_type=tag_config['data_type'],
-                        tag_type='confidence'
+                        tag_type='confidence',
+                        metadata=metadata
                     )
 
         return self.process_confidence(data, success, metadata)
@@ -142,8 +144,8 @@ class OPC(BaseActivity):
         if not success:
             data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
             self.debug(
-                "Some data could not be written to OPC servers, setting confidence to "
-                f"{OPC_WRITTING_ERROR_CONFIDENCE}."
+                f"Some data could not be written to OPC servers, setting confidence to  {OPC_WRITTING_ERROR_CONFIDENCE}.",
+                metadata
             )
 
         else:

laborious/utils/connectors_config.py

@@ -34,7 +34,7 @@ def build_opc_config():
         return json.loads(opc_raw)
 
     return {
-        'opc': {
+        getenv('OPC_ID', '1'): {
             'id': getenv('OPC_ID', '1'),
             'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'),
             'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
@@ -44,3 +44,15 @@ def build_opc_config():
             'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120'))
         }
     }
+
+
+def build_mongodb_config():
+    username = getenv('MONGODB_USERNAME', 'sientia')
+    password = getenv('MONGODB_PASSWORD', 'sientia')
+    uri = getenv('MONGODB_URL', 'localhost:27017')
+
+    connection_string = f'mongodb://{username}:{password}@{uri}'
+    return {
+        'connection_string': connection_string,
+        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia')
+    }

laborious/utils/repository/opc_repository.py

@@ -1,14 +1,14 @@
 import traceback
-from logging import Logger
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 from asyncua.sync import Client
 from asyncua.crypto.security_policies import SecurityPolicyBasic256
-from asyncua.ua import DataValue, Variant, VariantType
+from asyncua.ua import DataValue, Variant, VariantType, DateTime
 from regex import F
-from sientia_do.notifications.handlers import NotificationHandler
+from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
 from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.utils.logger import Logger
 
 data_type_map = {
     'float': {
@@ -104,7 +104,7 @@ class OpcRepository():
         self.client = Client(self.url)
         if self.cert_path:
             self.set_security()
-        self.logger.info('Starting connection...')
+        self.logger.info(f'Starting connection to OPC server {self.id}...')
         return self.try_connect()
 
     def try_connect(self):
@@ -114,6 +114,7 @@ class OpcRepository():
         Returns:
             bool: True if the connection was successful, False otherwise.
         """
+
         try:
             self.last_reconnection_time = datetime.now()
             self.client.connect()
@@ -136,9 +137,12 @@ class OpcRepository():
         """
         if self.client is None:
             return
-        self.client.disconnect()
+        try:
+            self.client.disconnect()
+            self.logger.info('Disconnected from OPC server')
+        except Exception as e:
+            self.logger.error(f"Failed to disconnect from OPC server: {e}")
         self.client = None
-        self.logger.info('Disconnected from OPC server')
 
     def __del__(self):
         """
@@ -182,17 +186,19 @@ class OpcRepository():
 
             self.logger.error(
                 f"OPC server {self.id} is not connected")
-            if (datetime.now() - self.last_reconnection_time).total_seconds(
+            if self.last_reconnection_time is None or (datetime.now() - self.last_reconnection_time).total_seconds(
             ) > self.reconnection_interval:
+                self.disconnect()
                 self.logger.error(
                     f"Trying to reconnect to OPC server {self.id}...")
-                return self.try_connect()
+                return self.connect()
 
             return False
 
         return True
 
-    def write_data(self, node: str, value: Any, data_type: str) -> bool:
+    def write_data(self, node: str, value: Any, data_type: str,
+                   logger: Logger, metadata: dict[str, Any]) -> bool:
         """
         Writes data to the OPC server.
         If the connection is not established, it attempts to reconnect.
@@ -211,28 +217,40 @@ class OpcRepository():
             trace = traceback.format_exc()
             self.notification_handler.build_and_send_notification(
                 notification_id=f"OPC_WRITE_GET_NODE_ERROR_{self.id}",
-                message=f"Failed to get node from OPC server: {e}",
+                message=f"Failed to get node from OPC server: {e} | metadata: {metadata}",
                 block="opc_repository",
                 level=NotificationLevel.ERROR,
                 attachment_content=trace
             )
-            self.logger.error(trace)
+            logger.custom_error(trace, metadata.get('schedule_name', 'N/A'))
             self.error_count += 1
             return False
 
         if data_type not in data_type_map:
             self.notification_handler.build_and_send_notification(
                 notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}",
-                message=f"Unsupported data type: {data_type}",
+                message=f"Unsupported data type: {data_type} | metadata: {metadata}",
                 block="opc_repository",
                 level=NotificationLevel.ERROR
             )
             return False
 
         data = data_type_map[data_type]['converter'](value)
-        self.logger.info(f'Writing {data} - {type(data)} to {node}')
+        logger.custom_info(
+            f'Writing {data} - {type(data)} to {node}', metadata.get('schedule_name', 'N/A'))
+        now = datetime.now()
         ua_data = DataValue(
-            Variant(data, data_type_map[data_type]['opc_type']))
+            Variant(data, data_type_map[data_type]['opc_type']),
+            SourceTimestamp=DateTime(
+                now.year,
+                now.month,
+                now.day,
+                now.hour,
+                now.minute,
+                now.second,
+                now.microsecond
+            )
+        )
 
         try:
             node.write_value(ua_data)
@@ -240,12 +258,12 @@ class OpcRepository():
             trace = traceback.format_exc()
             self.notification_handler.build_and_send_notification(
                 notification_id=f"OPC_WRITE_DATA_ERROR_{self.id}",
-                message=f"Failed to write data to OPC server: {e}",
+                message=f"Failed to write data to OPC server: {e} | metadata: {metadata}",
                 block="opc_repository",
                 level=NotificationLevel.ERROR,
                 attachment_content=trace
             )
-            self.logger.error(trace)
+            logger.custom_error(trace, metadata.get('schedule_name', 'N/A'))
             self.error_count += 1
             return False
         self.error_count = 0

laborious/worker/worker.py

@@ -14,9 +14,10 @@ with workflow.unsafe.imports_passed_through():
     from laborious.utils.connectors_config import (
         build_postgres_config,
         build_mlflow_config,
-        build_opc_config
+        build_opc_config,
+        build_mongodb_config
     )
-    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.temporal.utils.logger import get_logger
 
 
@@ -28,10 +29,12 @@ async def main():
 
     logger.info('Starting Notification Handler...')
 
+    mongo_config = build_mongodb_config()
     notification_handler = NotificationHandler(
-        servers=os.getenv('KAFKA_BOOTSTRAP_SERVERS', 'http://localhost:9092'),
+        connection_string=mongo_config['connection_string'],
+        database=mongo_config['database_name'],
         logger=logger,
-        project_name=os.getenv('PROJECT_NAME', 'laborious'),
+        project_name=os.getenv('PROJECT_NAME', 'laborious')
     )
 
     logger.info('Starting Activities...')
@@ -76,7 +79,12 @@ async def main():
                 activities.load_custom_query,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres
-            ]
+            ],
+            max_concurrent_workflow_tasks=100,
+            max_concurrent_activities=100,
+            max_concurrent_local_activities=100,
+            max_concurrent_workflow_task_polls=100,
+            max_cached_workflows=50,
         )
     ]
 
@@ -90,7 +98,7 @@ async def main():
         # This will run the workers and wait for them to complete.
         # If an exception occurs in any of the worker handlers, it will be propagated here.
         await asyncio.gather(*handlers)
-    except BaseException as e:
+    except BaseException as e:  # NOSONAR
         logger.error(f"An unhandled exception occurred: {e}")
     finally:
         if notification_handler:

requirements.txt

@@ -3,5 +3,5 @@ psycopg2-binary
 sqlalchemy
 asyncua
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.0
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.1
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.1

tests/laborious/activities/test_activities.py

@@ -90,65 +90,6 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
     )
 
 
-@mark.asyncio
-@patch('laborious.activities.activities.Postgres.__init__')
-@patch('laborious.activities.activities.MLFlow.__init__')
-@patch('laborious.activities.activities.OPC.__init__')
-async def test_prepare_activity(_mock_opc_init,
-                                _mock_mlflow_init, _mock_postgres_init):
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10
-    }
-
-    mlflow_config = {
-        'host': 'localhost',
-        'port': 5000,
-        'username': 'mlflow',
-        'password': 'mlflow'
-    }
-
-    opc_config = {
-        'bootstrap_servers': 'localhost:9092',
-        'polling_time': 1000,
-        'group_id': 'test-group'
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        opc_config=opc_config,
-        logger=logger,
-        notification_handler=notification_handler
-    )
-
-    input_data = {
-        'workflow_name': 'test-workflow-name',
-        'schedule_name': 'test-schedule-name',
-        'model_name': 'test-model-name',
-        'model_id': 'test-model-id'
-    }
-
-    await activities.prepare_activity(input_data)
-
-    assert activities.notification_handler.base_notification.pipeline == input_data[
-        'workflow_name']
-    assert activities.notification_handler.base_notification.trigger == input_data[
-        'schedule_name']
-    assert activities.notification_handler.base_notification.model_name == input_data[
-        'model_name']
-    assert activities.notification_handler.base_notification.model_id == input_data[
-        'model_id']
-
-
 @patch('laborious.activities.activities.Postgres', return_value=MagicMock())
 @patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
 @patch('laborious.activities.activities.OPC', return_value=MagicMock())

tests/laborious/activities/test_gates.py

@@ -6,16 +6,34 @@ from laborious.activities.gates import Gates
 
 @fixture
 def gates_activity():
-    return Gates(
+    gates = Gates(
         logger=MagicMock(),
         notification_handler=MagicMock(),
     )
+    gates.error = MagicMock()
+    gates.debug = MagicMock()
+    gates.info = MagicMock()
+    gates.warning = MagicMock()
+    gates.critical = MagicMock()
+    gates.send_notification = MagicMock()
+    return gates
+
+
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
 
 
 @mark.asyncio
 async def test_input_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'INVALID_FILTER': {'POLICY': 'STOP'}
         },
@@ -28,8 +46,8 @@ async def test_input_gate_invalid_filter(gates_activity):
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.logger.error.assert_called_once_with(
-        "Filter INVALID_FILTER not found"
+    gates_activity.error.assert_called_once_with(
+        "Filter INVALID_FILTER not found", metadata['metadata']
     )
 
 
@@ -41,6 +59,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
     mock_input_filter_functions.__getitem__.return_value = MagicMock(
         side_effect=Exception("Test error"))
     input_data = {
+        **metadata,
         'filters': {
             'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
         },
@@ -53,7 +72,8 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
+    gates_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
         notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
         message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
         block="input_gate",
@@ -66,6 +86,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
 async def test_input_gate_no_filters(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {},
         'data': {'value': [1, 2, 3]},
         'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
@@ -76,13 +97,14 @@ async def test_input_gate_no_filters(gates_activity):
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_input_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
         },
@@ -95,13 +117,14 @@ async def test_input_gate_with_filter(gates_activity):
 
     # Assert
     assert result == ('STOP', -1, "Input data with bad quality")
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_mlflow_response_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'INVALID_FILTER': {'POLICY': 'STOP'}
         },
@@ -126,6 +149,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
     mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
         side_effect=Exception("Test error"))
     input_data = {
+        **metadata,
         'filters': {
             'INVALID_FILTER': {'POLICY': 'STOP'}
         },
@@ -139,7 +163,8 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
+    gates_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
         notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
         message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
         block="mlflow_gate",
@@ -152,6 +177,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
 async def test_mlflow_response_gate_no_filters(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {},
         'data': {'content': {'message': 'success'}},
         'type': 'test',
@@ -163,13 +189,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_mlflow_response_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'API_ERROR': {'policy': 'STOP'}
         },
@@ -189,14 +216,15 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
 
     # Assert
     assert result == ('STOP', -1, "API error occurred")
-    gates_activity.logger.debug.assert_called()
-    gates_activity.notification_handler.build_and_send_notification.assert_called()
+    gates_activity.debug.assert_called()
+    gates_activity.send_notification.assert_called()
 
 
 @mark.asyncio
 async def test_mlflow_content_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'INVALID_FILTER': {'POLICY': 'STOP'}
         },
@@ -221,6 +249,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
     mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
         side_effect=Exception("Test error"))
     input_data = {
+        **metadata,
         'filters': {
             'API_ERROR': {'POLICY': 'STOP'}
         },
@@ -240,8 +269,9 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.logger.debug.assert_called()
-    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
+    gates_activity.debug.assert_called()
+    gates_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
         notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
         message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
         block="mlflow_gate",
@@ -254,6 +284,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
 async def test_mlflow_content_gate_no_filters(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {},
         'data': {'value': [1, 2, 3]},
         'type': 'test',
@@ -265,13 +296,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
 
     # Assert
     assert result == (None, 0, "")
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_mlflow_content_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'filters': {
             'NAN_VALUES': {'policy': 'STOP', 'config': {}}
         },
@@ -286,14 +318,15 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
     # Assert
     assert result == (
         'STOP', -1, "Transformed data not passed the content filter")
-    gates_activity.logger.debug.assert_called()
-    gates_activity.notification_handler.build_and_send_notification.assert_called()
+    gates_activity.debug.assert_called()
+    gates_activity.send_notification.assert_called()
 
 
 @mark.asyncio
 async def test_format_prediction(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'data': {'prediction': [1], 'response_time': [0.1]},
         'timestamp': '2023-05-26 11:12:27',
         'model_id': 'test_model',
@@ -311,13 +344,14 @@ async def test_format_prediction(gates_activity):
     assert result['prediction_confidence'] == {0: 0.9}
     assert result['prediction_status'] == {0: 'Good'}
     assert result['comments'] == {0: ""}
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_format_default_prediction(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'timestamp': '2023-05-26 11:12:27',
         'model_id': 'test_model',
         'prediction_confidence': 0.1,
@@ -335,13 +369,14 @@ async def test_format_default_prediction(gates_activity):
     assert result['prediction_confidence'] == {0: 0.1}
     assert result['prediction_status'] == {0: 'Bad'}
     assert result['comments'] == {0: 'Test comment'}
-    gates_activity.logger.debug.assert_called()
+    gates_activity.debug.assert_called()
 
 
 @mark.asyncio
 async def test_get_last_timestamp_with_data(gates_activity):
     # Arrange
     input_data = {
+        **metadata,
         'data': {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
         }

tests/laborious/activities/test_mlflow.py

@@ -39,6 +39,16 @@ def mlflow(mock_mlflow_repository):
     )
 
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
+
+
 @mark.asyncio
 @patch("laborious.activities.mlflow.DataFrame")
 @patch("laborious.activities.mlflow.max")
@@ -46,6 +56,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
+        **metadata,
         'data': [
             {'timestamp': '2024-01-01', 'variable': 'var1',
                 'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
@@ -100,6 +111,7 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
+        **metadata,
         'data': [
             {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
             {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},

tests/laborious/activities/test_opc.py

@@ -5,6 +5,15 @@ from laborious.activities.opc import NotificationLevel
 
 from laborious.activities.opc import OPC
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
+
 
 @patch("laborious.activities.opc.OpcRepository")
 def test___init__(mock_opc_repository):
@@ -15,6 +24,7 @@ def test___init__(mock_opc_repository):
     mock_notification_handler = MagicMock()
     servers = {
         'server1': {
+            'id': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
             'cert_path': '',
@@ -23,6 +33,7 @@ def test___init__(mock_opc_repository):
             'reconnection_interval': 60,
         },
         'server2': {
+            'id': 'server2',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
             'cert_path': '',
@@ -45,7 +56,7 @@ def test___init__(mock_opc_repository):
 
     mock_opc_repository.assert_has_calls([
         call(
-            name="server1",
+            id="server1",
             url="http://localhost:8080",
             logger=mock_logger,
             server_uri="opc.tcp://localhost:4840",
@@ -58,7 +69,7 @@ def test___init__(mock_opc_repository):
     ])
     mock_opc_repository.assert_has_calls([
         call(
-            name="server2",
+            id="server2",
             url="http://localhost:8080",
             logger=mock_logger,
             server_uri="opc.tcp://localhost:4840",
@@ -79,6 +90,7 @@ def test___init__(mock_opc_repository):
 def opc(mock_opc_repository):
     servers = {
         'server1': {
+            'id': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
             'cert_path': '',
@@ -90,12 +102,15 @@ def opc(mock_opc_repository):
     mock_opc_repository.write_data = MagicMock(
         return_value=True
     )
-    return OPC(
+    opc = OPC(
         opc_servers=servers,
         logger=MagicMock(),
         notification_handler=MagicMock()
     )
 
+    opc.send_notification = MagicMock()
+    return opc
+
 
 WRITE_DATA_CASES = [
     ('tag1', 'int', 50),
@@ -107,10 +122,10 @@ WRITE_DATA_CASES = [
 
 @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
 def test_write_data_success(opc, tag, data_type, data):
-    assert opc.write_data(server='server1', tag=tag, data=data,
-                          data_type=data_type, tag_type='prediction')
+    assert opc.write_data(server_id='server1', tag=tag, data=data,
+                          data_type=data_type, tag_type='prediction', metadata=metadata)
     opc.opc_repository['server1'].write_data.assert_called_once_with(
-        tag, data, data_type)
+        tag, data, data_type, opc.logger, metadata)
 
 
 def test_write_data_exception(opc):
@@ -118,11 +133,12 @@ def test_write_data_exception(opc):
         "Test error")
 
     try:
-        opc.write_data(server='server1', tag='tag1', data=50,
-                       data_type='int', tag_type='prediction')
+        opc.write_data(server_id='server1', tag='tag1', data=50,
+                       data_type='int', tag_type='prediction', metadata=metadata)
 
     except Exception:
-        opc.notification_handler.build_and_send_notification.assert_called_once_with(
+        opc.send_notification.assert_called_once_with(
+            metadata=metadata,
             notification_id="WRITE_OPC_PREDICTION_ERROR",
             message="Error writing data to OPC server: Test error",
             block="write_opc_data",
@@ -138,6 +154,7 @@ def test_write_data_exception(opc):
 async def test_write_opc_data_success(opc):
     # Arrange
     input_data = {
+        **metadata,
         'data': {
             'prediction': [0.75],
             'prediction_confidence': [0.95]
@@ -163,19 +180,21 @@ async def test_write_opc_data_success(opc):
     assert output == {'data': 'data'}
     opc.write_data.assert_has_calls([
         call(
-            server='server1',
+            server_id='server1',
             tag='tag1',
             data=0.75,
             data_type='float',
-            tag_type='prediction'
+            tag_type='prediction',
+            metadata=metadata['metadata']
         )])
     opc.write_data.assert_has_calls([
         call(
-            server='server1',
+            server_id='server1',
             tag='tag2',
             data=0.95,
             data_type='float',
-            tag_type='confidence'
+            tag_type='confidence',
+            metadata=metadata['metadata']
         )
     ])
     assert opc.write_data.call_count == 2
@@ -185,6 +204,7 @@ async def test_write_opc_data_success(opc):
 async def test_write_opc_data_empty_config(opc):
     # Arrange
     input_data = {
+        **metadata,
         'data': {
             'prediction': [0.75],
             'prediction_confidence': [0.95]
@@ -209,7 +229,7 @@ async def test_write_opc_data_empty_config(opc):
 ])
 def test_process_confidence(opc, data, success, expected):
     # Act
-    result = opc.process_confidence(data, success)
+    result = opc.process_confidence(data, success, metadata)
 
     # Assert
     assert result['prediction_confidence'][0] == expected

tests/laborious/utils/repository/test_opc_repository.py

@@ -14,7 +14,7 @@ def mock_logger():
 @fixture
 def opc_repository(mock_logger):
     return OpcRepository(
-        name="test_repo",
+        id="test_repo",
         url="opc.tcp://localhost:4840",
         logger=mock_logger,
         notification_handler=Mock(),
@@ -34,8 +34,18 @@ def mock_client():
         yield client_instance
 
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
+
+
 def test_init(opc_repository):
-    assert opc_repository.name == "test_repo"
+    assert opc_repository.id == "test_repo"
     assert opc_repository.url == "opc.tcp://localhost:4840"
     assert opc_repository.server_uri == "urn:test:server"
     assert opc_repository.cert_path == "/path/to/cert.pem"
@@ -110,7 +120,7 @@ def test_try_connect_fail(opc_repository):
     opc_repository.client.connect.assert_called_once()
     assert opc_repository.last_reconnection_time is not None
     opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}",
+        notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.id}",
         message="Failed to connect to OPC server: Test error",
         block="opc_repository",
         level=NotificationLevel.ERROR,
@@ -126,6 +136,17 @@ def test_disconnect(opc_repository, mock_client):
     assert opc_repository.client is None
 
 
+def test_disconnect_error(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.side_effect = Exception("Test error")
+    opc_repository.disconnect()
+
+    opc_repository.logger.error.assert_called_once_with(
+        "Failed to disconnect from OPC server: Test error"
+    )
+    assert opc_repository.client is None
+
+
 def test_validate_connection_none_client(opc_repository):
     opc_repository.client = None
     opc_repository.connect = MagicMock()
@@ -174,11 +195,11 @@ def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repositor
     opc_repository.client = MagicMock()
     opc_repository.client.aio_obj.uaclient.protocol = None
     opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
-    opc_repository.try_connect = MagicMock()
+    opc_repository.connect = MagicMock()
 
     response = opc_repository.validate_connection()
-    opc_repository.try_connect.assert_called_once()
-    assert response == opc_repository.try_connect.return_value
+    opc_repository.connect.assert_called_once()
+    assert response == opc_repository.connect.return_value
 
 
 def test_validate_connection_failed(opc_repository):
@@ -192,7 +213,8 @@ def test_validate_connection_failed(opc_repository):
 def test_write_data_validate_connection_do_nothing(opc_repository):
     opc_repository.validate_connection = MagicMock(return_value=True)
     opc_repository.client = MagicMock()
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "float", opc_repository.logger, metadata)
     opc_repository.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
 
@@ -201,7 +223,8 @@ def test_write_data_validate_connection_failed(opc_repository):
     opc_repository.validate_connection = MagicMock(return_value=False)
     opc_repository.client = MagicMock()
     opc_repository.error_count = 0
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "float", opc_repository.logger, metadata)
     opc_repository.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_not_called()
 
@@ -211,12 +234,13 @@ def test_write_data_get_node_failed(opc_repository):
     opc_repository.client = MagicMock()
     opc_repository.error_count = 0
     opc_repository.client.get_node.side_effect = Exception("Test error")
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "float", opc_repository.logger, metadata)
     opc_repository.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
     opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}",
-        message="Failed to get node from OPC server: Test error",
+        notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}",
+        message="Failed to get node from OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
         block="opc_repository",
         level=NotificationLevel.ERROR,
         attachment_content=ANY
@@ -230,13 +254,14 @@ def test_write_data_invalid_data_type(opc_repository, mock_client):
     mock_node = MagicMock()
     mock_client.get_node.return_value = mock_node
 
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "invalid_type")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "invalid_type", opc_repository.logger, metadata)
     opc_repository.validate_connection.assert_called_once()
     mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
 
     opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.name}",
-        message="Unsupported data type: invalid_type",
+        notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}",
+        message="Unsupported data type: invalid_type | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
         block="opc_repository",
         level=NotificationLevel.ERROR
     )
@@ -248,12 +273,11 @@ def test_write_data(opc_repository, mock_client):
     mock_node = MagicMock()
     mock_client.get_node.return_value = mock_node
 
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "float", opc_repository.logger, metadata)
 
     mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
     mock_node.write_value.assert_called_once()
-    opc_repository.logger.info.assert_called_once_with(
-        "Writing 42.0 - <class 'float'> to " + str(mock_node))
 
 
 def test_write_data_write_value_failed(opc_repository, mock_client):
@@ -263,13 +287,14 @@ def test_write_data_write_value_failed(opc_repository, mock_client):
     opc_repository.error_count = 0
     mock_client.get_node.return_value = mock_node
     mock_node.write_value.side_effect = Exception("Test error")
-    opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
+    opc_repository.write_data("ns=2;s=TestNode", 42.0,
+                              "float", opc_repository.logger, metadata)
     opc_repository.validate_connection.assert_called_once()
     mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
     mock_node.write_value.assert_called_once()
     opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}",
-        message="Failed to write data to OPC server: Test error",
+        notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.id}",
+        message="Failed to write data to OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
         block="opc_repository",
         level=NotificationLevel.ERROR,
         attachment_content=ANY

tests/laborious/utils/test_connectors_config.py

@@ -54,7 +54,7 @@ def test_build_opc_config_with_env_vars():
 def test_build_opc_config_with_individual_env_vars():
     # Arrange
     environ.pop('OPC_CONFIG', None)
-    environ['OPC_NAME'] = 'test-name'
+    environ['OPC_ID'] = '1'
     environ['OPC_URL'] = 'opc.tcp://test:4840'
     environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
     environ['OPC_RECONNECTION_INTERVAL'] = '300'
@@ -63,16 +63,16 @@ def test_build_opc_config_with_individual_env_vars():
     config = build_opc_config()
 
     # Assert
-    assert config['opc']['name'] == 'test-name'
-    assert config['opc']['url'] == 'opc.tcp://test:4840'
-    assert config['opc']['server_uri'] == 'opc.tcp://test:4840'
-    assert config['opc']['reconnection_interval'] == 300
+    assert config['1']['id'] == '1'
+    assert config['1']['url'] == 'opc.tcp://test:4840'
+    assert config['1']['server_uri'] == 'opc.tcp://test:4840'
+    assert config['1']['reconnection_interval'] == 300
 
 
 def test_build_opc_config_with_defaults():
     # Arrange
     environ.pop('OPC_CONFIG', None)
-    environ.pop('OPC_NAME', None)
+    environ.pop('OPC_ID', None)
     environ.pop('OPC_URL', None)
     environ.pop('OPC_SERVER_URI', None)
     environ.pop('OPC_RECONNECTION_INTERVAL', None)
@@ -81,10 +81,10 @@ def test_build_opc_config_with_defaults():
     config = build_opc_config()
 
     # Assert
-    assert config['opc']['name'] == 'opc'
-    assert config['opc']['url'] == 'opc.tcp://localhost:4840'
-    assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840'
-    assert config['opc']['reconnection_interval'] == 120
+    assert config['1']['id'] == '1'
+    assert config['1']['url'] == 'opc.tcp://localhost:4840'
+    assert config['1']['server_uri'] == 'opc.tcp://localhost:4840'
+    assert config['1']['reconnection_interval'] == 120
 
 
 def test_build_postgres_config_with_env_vars():

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -10,11 +10,22 @@ def format_and_export_prediction():
     return FormatAndExportPrediction()
 
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
+
+
 @mark.asyncio
 @patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
 async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
 
     input_data = {
+        'metadata': metadata,
         "path_flag": None,
         "data": {"test": "data"},
         "timestamp": "2021-01-01",
@@ -35,7 +46,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                 'data': input_data['data'],
                 'timestamp': input_data['timestamp'],
                 'model_id': input_data['model_id'],
-                'prediction_confidence': input_data['prediction_confidence']
+                'prediction_confidence': input_data['prediction_confidence'],
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
@@ -46,7 +58,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
             Activities.write_opc_data,
             {
                 'opc_output_config': input_data['opc_output_config'],
-                'data': workflow_mock.execute_local_activity_method.return_value
+                'data': workflow_mock.execute_local_activity_method.return_value,
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
@@ -59,7 +72,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
             {
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_activity_method.return_value
+                'data': workflow_mock.execute_activity_method.return_value,
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
@@ -74,6 +88,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
 async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
 
     input_data = {
+        'metadata': metadata,
         "path_flag": "default",
         "data": {"test": "data"},
         "timestamp": "2021-01-01",
@@ -95,7 +110,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                 'timestamp': input_data['timestamp'],
                 'model_id': input_data['model_id'],
                 'prediction_confidence': input_data['prediction_confidence'],
-                'comment': input_data['comment']
+                'comment': input_data['comment'],
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
@@ -107,7 +123,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
             Activities.write_opc_data,
             {
                 'opc_output_config': input_data['opc_output_config'],
-                'data': workflow_mock.execute_local_activity_method.return_value
+                'data': workflow_mock.execute_local_activity_method.return_value,
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
@@ -120,7 +137,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
             {
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_activity_method.return_value
+                'data': workflow_mock.execute_activity_method.return_value,
+                **metadata
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -9,12 +9,23 @@ def prediction_process():
     return PredictionProcess()
 
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model",
+        "model_name": "test_model",
+        "workflow_name": "test_workflow",
+        "schema_name": "test_schedule",
+    },
+}
+
+
 @mark.asyncio
 @patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
 async def test_run(workflow_mock, prediction_process):
     prediction_process.path_flag_handler = AsyncMock(return_value=False)
     # Arrange
     input_data = {
+        'metadata': metadata,
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
@@ -49,51 +60,61 @@ async def test_run(workflow_mock, prediction_process):
     assert workflow_mock.execute_local_activity_method.call_count == 7
 
     workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {'data': input_data['data']},
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+        call(Activities.get_last_timestamp, {
+            'data': input_data['data'],
+            **metadata
+        },
+            retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.input_gate, {
             'filters': input_data['input_filters'],
             'data': input_data['data'],
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_transform, {
             'data': input_data['data'],
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']
+            'model_retention': input_data['model_retention'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_response_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_content_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': 'transformed_data',
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_predict, {
             'data': 'transformed_data',
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']
+            'model_retention': input_data['model_retention'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_response_gate, {
             'filters': input_data['mlflow_predict_filters'],
             'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
             'type': 'predict',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
 
     workflow_mock.execute_child_workflow.assert_called_once_with(
         'format_and_export_prediction',
         {
+            'metadata': metadata,
             'path_flag': 'continue',
             'data': 'predicted_data',
             'prediction_confidence': 0.95,
@@ -115,6 +136,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
     prediction_process.path_flag_handler = AsyncMock(return_value=True)
     # Arrange
     input_data = {
+        'metadata': metadata,
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
@@ -141,11 +163,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
     assert workflow_mock.execute_local_activity_method.call_count == 2
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.get_last_timestamp, {
-             'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
+             'data': input_data['data'],
+             **metadata
+             }, retry_policy=ANY, start_to_close_timeout=ANY),
         call(Activities.input_gate, {
              'filters': input_data['input_filters'],
              'data': input_data['data'],
-             'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
+             'path_priority': input_data['path_priority'],
+             **metadata
+             }, retry_policy=ANY, start_to_close_timeout=ANY)
     ])
     workflow_mock.execute_child_workflow.assert_not_called()
 
@@ -156,6 +182,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
     prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
     # Arrange
     input_data = {
+        'metadata': metadata,
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
@@ -183,19 +210,26 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 4
     workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {'data': input_data['data']},
+        call(Activities.get_last_timestamp, {
+             'data': input_data['data'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.input_gate, {
              'filters': input_data['input_filters'],
              'data': input_data['data'],
-             'path_priority': input_data['path_priority']},
+             'path_priority': input_data['path_priority'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_transform, {
             'data': input_data['data'],
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']},
+            'model_retention': input_data['model_retention'],
+            **metadata
+        },
             retry_policy=ANY, start_to_close_timeout=ANY)
     ])
     workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -203,7 +237,8 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
             'filters': input_data['mlflow_transform_filters'],
             'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)
     ])
     workflow_mock.execute_child_workflow.assert_not_called()
@@ -216,6 +251,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
         side_effect=[False, False, True])
     # Arrange
     input_data = {
+        'metadata': metadata,
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
@@ -247,33 +283,41 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
     assert workflow_mock.execute_local_activity_method.call_count == 5
 
     workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {'data': input_data['data']},
+        call(Activities.get_last_timestamp, {
+             'data': input_data['data'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.input_gate, {
              'filters': input_data['input_filters'],
              'data': input_data['data'],
-             'path_priority': input_data['path_priority']},
+             'path_priority': input_data['path_priority'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_transform, {
             'data': input_data['data'],
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']
+            'model_retention': input_data['model_retention'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_response_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_content_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': 'transformed_data',
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_child_workflow.assert_not_called()
 
@@ -285,6 +329,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
         side_effect=[False, False, False, True])
     # Arrange
     input_data = {
+        'metadata': metadata,
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
@@ -317,46 +362,56 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 7
     workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {'data': input_data['data']},
+        call(Activities.get_last_timestamp, {
+             'data': input_data['data'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.input_gate, {
              'filters': input_data['input_filters'],
              'data': input_data['data'],
-             'path_priority': input_data['path_priority']},
+             'path_priority': input_data['path_priority'],
+             **metadata
+             },
              retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_transform, {
             'data': input_data['data'],
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']
+            'model_retention': input_data['model_retention'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_response_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_content_gate, {
             'filters': input_data['mlflow_transform_filters'],
             'data': 'transformed_data',
             'type': 'transform',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.request_predict, {
             'data': 'transformed_data',
             'model_name': input_data['model_name'],
-            'model_retention': input_data['model_retention']
+            'model_retention': input_data['model_retention'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(Activities.mlflow_response_gate, {
             'filters': input_data['mlflow_predict_filters'],
             'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
             'type': 'predict',
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
+            **metadata
         }, retry_policy=ANY, start_to_close_timeout=ANY)])
     workflow_mock.execute_child_workflow.assert_not_called()
 
@@ -378,6 +433,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
     # Act
     result = await prediction_process.path_flag_handler(
         data, path_flag, {
+            'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
@@ -410,6 +466,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     # Act
     result = await prediction_process.path_flag_handler(
         data, path_flag, {
+            'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
@@ -424,6 +481,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     workflow_mock.execute_activity_method.assert_called_once_with(
         Activities.repeat_last_prediction,
         {
+            **metadata,
             'schema': schema,
             'table_name': table_name,
             'model': model,
@@ -452,6 +510,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     # Act
     result = await prediction_process.path_flag_handler(
         data, path_flag, {
+            'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
@@ -468,6 +527,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     workflow_mock.execute_child_workflow.assert_called_once_with(
         'format_and_export_prediction',
         {
+            'metadata': metadata,
             'path_flag': path_flag,
             'data': data,
             'prediction_confidence': confidence,
@@ -500,6 +560,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
     # Act
     result = await prediction_process.path_flag_handler(
         data, path_flag, {
+            **metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,

tests/laborious/workflows/test_predictions_batch.py

@@ -9,6 +9,16 @@ def predictions_batch() -> PredictionsBatch:
     return PredictionsBatch()
 
 
+metadata = {
+    "metadata": {
+        "model_id": "test_model_id",
+        "model_name": "test_model",
+        "workflow_name": "predictions_batch",
+        "schedule_name": "test_schedule",
+    },
+}
+
+
 @mark.asyncio
 @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
@@ -29,27 +39,17 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
 
     workflow_mock.execute_local_activity_method.assert_has_calls([
         call(
-            Activities.prepare_activity,
+            Activities.load_custom_query,
             {
-                'schedule_name': input_data['schedule_name'],
-                'model_name': input_data['model_name'],
-                'model_id': input_data['model_id'],
-                'workflow_name': 'predictions_batch'
+                **metadata,
+                'query': input_data['query'],
             },
             retry_policy=ANY,
             start_to_close_timeout=ANY
         )
     ])
-
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.load_custom_query,
-            input_data['query'],
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
     prediction_input = {
+        'metadata': metadata,
         'data': {'data': 'test_data'},
         'schema': input_data['schema'],
         'table_name': input_data['table_name'],

values.yaml

@@ -3,7 +3,7 @@
 # Declare variables to be passed into your templates.
 
 # This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/
-replicaCount: 1
+replicaCount: 5
 
 # This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/
 image:
@@ -11,7 +11,7 @@ image:
   # This sets the pull policy for images.
   pullPolicy: Always
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.2.2"
+  tag: "0.2.4"
 
 # This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
 imagePullSecrets:
@@ -123,7 +123,7 @@ env:
   - name: GITHUB_REPO_URL
     value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git"
   - name: GITHUB_BRANCH
-    value: "SIENTIAPDE-1110-criar-testes-e-2-e"
+    value: "SIENTIAPDE-1154-otimizar-conexao-com-banco-de-dados-e-paralelismo"
   - name: PYTHON_APP
     value: "laborious.worker.worker"
 
@@ -141,7 +141,7 @@ env:
   - name: POSTGRES_MIN_CONNECTIONS
     value: "10"
   - name: POSTGRES_MAX_CONNECTIONS
-    value: "20"
+    value: "30"
 
   - name: MLFLOW_HOST
     value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
@@ -168,7 +168,7 @@ env:
   - name: TEMPORAL_HOST
     value: "temporal-frontend.temporal.svc.cluster.local:7233"
   - name: TEMPORAL_NAMESPACE
-    value: "default"
+    value: "laborious"
 
 ssh:
   enabled: true