SIENTIAPDE-1030

Add unit tests for connectors configuration, logger, workflows, and predictions batch

- Implement tests for MLflow, OPC, and Postgres configuration builders to validate environment variable handling and default values.
- Create tests for the logger to ensure default settings and handler configurations are correct.
- Add comprehensive tests for the FormatAndExportPrediction and PredictionProcess workflows, covering various scenarios including path flags and activity execution.
- Introduce tests for the PredictionsBatch workflow to verify the execution of local activities and child workflows.
- Include a values.yaml file for Kubernetes deployment configuration, specifying image details, service account settings, environment variables, and resource limits.

laborious/activities/postgres.py

@@ -0,0 +1,181 @@
+import traceback
+from temporalio import workflow, activity
+
+from laborious.activities.base import BaseActivity
+with workflow.unsafe.imports_passed_through():
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import sessionmaker
+    from sqlalchemy.pool import QueuePool
+    from psycopg2.pool import ThreadedConnectionPool
+    from pandas import read_sql_query, DataFrame
+    from logging import Logger
+    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.models import NotificationLevel
+    from typing import Any
+
+
+class Postgres(BaseActivity):
+    def __init__(self, host: str, port: int,
+                 user: str, password: str, dbname: str,
+                 min_connections: int, max_connections: int,
+                 logger: Logger, notification_handler: NotificationHandler):
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        self.dbname = dbname
+
+        # Create SQLAlchemy engine with connection pooling
+        self.engine = create_engine(
+            f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
+            poolclass=QueuePool,
+            pool_size=min_connections,
+            max_overflow=max_connections - min_connections,
+            pool_pre_ping=True
+        )
+        self.session_factory = sessionmaker(bind=self.engine)
+
+        BaseActivity.__init__(self, logger, notification_handler)
+
+    def close(self):
+        self.engine.dispose()
+
+    def __del__(self):
+        self.close()
+
+    @activity.defn(name="load_custom_query")
+    async def load_custom_query(self, query: str) -> dict[str, Any]:
+        """
+        Loads data from a custom query.
+
+        Args:
+            query (str): The query to load data from.
+
+        Returns:
+            dict[str, dict]: The data from the query.
+        """
+        self.logger.info(f"Fetching data from query: {query}")
+
+        data = None
+        with self.session_factory() as session:
+            try:
+                data = read_sql_query(query, self.engine)
+
+            except Exception as e:
+                trace = traceback.format_exc()
+                self.notification_handler.build_and_send_notification(
+                    notification_id="ERROR_LOADING_CUSTOM_QUERY",
+                    message=f"Error fetching data from query: {e}",
+                    block="load_custom_query",
+                    level=NotificationLevel.ERROR,
+                    attachment_content=trace
+                )
+
+                self.logger.error(trace)
+
+                return {}
+            finally:
+                session.close()
+
+        if data is None:
+            return {}
+
+        # Converts any datetime datatype columns to string
+        for col in data.select_dtypes(include=['datetime64']).columns:
+            data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S')
+
+        self.logger.info(f"Fetched {len(data)} rows")
+        self.logger.debug(f"Data: \n{data.to_string()}")
+
+        return data.to_dict()
+
+    @activity.defn(name="repeat_last_prediction")
+    async def repeat_last_prediction(self, query_items: dict[str, str]):
+        """
+        Repeats the last prediction for a given model.
+
+        Args:
+            query_items (dict[str, str]): The query items. Contains:
+                schema (str): The schema of the table.
+                table_name (str): The name of the table.
+                model (int): The model to repeat the prediction for.
+
+        Returns:
+            None
+        """
+        schema = query_items["schema"]
+        table_name = query_items["table_name"]
+        model = query_items["model"]
+
+        repeat_query = f"""
+        INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
+        SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
+        FROM \"{schema}\".{table_name}
+        WHERE model_id = {model}
+        ORDER BY timestamp DESC
+        LIMIT 1;
+        """
+        self.logger.info(f"Repeating last prediction for model {model}")
+        self.logger.debug(f"Query: {repeat_query}")
+
+        with self.session_factory() as session:
+            try:
+                session.execute(repeat_query)
+                session.commit()
+
+            except Exception as e:
+                trace = traceback.format_exc()
+                self.notification_handler.build_and_send_notification(
+                    notification_id="ERROR_REPEATING_LAST_PREDICTION",
+                    message=f"Error repeating last prediction: {e}",
+                    block="repeat_last_prediction",
+                    level=NotificationLevel.ERROR,
+                    attachment_content=trace
+                )
+
+                self.logger.error(trace)
+
+            finally:
+                session.close()
+
+    @activity.defn(name="export_data_to_postgres")
+    async def export_data_to_postgres(self, input_data: dict[str, Any]):
+        """
+        Exports data to a postgres table.
+
+        Args:
+            input_data (dict[str, Any]): The data to export. Contains:
+                schema (str): The schema of the table.
+                table_name (str): The name of the table.
+                data (DataFrame): The data to export.
+        """
+
+        self.logger.debug(
+            f"Exporting data to postgres: {input_data['data']}")
+
+        schema = input_data["schema"]
+        table_name = input_data["table_name"]
+        data = DataFrame(input_data["data"])
+
+        with self.session_factory() as session:
+            try:
+                data.to_sql(table_name, self.engine, schema=schema,
+                            if_exists="append", index=False)
+                session.commit()
+
+            except Exception as e:
+                trace = traceback.format_exc()
+                self.notification_handler.build_and_send_notification(
+                    notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
+                    message=f"Error exporting data to postgres: {e}",
+                    block="export_data_to_postgres",
+                    level=NotificationLevel.ERROR,
+                    attachment_content=trace
+                )
+
+                self.logger.error(trace)
+
+            else:
+                self.logger.debug("Data exported to postgres")
+            finally:
+                session.close()