SIENTIAPDE-1273

SIENTIAPDE-1273
Enhance security analysis and SQL injection handling

- Added skip for potential SQL injection false positives in Bandit configuration.
- Updated validate.sh to use the pyproject.toml configuration for Bandit security analysis.
- Refactored code to replace ensure_dataframe utility with direct DataFrame usage in multiple activities, improving clarity and reducing dependencies.
- Removed the deprecated dataframe_utils module to streamline the codebase.

laborious/activities/activities.py

@@ -9,9 +9,10 @@ with workflow.unsafe.imports_passed_through():
 
     from laborious.activities.gates import Gates
     from laborious.activities.mlflow import MLFlow
+    from laborious.activities.model_metrics import ModelMetrics
     from laborious.activities.opc import OPC
     from laborious.activities.storage import Storage
-    from laborious.activities.model_metrics import ModelMetrics
+
 
 class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
     """

laborious/activities/gates.py

@@ -15,7 +15,6 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
 
     from laborious import metrics
-    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.filters.conditional_filters import (
         filter_empty_data,
         filter_specific_variables_null_values,
@@ -150,7 +149,7 @@ class Gates(SientiaMonitoring):
         self.info('Performing input gate...', metadata)
 
         filters = input_data['filters']
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
         path_priority = input_data['path_priority']
 
         filter_output = []
@@ -403,7 +402,6 @@ class Gates(SientiaMonitoring):
 
         return policy_type, int(policy_value)
 
-
     @activity.defn(name='format_transformed_data')
     async def format_transformed_data(self, input_data: dict[str, Any]) -> dict:
         """
@@ -415,7 +413,7 @@ class Gates(SientiaMonitoring):
 
         self.info('Formatting transformed data...', metadata)
 
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
 
         data['timestamp'] = data.index
         data = data.reset_index(drop=True)
@@ -454,7 +452,7 @@ class Gates(SientiaMonitoring):
         prediction_store_policy = input_data['prediction_store_policy']
         self.info('Formatting prediction...', metadata)
 
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
 
         # Create timestamp column from index and reset index
         data['timestamp'] = data.index
@@ -602,7 +600,7 @@ class Gates(SientiaMonitoring):
 
         self.info('Getting last timestamp...', metadata)
 
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
 
         self.debug(f'Input data: {data.head(5).to_string()}', metadata)
 

laborious/activities/mlflow.py

@@ -1,4 +1,3 @@
-from typing import Hashable
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -20,7 +19,6 @@ with workflow.unsafe.imports_passed_through():
         now,
     )
 
-    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.repository.minio_repository import MinioRepository
     from laborious.utils.repository.model_repository import MLFlowRepository
 
@@ -140,7 +138,7 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Transforming data...', metadata)
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
@@ -212,7 +210,7 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Predicting data...', metadata)
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
@@ -421,7 +419,6 @@ class MLFlow(SientiaMonitoring):
             self.error(trace, metadata=metadata)
             raise e
 
-
     @activity.defn(name='get_reference_data')
     async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
         """
@@ -439,7 +436,7 @@ class MLFlow(SientiaMonitoring):
 
         metadata = input_data['metadata']
         model_name = input_data['model_name']
-        artifact = "evaluation_data.csv"
+        artifact = 'evaluation_data.csv'
 
         reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
             model_name=model_name, artifact_path=artifact, metadata=metadata

laborious/activities/model_metrics.py

@@ -1,25 +1,28 @@
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
+    import time
+    import traceback
+    import warnings
     from typing import Any
 
+    import numpy as np
     from pandas import DataFrame, Index, to_datetime
-
+    from sientia.ModelAnalysis import ModelAnalysis
-    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
-    from sientia_do.observability.metrics_controller import MetricsController
-    from sientia_do.observability.logger import Logger
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from sientia.ModelAnalysis import ModelAnalysis
+    from sientia_do.observability.logger import Logger
-    from laborious import metrics
+    from sientia_do.observability.metrics_controller import MetricsController
-    import time
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
-    import numpy as np
     from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
-    import warnings
+
-    import traceback
+    from laborious import metrics
 
 warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
-warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide')
+warnings.filterwarnings(
+    'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
+)
+
 
 class ModelMetrics(SientiaMonitoring):
     """
@@ -28,12 +31,12 @@ class ModelMetrics(SientiaMonitoring):
     This class provides activities for writing metrics to the Prometheus monitoring system.
     """
 
-    def __init__(self,
+    def __init__(
+        self,
         logger: Logger,
         notification_handler: NotificationHandler,
         metrics_controller: MetricsController,
     ):
-
         SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
 
     def close(self) -> None:
@@ -42,11 +45,11 @@ class ModelMetrics(SientiaMonitoring):
         """
         SientiaMonitoring.shutdown(self)
 
-
     def __del__(self):
         self.close()
 
-    async def get_drift_metrics(self,
+    async def get_drift_metrics(
+        self,
         reference_data: DataFrame,
         target_data: DataFrame,
         target_name: str,
@@ -73,34 +76,37 @@ class ModelMetrics(SientiaMonitoring):
             'features': reference_columns,
         }
 
-
         model_analysis = ModelAnalysis(config=config)
 
-        self.debug(f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', metadata)
+        self.debug(
+            f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}',
+            metadata,
+        )
 
-        self.debug(f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata)
+        self.debug(
+            f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata
+        )
 
         core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
         start_time = time.time()
         try:
-
             univariate_drift = model_analysis.detect_univariate_drift(
                 reference_df=reference_data,
                 analysis_df=target_data,
                 features=reference_columns,
                 timestamp_col=config['timestamp'],
                 methods=drift_metrics,
-                chunk_period=chunk_period
+                chunk_period=chunk_period,
             )
         except Exception as e:
             self.error(f'Error detecting univariate drift: {e}', metadata)
             await self.emit_metric(
-                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
             raise e
         await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
 
-
         core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
         start_time = time.time()
         try:
@@ -109,11 +115,13 @@ class ModelMetrics(SientiaMonitoring):
                 analysis_df=target_data,
                 features=reference_columns,
                 timestamp_col=config['timestamp'],
-                chunk_period=chunk_period
+                chunk_period=chunk_period,
             )
         except Exception as e:
             self.error(f'Error detecting multivariate drift: {e}', metadata)
-            await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
+            await self.emit_metric(
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
             raise e
         await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
@@ -127,16 +135,19 @@ class ModelMetrics(SientiaMonitoring):
             )
         except Exception as e:
             self.error(f'Error getting drift metrics: {e}', metadata)
-            await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
+            await self.emit_metric(
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
             raise e
         await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
 
-        self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
+        self.debug(
+            f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
+        )
 
         return drift_df
 
-
     @activity.defn(name='calculate_drift')
     async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
         """
@@ -194,8 +205,8 @@ class ModelMetrics(SientiaMonitoring):
             )
 
         reference_columns = reference_data.drop(
-            columns=[target_name, 'timestamp', 'target', 'prediction'],
+            columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
-            errors='ignore').columns
+        ).columns
 
         try:
             drift_df = await self.get_drift_metrics(
@@ -224,12 +235,11 @@ class ModelMetrics(SientiaMonitoring):
             return []
 
         # Drop unnecessary columns
-        drift_df.drop(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+        drift_df.drop(columns=['p_value'], inplace=True)
 
         # Extract timestamps only until minutes
         if chunk_period == 'min':
-            target_timestamps = target_data['timestamp'].apply(
+            target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
-                lambda x: x[:16])
         else:
             target_timestamps = target_data['timestamp']
 
@@ -237,29 +247,36 @@ class ModelMetrics(SientiaMonitoring):
         drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
 
         if drift_df.empty:
-            self.warning('No drift metrics found after dropping rows where timestamp is not in target data', metadata)
+            self.warning(
+                'No drift metrics found after dropping rows where timestamp is not in target data',
+                metadata,
+            )
             return []
 
         # Rename columns to match database columns
-        drift_df.rename(columns={
+        drift_df.rename(
+            columns={
                 'metric': 'method',
                 'statistic': 'value',
-        }, inplace=True)
+            },
+            inplace=True,
+        )
 
         # Drop duplicates
         drift_df.drop_duplicates(
-            subset=['timestamp', 'method', 'feature'],
+            subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
-            keep='first', inplace=True)
+        )
 
         drift_df['model_id'] = model_id
         drift_df['accurate'] = accurate
 
-        
         drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
         drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
         drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
 
-        self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
+        self.debug(
+            f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
+        )
 
         self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata)
 
@@ -298,26 +315,17 @@ class ModelMetrics(SientiaMonitoring):
         output_data = []
 
         diff = target_data['target'] - target_data['prediction']
-        diff_squared = diff ** 2
+        diff_squared = diff**2
 
         self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
 
         for metric in metrics:
             if metric == 'rmse':
-                output_data.append({
+                output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
-                    'metric': 'rmse',
-                    'value': np.sqrt(np.mean(diff_squared))
-                })
             elif metric == 'mse':
-                output_data.append({
+                output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
-                    'metric': 'mse',
-                    'value': np.mean(diff_squared)
-                })
             elif metric == 'mae':
-                output_data.append({
+                output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
-                    'metric': 'mae',
-                    'value': np.mean(np.abs(diff))
-                })
             elif metric == 'r2':
                 y_true = target_data['target']
                 y_mean = np.mean(y_true)
@@ -331,10 +339,7 @@ class ModelMetrics(SientiaMonitoring):
                 else:
                     r2_score = 1 - (ss_res / ss_tot)
 
-                output_data.append({
+                output_data.append({'metric': 'r2', 'value': r2_score})
-                    'metric': 'r2',
-                    'value': r2_score
-                })
 
         data = DataFrame(output_data)
         data['model_id'] = model_id
@@ -342,8 +347,8 @@ class ModelMetrics(SientiaMonitoring):
         data['data_size'] = data_size
         data['interval_minutes'] = interval_minutes
 
-        self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata)
+        self.debug(
+            f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata
+        )
 
         return data.to_dict(orient='records')
-
-

laborious/activities/opc.py

@@ -1,4 +1,5 @@
-from typing import Hashable
+from collections.abc import Hashable
+
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -12,7 +13,6 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.observability.metrics_controller import MetricsController
     from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 
-    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.repository.opc_repository import OpcRepository
 
 OPC_WRITTING_ERROR_CONFIDENCE = 12
@@ -297,7 +297,7 @@ class OPC(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Writing data to OPC servers...', metadata)
-        data = ensure_dataframe(input_data['data'])
+        data = DataFrame(input_data['data'])
         opc_output_config = input_data['opc_output_config']
         self.info(f'Data to write: {data.size} rows', metadata)
 

laborious/utils/dataframe_utils.py

@@ -1,31 +0,0 @@
-"""
-DataFrame utility functions for handling serialized DataFrames.
-
-This module provides helper functions to work with DataFrames that may
-come from Temporal serialization (already as DataFrame) or from legacy
-code (as dict).
-"""
-
-from typing import Any
-
-from pandas import DataFrame
-
-
-def ensure_dataframe(data: Any) -> DataFrame:
-    """
-    Ensure that data is a DataFrame, converting from dict if necessary.
-
-    This function handles both cases:
-    - Data already deserialized as DataFrame (from Temporal codec)
-    - Data as dict (legacy format or non-DataFrame serialization)
-
-    Args:
-        data: Data that should be a DataFrame (can be DataFrame or dict)
-
-    Returns:
-        DataFrame: The data as a pandas DataFrame
-    """
-    if isinstance(data, DataFrame):
-        return data
-    return DataFrame(data)
-

laborious/utils/repository/model_repository.py

@@ -16,11 +16,11 @@ Capabilities:
 
 import ctypes
 import gc
-from io import StringIO
 import threading
 import time
 import traceback
 from datetime import datetime, timedelta
+from io import StringIO
 from os import environ, makedirs, path
 from shutil import rmtree
 from typing import Any, Literal, overload
@@ -34,7 +34,6 @@ from sientia_do.observability.logger import Logger
 from sientia_do.observability.metrics_controller import MetricsController
 from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
-from sientia.ModelAnalysis import ModelAnalysis
 
 from laborious import metrics
 
@@ -217,8 +216,9 @@ class MLFlowRepository(SientiaMonitoring):
         run_info = mlflow.get_run(run_id)
         return run_info.data.params
 
-    def check_artifact_exists(self, run_id: str,
+    def check_artifact_exists(
-    artifact_path: str, metadata: dict[str, Any]) -> bool:
+        self, run_id: str, artifact_path: str, metadata: dict[str, Any]
+    ) -> bool:
         """
         Check if an artifact exists in the MLflow Model Registry.
 
@@ -235,6 +235,7 @@ class MLFlowRepository(SientiaMonitoring):
         self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata)
 
         return any(artifact.path == artifact_path for artifact in artifacts)
+
     """
     Functions related to download and load models
     """
@@ -279,10 +280,9 @@ class MLFlowRepository(SientiaMonitoring):
 
         return artifacts
 
-
+    async def load_artifact_dataframe(
-    async def load_artifact_dataframe(self, model_name: str, artifact_path: str,
+        self, model_name: str, artifact_path: str, metadata: dict[str, Any]
-                                         metadata: dict[str, Any]) -> pd.DataFrame | None:
+    ) -> pd.DataFrame | None:
-        
         """
         Load the dataframe content of an artifact from the MLflow Model Registry.
 
@@ -300,7 +300,7 @@ class MLFlowRepository(SientiaMonitoring):
         if not self.check_artifact_exists(run_id, artifact_path, metadata):
             return None
 
-        artifact_path = path.join("runs:/", run_id, artifact_path)
+        artifact_path = path.join('runs:/', run_id, artifact_path)
 
         start_time = time.time()
         try:
@@ -704,8 +704,9 @@ class MLFlowRepository(SientiaMonitoring):
     Functions related to model retraining
     """
 
-    def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame,
+    def get_prediction_data(
-    target_name: str) -> pd.DataFrame:
+        self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str
+    ) -> pd.DataFrame:
         """
         Get prediction data from prediction model.
         """
