from datetime import UTC from sientia_do.observability.metrics_controller import MetricsController from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): import traceback from datetime import datetime from typing import Any 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.observability.sientia_monitoring import SientiaMonitoring from sientia_do.repository.mongodb_repository import MongoDBRepository 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, metrics_controller: MetricsController, ): """ 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, metrics_controller=metrics_controller, ) SientiaMonitoring.__init__( self, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) 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]) -> list[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 = await 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 data except Exception as e: trace = traceback.format_exc() await self.send_notification_async( 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