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

from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from typing import Any
    import traceback
    from datetime import datetime, timezone
    from pymongo import MongoClient
    from pandas import DataFrame
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.observability.logger import Logger
    from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ


def clear_mongo_id(docs: list) -> list:
    """
    Remove MongoDB internal `_id` fields from documents.

    This utility function recursively removes the MongoDB `_id` field from
    documents and nested structures. It's used to clean data before
    processing or export operations.

    Args:
        docs (list): List of documents to clean

    Returns:
        list: Documents with `_id` fields removed

    Note:
        This function modifies the input list in-place and returns the same reference
    """
    for doc in docs:
        if isinstance(doc, list):
            clear_mongo_id(doc)

        elif isinstance(doc, dict):
            if "_id" in doc:
                del doc["_id"]

            for key, value in doc.items():
                if isinstance(value, list):
                    clear_mongo_id(value)
                elif isinstance(value, dict):
                    clear_mongo_id([value])

    return docs


class MongoDB(BaseActivity):
    """
    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.connection_string = connection_string
        self.database_name = database_name

        self.client = MongoClient(
            self.connection_string, serverSelectionTimeoutMS=5000)
        self.client.server_info()  # Trigger an exception if connection fails

        self.database = self.client[self.database_name]

        # Initialize MongoDB client here (omitted for brevity)
        logger.info("MongoDB connection initialized")

        BaseActivity.__init__(self,
                              logger=logger,
                              notification_handler=notification_handler,
                              set_error_counter=True)

    def shutdown(self):
        """
        Gracefully close MongoDB client connection.

        This method ensures proper cleanup of MongoDB connections to prevent
        connection leaks and ensure graceful application termination.
        """
        try:
            if self.client:
                self.logger.info("Closing MongoDB connection...")
                self.client.close()
                self.logger.info("MongoDB connection closed successfully")
        except Exception as e:
            self.logger.error(f"Failed to close MongoDB connection: {e}")

    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.shutdown()

    @activity.defn(name="load_latest_data")
    async def load_latest_data(self, input_data: dict[str, Any]) -> dict[str, 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 = list(self.database[collection_name].find(
                data_filter, {"_id": 0}))

            data = clear_mongo_id(data)

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

            for item in data:
                item['inserted_at'] = item['inserted_at'].replace(
                    tzinfo=timezone.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