@@ -714,7 +715,6 @@ class MLFlowRepository(SientiaMonitoring):
         prediction_data = prediction_model.predict(retrain_dataset)
 
         if isinstance(prediction_data, pd.DataFrame):
-
             prediction_data.columns = pd.Index(['prediction'])
 
         else:
@@ -724,7 +724,8 @@ class MLFlowRepository(SientiaMonitoring):
 
         # Merge prediction data with retrain_dataset on index
         prediction_data = pd.merge(
-            retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')
+            retrain_dataset, prediction_data, left_index=True, right_index=True, how='left'
+        )
 
         # Rename column "target_name" to "target"
         prediction_data.rename(columns={target_name: 'target'}, inplace=True)
@@ -733,9 +734,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_data.reset_index(drop=True, inplace=True)
 
-        prediction_data.sort_values(
+        prediction_data.sort_values(by='timestamp', ascending=True, inplace=True)
-            by='timestamp', ascending=True, inplace=True
-        )
 
         return prediction_data
 
@@ -859,8 +858,7 @@ class MLFlowRepository(SientiaMonitoring):
         prediction_model.fit(retrain_dataset)
 
         # get prediction data
-        prediction_data = self.get_prediction_data(
+        prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name)
-            prediction_model, retrain_dataset, target_name)
 
         self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
 
