sientia-dataops-scouter_temporal/scouter/activities/mongodb.py

from collections.abc import Hashable
from datetime import UTC

from temporalio import activity, workflow

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

    from pandas import DataFrame
    from pymongo import MongoClient
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.repository.mongodb_repository import MongoDBRepository
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ



class MongoDB(SientiaMonitoring):
    """
    MongoDB operations for data retrieval and storage.

    This class provides MongoDB connectivity and operations for the Scouter system.
    It handles:
    - Connection management with automatic reconnection
    - Data retrieval with timestamp-based filtering
    - Document cleaning and preprocessing
    - Error handling and notification integration

    The class implements Temporal activities for MongoDB operations, enabling
    distributed data processing with fault tolerance and monitoring.
    """

    def __init__(
        self,
        connection_string: str,
        database_name: str,
        logger: Logger,
        notification_handler: NotificationHandler,
    ):
        """
        Initialize MongoDB connection and services.

        Args:
            connection_string (str): MongoDB connection URI string
            database_name (str): Name of the target database
            logger (Logger): Logger instance for operation logging
            notification_handler (NotificationHandler): Handler for system notifications

        Raises:
            ConnectionError: If MongoDB connection fails
        """
        
        self.mongodb_repository = MongoDBRepository(
            connection_string=connection_string,
            database_name=database_name,
            logger=logger,
            notification_handler=notification_handler,
        )

        SientiaMonitoring.__init__(self, logger=logger, notification_handler=notification_handler)

    def close(self):
        """
        Close the MongoDB connection.
        """
        self.mongodb_repository.close()
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        """
        Destructor to ensure MongoDB client is closed.

        This destructor ensures that MongoDB connections are properly closed
        when the object is garbage collected, preventing resource leaks.
        """
        self.close()

    @activity.defn(name='load_latest_data')
    async def load_latest_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
        """
        Load the latest data from MongoDB collection since a specified timestamp.

        This activity retrieves data from a MongoDB collection, optionally
        filtering by timestamp to enable incremental data processing. It
        handles connection management and provides comprehensive error reporting.

        Args:
            input_data (dict[str, Any]): Activity input parameters.
                Required fields:
                - metadata (dict[str, Any]): Workflow execution metadata
                - collection_name (str): Name of the MongoDB collection
                - last_data_timestamp (str | None): Last processed timestamp for filtering

        Returns:
            dict[str, Any]: Retrieved data, or empty dict if no data found

        Raises:
            Exception: If MongoDB operation fails
        """
        metadata = input_data['metadata']
        collection_name = input_data['collection_name']
        last_data_timestamp = input_data['last_data_timestamp']

        self.info(f'Loading data from MongoDB: {input_data}', metadata=metadata)

        try:
            if last_data_timestamp is None:
                data_filter = {}
            else:
                data_filter = {
                    'inserted_at': {
                        '$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ)
                    }
                }

            self.debug(f'Data filter: {data_filter}', metadata=metadata)

            data = self.mongodb_repository.find(
                collection_name=collection_name,
                filters=data_filter,
                metadata=metadata,
            )

            self.debug(f'Collected: {data}', metadata=metadata)

            for item in data:
                item['inserted_at'] = (
                    item['inserted_at'].replace(tzinfo=UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ)
                )

            self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata)

            self.debug(f'Loaded data: {data}', metadata=metadata)

            return DataFrame(data).to_dict()
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='MONGO_LOAD_ERROR',
                message=f'Error loading data from MongoDB: {e}',
                block='load_latest_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise e