sientia-dataops-laborious-legacy/laborious/activities/activities.py

from sientia_do.observability.metrics_controller import MetricsController
from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from typing import Any

    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import Logger
    from sientia_do.repository.minio_repository import MinioRepository
    from sientia_do.repository.mongodb_repository import MongoDBRepository

    from laborious.activities.api import API
    from laborious.activities.gates import Gates
    from laborious.activities.mlflow import MLFlow
    from laborious.activities.model_import import ModelImport
    from laborious.activities.model_metrics import ModelMetrics
    from laborious.activities.opc import OPC
    from laborious.activities.storage import Storage
    from laborious.utils.connectors_config import (
        build_import_config,
        build_import_status_config,
    )
    from laborious.utils.repository.model_repository import MLFlowRepository


class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
    """
    Main activities orchestrator for the Laborious system.

    This class combines functionality from multiple activity classes to provide
    a unified interface for all workflow operations. It manages database connections,
    MLFlow model interactions, data quality validation, and OPC server communications.

    The class implements multiple inheritance to combine specialized functionality:
    - Storage: Database operations and data persistence
    - MLFlow: Model inference and transformation operations
    - Gates: Data quality validation and filtering mechanisms
    - OPC: Real-time data export to OPC servers
    - ModelMetrics: Model performance metrics and drift detection
    - API: PI Web API export operations for industrial systems

    Attributes:
        postgres_config (dict): PostgreSQL connection configuration
        mlflow_config (dict): MLFlow server configuration
        opc_config (dict): OPC server configuration
        pi_web_api_config (dict): PI Web API server configuration
        logger (Logger): Logging and observability instance
        notification_handler (NotificationHandler): Notification management instance
    """

    def __init__(
        self,
        postgres_config: dict[str, Any],
        mlflow_config: dict[str, Any],
        minio_config: dict[str, Any],
        opc_config: dict[str, Any],
        pi_web_api_config: dict[str, Any],
        logger: Logger,
        notification_handler: NotificationHandler,
        mongo_config: dict[str, Any] | None = None,
        import_config: dict[str, Any] | None = None,
        import_status_config: dict[str, Any] | None = None,
    ):
        """
        Initialize the Activities orchestrator with all required configurations.

        This constructor initializes all parent classes with their respective
        configurations and sets up the foundation for all activity operations.

        Args:
            postgres_config: PostgreSQL connection configuration dictionary
                Required keys: host, port, user, password, dbname, min_connections, max_connections
            mlflow_config: MLFlow server configuration dictionary
                Required keys: host, port, username, password
            opc_config: OPC server configuration dictionary
                Can contain multiple server configurations
            pi_web_api_config: PI Web API server configuration dictionary
                Required keys: base_url, auth_type, auth_token
            logger: Logger instance for observability and debugging
            notification_handler: Notification handler for alerts and monitoring
            mongo_config: MongoDB configuration for the model listing the import writes. When
                absent, the import activities that need it refuse rather than guess — the
                notification handler keeps its own client either way.
            import_config: `build_import_config()`; read from the environment when absent
            import_status_config: `build_import_status_config()`; read from the environment when
                absent. This is the BFF database holding the import log: another database on the
                `POSTGRES_*` server, and a second connection that is never the one Storage holds.

        Raises:
            Exception: If any parent class initialization fails
        """
        metrics_controller = MetricsController(logger=logger)

        minio_repository = MinioRepository(
            endpoint=minio_config['endpoint_url'],
            access_key=minio_config['access_key'],
            secret_key=minio_config['secret_key'],
            bucket=minio_config['default_bucket'],
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
            secure=minio_config['secure'],
        )

        mlflow_repository = MLFlowRepository(
            f'{mlflow_config["host"]}:{mlflow_config["port"]}',
            mlflow_config['username'],
            mlflow_config['password'],
            logger,
            notification_handler,
            metrics_controller,
        )

        # Initialize parent classes
        Storage.__init__(
            self,
            host=postgres_config['host'],
            port=postgres_config['port'],
            user=postgres_config['user'],
            password=postgres_config['password'],
            dbname=postgres_config['dbname'],
            min_connections=postgres_config['min_connections'],
            max_connections=postgres_config['max_connections'],
            retention_hours=minio_config['retention_hours'],
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        MLFlow.__init__(
            self,
            mlflow_repository=mlflow_repository,
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        Gates.__init__(
            self,
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        OPC.__init__(
            self,
            opc_servers=opc_config,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        ModelMetrics.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        API.__init__(
            self,
            base_url=pi_web_api_config['base_url'],
            auth_type=pi_web_api_config['auth_type'],
            auth_token=pi_web_api_config['auth_token'],
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        mongo_repository = (
            MongoDBRepository(
                connection_string=mongo_config['connection_string'],
                database_name=mongo_config['database_name'],
                logger=logger,
                notification_handler=notification_handler,
                metrics_controller=metrics_controller,
            )
            if mongo_config
            else None
        )

        ModelImport.__init__(
            self,
            status_config=import_status_config or build_import_status_config(),
            import_config=import_config or build_import_config(),
            minio_repository=minio_repository,
            mlflow_repository=mlflow_repository,
            mongo_repository=mongo_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    async def shutdown(self):
        """
        Gracefully shutdown all activities and clean up resources.

        This method ensures proper cleanup of all resources including:
        - PostgreSQL connection pools
        - OPC server connections
        - PI Web API client connections
        - MLFlow model repositories
        - Any other resources that need explicit cleanup

        The method should be called before the application terminates to ensure
        proper resource cleanup and prevent resource leaks.
        """
        Storage.close(self)
        MLFlow.close(self)
        Gates.close(self)
        await OPC.aclose(self)
        ModelMetrics.close(self)
        API.close(self)
        ModelImport.close(self)