@@ -1439,4 +1437,3 @@ class MLFlowRepository(SientiaMonitoring):
         metadata_result['mlflow_experiment_id'] = experiment_id
 
         return metadata_result
-

laborious/worker/worker.py

@@ -29,7 +29,6 @@ from temporalio import client, workflow
 from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 from temporalio.worker import PollerBehaviorAutoscaling, Worker
 
-
 with workflow.unsafe.imports_passed_through():
     import asyncio
     import os
@@ -48,11 +47,10 @@ with workflow.unsafe.imports_passed_through():
         build_opc_config,
         build_postgres_config,
     )
+    from laborious.workflows.drift import Drift
     from laborious.workflows.minimal_retrain import MinimalRetrain
     from laborious.workflows.predictions_batch import PredictionsBatch
-    from laborious.workflows.drift import Drift
     from laborious.workflows.simple_metrics import SimpleMetrics
-    
     from laborious.workflows.sub_workflows.format_and_export_prediction import (
         FormatAndExportPrediction,
     )

laborious/workflows/drift.py

@@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through():
     from datetime import timedelta
     from typing import Any
 
+    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
     from sientia_do.temporal.policies import retry_policy
 
     from laborious.activities.activities import Activities
-    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @workflow.defn(name='drift')
@@ -39,10 +39,10 @@ class Drift:
 
         gathering_query = f"""
             SELECT *
-            FROM {input_data['schema']}.{input_data['source_table_name']}
+            FROM "{input_data['schema']}"."{input_data['source_table_name']}"
             WHERE
-                model_id = {input_data['model_id']} AND
+                model_id = '{input_data['model_id']}' AND
-                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+                timestamp > NOW() - INTERVAL {input_data['interval']} minutes
             ORDER BY timestamp ASC
         """
 
@@ -60,10 +60,7 @@ class Drift:
 
         reference_data_handler = workflow.start_local_activity_method(
             Activities.get_reference_data,
-            {
+            {**metadata, 'model_name': input_data['model_name']},
-                **metadata,
-                'model_name': input_data['model_name']
-            },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=300),
         )
@@ -83,8 +80,9 @@ class Drift:
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
                 'target_name': target_name,
-                'drift_metrics': input_data.get('drift_metrics',
+                'drift_metrics': input_data.get(
-                    ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']),
+                    'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
+                ),
                 'chunk_period': input_data.get('chunk_period', 'min'),
             },
             retry_policy=retry_policy,
@@ -99,7 +97,10 @@ class Drift:
                     'data': drift_data,
                     'schema': input_data['schema'],
                     'table_name': input_data['target_table_name'],
-                    'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
                 },
                 retry_policy=retry_policy,
                 start_to_close_timeout=timedelta(seconds=300),

laborious/workflows/simple_metrics.py

@@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through():
     from datetime import timedelta
     from typing import Any
 
+    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
     from sientia_do.temporal.policies import retry_policy
 
     from laborious.activities.activities import Activities
-    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @workflow.defn(name='simple_metrics')
@@ -19,10 +19,10 @@ class SimpleMetrics:
         """
         metadata = {
             'metadata': {
-                'schedule_name': input_data['schedule_name'],
-                'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
+                'model_name': input_data['model_name'],
                 'workflow_name': 'simple_metrics',
+                'schedule_name': input_data['schedule_name'],
             }
         }
 
@@ -34,15 +34,15 @@ class SimpleMetrics:
 
         query = f"""
             select p."timestamp", p.prediction, ld.value as "target"
-            from {input_data['schema']}.{input_data['predictions_table_name']} p
+            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
-            inner join {input_data['schema']}.{input_data['data_table_name']} ld
+            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
                 on p."timestamp" = ld."timestamp"
             where
-                p.model_id = {model_id} and
+                p.model_id = '{model_id}' and
                 p.prediction is not null and
                 ld.variable = '{target_name}' and
                 ld.value is not null and
-                p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
+                p."timestamp" >= NOW() - INTERVAL {interval_minutes} minutes
             order by
                 p."timestamp" desc;
         """
@@ -75,7 +75,6 @@ class SimpleMetrics:
             start_to_close_timeout=timedelta(seconds=300),
         )
 
-
         if not simple_metrics:
             return
 

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -89,7 +89,6 @@ class FormatAndExportPrediction:
             )
 
             if transformed_data is not None:
-
                 transformed = await workflow.execute_local_activity_method(
                     Activities.format_transformed_data,
                     {

pyproject.toml

@@ -48,6 +48,7 @@ ignore = [
     "S101",  # use of assert (needed for tests)
     "S105",  # possible hardcoded password (false positives)
     "S106",  # possible hardcoded password (false positives)
+    "S608",  # potential sql injection (false positives)
     "N802",  # function name should be lowercase (temporal decorators)
     "N806",  # variable in function should be lowercase
 ]
@@ -152,4 +153,4 @@ directory = "htmlcov"
 
 [tool.bandit]
 exclude_dirs = ["tests", "venv", ".venv"]
-skips = ["B101", "B601"]  # Skip assert and shell injection in controlled environments
+skips = ["B101", "B601", "B608"]  # Skip assert, shell injection, and SQL injection (false positives)

tests/laborious/activities/test_model_metrics.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
 from pandas import DataFrame
 from pytest import fixture, mark
@@ -22,7 +22,13 @@ def model_metrics_activity():
     model_metrics.send_notification = MagicMock()
     model_metrics.send_notification_async = AsyncMock()
     model_metrics.emit_metric = AsyncMock()
-    model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'})
+    model_metrics.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
     model_metrics.observe_lag = AsyncMock()
     model_metrics.pod_id = 'test_pod'
     return model_metrics
@@ -60,7 +66,7 @@ async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
     try:
         await model_metrics_activity.calculate_drift(input_data)
     except ValueError as e:
-        assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\""
+        assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
         model_metrics_activity.error.assert_called_once_with(
             'Invalid chunk period: invalid', metadata['metadata']
         )
@@ -76,25 +82,40 @@ async def test_calculate_drift_with_reference_data(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
 
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -124,16 +145,20 @@ async def test_calculate_drift_with_reference_data(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
+    assert isinstance(result, list)
-    assert result == mock_drift_df.to_dict.return_value
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
+    mock_drift_df.rename.assert_called_once_with(
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
-    mock_drift_df.to_dict.assert_called_once()
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -144,17 +169,30 @@ async def test_calculate_drift_without_reference_data(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
 
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]}
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': False,
+        }
+    ]
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
@@ -170,8 +208,14 @@ async def test_calculate_drift_without_reference_data(
     mock_target_df.reset_index.return_value = mock_target_df
     mock_target_df.dropna.return_value = mock_target_df
     mock_target_df.sort_values.return_value = mock_target_df
-    mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]})
+    mock_target_df.head.return_value = DataFrame(
-    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
+        {'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}
+    )
+    mock_target_df.__getitem__.return_value.apply.return_value = [
+        '2023-05-26 11:12:27',
+        '2023-05-26 11:12:28',
+        '2023-05-26 11:12:29',
+    ]
     mock_target_df.drop.return_value.columns = ['feature1']
     mock_dataframe.return_value = mock_target_df
     mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df
@@ -191,7 +235,7 @@ async def test_calculate_drift_without_reference_data(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
+    assert isinstance(result, list)
     assert result == mock_drift_df.to_dict.return_value
     model_metrics_activity.warning.assert_called()
     model_metrics_activity.send_notification_async.assert_called_once_with(
@@ -203,11 +247,15 @@ async def test_calculate_drift_without_reference_data(
         attachment_content=ANY,
     )
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
+    mock_drift_df.rename.assert_called_once_with(
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
-    mock_drift_df.to_dict.assert_called_once()
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -221,11 +269,13 @@ async def test_calculate_drift_empty_drift_df(
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame())
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -255,8 +305,10 @@ async def test_calculate_drift_empty_drift_df(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
+    assert result == []
-    model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata'])
+    model_metrics_activity.warning.assert_called_with(
+        'No drift metrics found', metadata['metadata']
+    )
 
 
 @mark.asyncio
@@ -290,11 +342,13 @@ async def test_calculate_drift_empty_after_timestamp_filter(
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -324,13 +378,13 @@ async def test_calculate_drift_empty_after_timestamp_filter(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
+    assert result == []
     model_metrics_activity.warning.assert_called_with(
         'No drift metrics found after dropping rows where timestamp is not in target data',
-        metadata['metadata']
+        metadata['metadata'],
     )
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
 
 
@@ -342,25 +396,40 @@ async def test_calculate_drift_success_min(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
 
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -390,45 +459,62 @@ async def test_calculate_drift_success_min(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
+    assert isinstance(result, list)
-    assert result == mock_drift_df.to_dict.return_value
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
+    mock_drift_df.rename.assert_called_once_with(
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
-    mock_drift_df.to_dict.assert_called_once()
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
 @patch('laborious.activities.model_metrics.DataFrame')
 @patch('laborious.activities.model_metrics.to_datetime')
-async def test_calculate_drift_success_s(
+async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity):
-    mock_to_datetime, mock_dataframe, model_metrics_activity
-):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
 
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -458,16 +544,20 @@ async def test_calculate_drift_success_s(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
+    assert isinstance(result, list)
-    assert result == mock_drift_df.to_dict.return_value
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
+    mock_drift_df.rename.assert_called_once_with(
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
-    mock_drift_df.to_dict.assert_called_once()
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -479,13 +569,17 @@ async def test_calculate_drift_get_drift_metrics_error(
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
 
-    model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error'))
+    model_metrics_activity.get_drift_metrics = AsyncMock(
+        side_effect=Exception('Get drift metrics error')
+    )
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
@@ -515,10 +609,9 @@ async def test_calculate_drift_get_drift_metrics_error(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
+    assert result == []
     model_metrics_activity.error.assert_called_once_with(
-        'Error getting drift metrics: Get drift metrics error',
+        'Error getting drift metrics: Get drift metrics error', metadata['metadata']
-        metadata['metadata']
     )
     model_metrics_activity.send_notification_async.assert_called_once_with(
         metadata=metadata['metadata'],
@@ -541,28 +634,34 @@ async def test_get_drift_metrics_success(
     # Arrange
     mock_time.return_value = 1000.0
 
-    mock_drift_df = DataFrame({
+    mock_drift_df = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'metric': ['ks_test'],
             'statistic': [0.5],
             'feature': ['feature1'],
-    })
+        }
+    )
 
     mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
     mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
     mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     reference_columns = reference_data.drop(
         columns=['target', 'timestamp'], errors='ignore'
@@ -597,19 +696,25 @@ async def test_get_drift_metrics_univariate_error(
     # Arrange
     mock_time.return_value = 1000.0
 
-    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error')
+    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
+        'Univariate drift error'
+    )
 
-    reference_data = DataFrame({
+    reference_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27'],
             'target': [1.0],
             'feature1': [1.0],
-    })
+        }
+    )
 
     reference_columns = reference_data.drop(
         columns=['target', 'timestamp'], errors='ignore'
@@ -629,27 +734,25 @@ async def test_get_drift_metrics_univariate_error(
     except Exception as e:
         assert str(e) == 'Univariate drift error'
         model_metrics_activity.error.assert_called_once_with(
-            'Error detecting univariate drift: Univariate drift error',
+            'Error detecting univariate drift: Univariate drift error', metadata['metadata']
-            metadata['metadata']
         )
         model_metrics_activity.emit_metric.assert_called_with(
-            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT,
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
-            tags=ANY
         )
     else:
         raise AssertionError('Expected Exception')
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_all_metrics(
+async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
             'target': [1.0, 2.0, 3.0],
             'prediction': [1.1, 2.1, 2.9],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -673,22 +776,22 @@ async def test_calculate_simple_metrics_success_all_metrics(
     assert all(data_size == 3 for data_size in result['data_size'].values)
     assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']',
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
-        metadata['metadata']
+        metadata['metadata'],
     )
     model_metrics_activity.debug.assert_called_once()
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_rmse_only(
+async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
             'target': [1.0, 2.0],
             'prediction': [1.1, 2.1],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -709,21 +812,20 @@ async def test_calculate_simple_metrics_success_rmse_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\']',
+        "Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
-        metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_mse_only(
+async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
             'target': [1.0, 2.0],
             'prediction': [1.1, 2.1],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -744,21 +846,20 @@ async def test_calculate_simple_metrics_success_mse_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'mse\']',
+        "Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
-        metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_mae_only(
+async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
             'target': [1.0, 2.0],
             'prediction': [1.1, 2.1],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -779,21 +880,20 @@ async def test_calculate_simple_metrics_success_mae_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'mae\']',
+        "Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
-        metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_r2_only(
+async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
             'target': [1.0, 2.0],
             'prediction': [1.1, 2.1],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -814,22 +914,21 @@ async def test_calculate_simple_metrics_success_r2_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'r2\']',
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
-        metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_r2_zero_ss_tot(
+async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
     # All target values are the same, so ss_tot will be 0
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
             'target': [1.0, 1.0],
             'prediction': [1.1, 1.1],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -851,21 +950,20 @@ async def test_calculate_simple_metrics_r2_zero_ss_tot(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'r2\']',
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
-        metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_multiple_metrics_subset(
+async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
-    model_metrics_activity
-):
     # Arrange
-    target_data = DataFrame({
+    target_data = DataFrame(
+        {
             'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
             'target': [1.0, 2.0, 3.0],
             'prediction': [1.1, 2.1, 2.9],
-    })
+        }
+    )
 
     input_data = {
         **metadata,
@@ -887,7 +985,5 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(
     assert all(data_size == 3 for data_size in result['data_size'].values)
     assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']',
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
-        metadata['metadata']
     )
-

tests/laborious/utils/repository/test_model_repository.py

@@ -194,7 +194,9 @@ def test_check_artifact_exists_true(mlflow_repository):
     artifact = MagicMock(path='test_artifact')
     mlflow_repository.client.list_artifacts.return_value = [artifact]
 
-    result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
+    result = mlflow_repository.check_artifact_exists(
+        'run_id', 'test_artifact', metadata['metadata']
+    )
 
     assert result is True
     mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -204,7 +206,9 @@ def test_check_artifact_exists_false(mlflow_repository):
     artifact = MagicMock(path='other_artifact')
     mlflow_repository.client.list_artifacts.return_value = [artifact]
 
-    result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
+    result = mlflow_repository.check_artifact_exists(
+        'run_id', 'test_artifact', metadata['metadata']
+    )
 
     assert result is False
     mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -299,16 +303,22 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor
 @patch('laborious.utils.repository.model_repository.mlflow')
 @patch('laborious.utils.repository.model_repository.pd')
 @patch('laborious.utils.repository.model_repository.StringIO')
-async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository):
+async def test_load_artifact_dataframe_success(_stringio, pd, mlflow, mlflow_repository):
     mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
     mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
 
     mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4'
 
-    result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+    result = await mlflow_repository.load_artifact_dataframe(
+        'model_name', 'artifact_path', metadata['metadata']
+    )
 
-    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
+    mlflow_repository.get_model_run_id.assert_called_once_with(
-    mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
+        model_name='model_name', stage='Production'
+    )
+    mlflow_repository.check_artifact_exists.assert_called_once_with(
+        'run_id', 'artifact_path', metadata['metadata']
+    )
     mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path')
     assert result == pd.read_csv.return_value
     mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
@@ -322,11 +332,17 @@ async def test_load_artifact_dataframe_not_exists(mlflow_repository):
     mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
     mlflow_repository.check_artifact_exists = MagicMock(return_value=False)
 
-    result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+    result = await mlflow_repository.load_artifact_dataframe(
+        'model_name', 'artifact_path', metadata['metadata']
+    )
 
     assert result is None
-    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
+    mlflow_repository.get_model_run_id.assert_called_once_with(
-    mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
+        model_name='model_name', stage='Production'
+    )
+    mlflow_repository.check_artifact_exists.assert_called_once_with(
+        'run_id', 'artifact_path', metadata['metadata']
+    )
 
 
 @pytest.mark.asyncio
@@ -337,7 +353,9 @@ async def test_load_artifact_dataframe_error(mlflow, mlflow_repository):
     mlflow.artifacts.load_text.side_effect = ValueError('error')
 
     with pytest.raises(ValueError):
-        await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+        await mlflow_repository.load_artifact_dataframe(
+            'model_name', 'artifact_path', metadata['metadata']
+        )
 
     mlflow_repository.emit_metric.assert_called_once_with(
         metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
@@ -1059,7 +1077,9 @@ async def test_create_new_experiment(
     data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False)
     # Verify prediction_data.to_csv was called with correct arguments
     prediction_data.to_csv.assert_called_once()
-    assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
+    assert (
+        prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
+    )
     assert prediction_data.to_csv.call_args[1]['index'] is False
 
     mlflow.start_run.assert_called_once_with(
@@ -1089,10 +1109,12 @@ async def test_create_new_experiment(
         }
     )
 
-    mlflow.log_artifact.assert_has_calls([
+    mlflow.log_artifact.assert_has_calls(
+        [
             call('./tmp/artifacts/model_name/retrain_data.csv'),
             call('./tmp/artifacts/model_name/evaluation_data.csv'),
-    ])
+        ]
+    )
 
     force_memory_release.assert_called_once_with(mlflow_repository.logger)
 

tests/laborious/workflows/test_drift.py

@@ -1,10 +1,10 @@
 from unittest.mock import ANY, AsyncMock, call, patch
 
 from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from laborious.activities.activities import Activities
 from laborious.workflows.drift import Drift
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @fixture
@@ -45,9 +45,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
     reference_data = {'data': 'test_reference_data'}
     drift_data = {'drift': 'test_drift_data'}
 
-    workflow_mock.start_local_activity_method.side_effect = [
+    workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
-        target_data, reference_data
-    ]
 
     workflow_mock.execute_local_activity_method.return_value = drift_data
     workflow_mock.execute_activity_method = AsyncMock()
@@ -56,12 +54,13 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
     await drift.run(input_data)
 
     # Assert - Check start_local_activity_method calls
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
     expected_gathering_query = f"""
             SELECT *
-            FROM {input_data['schema']}.{input_data['source_table_name']}
+            FROM "{input_data['schema']}"."{input_data['source_table_name']}"
             WHERE
-                model_id = {input_data['model_id']} AND
+                model_id = '{input_data['model_id']}' AND
-                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+                timestamp > NOW() - INTERVAL {input_data['interval']} minutes
             ORDER BY timestamp ASC
         """
 
@@ -73,6 +72,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
                     **metadata,
                     'query': expected_gathering_query,
                     'datetime_columns': ['timestamp', 'created_at'],
+                    'orient': 'records',
                 },
                 retry_policy=ANY,
                 start_to_close_timeout=ANY,
@@ -250,4 +250,3 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
         retry_policy=ANY,
         start_to_close_timeout=ANY,
     )
-

tests/laborious/workflows/test_simple_metrics.py

@@ -1,10 +1,10 @@
 from unittest.mock import ANY, AsyncMock, call, patch
 
 from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from laborious.activities.activities import Activities
 from laborious.workflows.simple_metrics import SimpleMetrics
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @fixture
@@ -42,9 +42,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
     target_data = {'data': 'test_target_data'}
     simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
 
-    workflow_mock.execute_local_activity_method.side_effect = [
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
-        target_data, simple_metrics_data
-    ]
 
     workflow_mock.execute_activity_method = AsyncMock()
 
@@ -52,17 +50,18 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
     await simple_metrics.run(input_data)
 
     # Assert - Check load_custom_query call
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
     expected_query = f"""
             select p."timestamp", p.prediction, ld.value as "target"
-            from {input_data['schema']}.{input_data['predictions_table_name']} p
+            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
-            inner join {input_data['schema']}.{input_data['data_table_name']} ld
+            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
                 on p."timestamp" = ld."timestamp"
             where
-                p.model_id = {input_data['model_id']} and
+                p.model_id = '{input_data['model_id']}' and
                 p.prediction is not null and
                 ld.variable = '{input_data['model_config']['target']}' and
                 ld.value is not null and
-                p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
+                p."timestamp" >= NOW() - INTERVAL {input_data['interval_minutes']} minutes
             order by
                 p."timestamp" desc;
         """
@@ -75,6 +74,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
                     **metadata,
                     'query': expected_query,
                     'datetime_columns': ['timestamp'],
+                    'orient': 'records',
                 },
                 retry_policy=ANY,
                 start_to_close_timeout=ANY,
@@ -159,9 +159,7 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics
     target_data = {'data': 'test_target_data'}
     simple_metrics_data = None
 
-    workflow_mock.execute_local_activity_method.side_effect = [
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
-        target_data, simple_metrics_data
-    ]
 
     workflow_mock.execute_activity_method = AsyncMock()
 
@@ -193,9 +191,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
     target_data = {'data': 'test_target_data'}
     simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
 
-    workflow_mock.execute_local_activity_method.side_effect = [
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
-        target_data, simple_metrics_data
-    ]
 
     workflow_mock.execute_activity_method = AsyncMock()
 
@@ -225,4 +221,3 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
             ),
         ]
     )
-

validate.sh

@@ -87,7 +87,7 @@ if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then
 fi
 
 # Step 4: Security Analysis (Bandit)
-if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then
+if ! run_step "4. Security Analysis (Bandit)" "bandit -c pyproject.toml -r laborious/ -ll -q"; then
     FAILED_STEPS+=("Security Analysis")
 fi