import json import os import traceback from time import sleep from kafka import KafkaProducer from kafka.errors import NoBrokersAvailable 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.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.repository.mongodb_repository import MongoDBRepository import ingestor.metrics as metrics class DataManager(SientiaMonitoring): """ Manages data persistence and export operations for the OPC Ingestor. The DataManager is responsible for: - Storing OPC data in MongoDB for historical analysis and persistence - Exporting data to Kafka for real-time streaming and downstream processing - Managing database connections and ensuring data integrity - Providing data access interfaces for other components The manager supports both MongoDB and Kafka operations, with Kafka export being optional and configurable. It implements retry logic for connection failures and provides comprehensive error handling and notification. Args: kafka_servers (str): Comma-separated string of Kafka server addresses mongo_connection_string (str): MongoDB connection string mongo_database (str): MongoDB database name export_to_kafka (bool): Whether to enable Kafka export functionality metadata (dict): Application metadata for notifications and tracking logger (Logger): Logger instance for application logging notification_handler (NotificationHandler): Handler for sending notifications Attributes: pod_id (str): Pod identifier for metrics labeling kafka_producer (KafkaProducer): Kafka producer instance for data export export_to_kafka (bool): Whether Kafka export is enabled connection_string (str): MongoDB connection string database (str): MongoDB database name mongo_client (MongoClient): MongoDB client instance metadata (dict): Application metadata """ def __init__( self, kafka_servers: str, mongo_connection_string: str, mongo_database: str, export_to_kafka: bool, metadata: dict, logger: Logger, notification_handler: NotificationHandler, metrics_controller: MetricsController, ) -> None: """ Initializes the DataManager instance with Kafka and MongoDB connections. This constructor attempts to establish connections to the specified services: 1. Kafka: Initializes producer with retry logic (up to 3 attempts) 2. MongoDB: Establishes connection and verifies server availability The initialization process includes: - Kafka producer setup with JSON serialization - MongoDB client initialization and connection testing - Metrics recording for connection status - Error handling with notifications Args: kafka_servers (str): Comma-separated string of Kafka server addresses mongo_connection_string (str): MongoDB connection string mongo_database (str): MongoDB database name export_to_kafka (bool): Whether to enable Kafka export metadata (dict): Application metadata logger (Logger): Logger instance notification_handler (NotificationHandler): Notification handler Raises: NoBrokersAvailable: If the connection to Kafka servers fails after 3 attempts. Metrics: - KAFKA_CONNECTION_STATUS: Set to 1 on successful connection, 0 on failure """ self.pod_id = os.getenv('HOSTNAME', 'localhost') self.kafka_producer = None self.export_to_kafka = export_to_kafka SientiaMonitoring.__init__( self, logger=logger, metrics_controller=metrics_controller, notification_handler=notification_handler, ) if self.export_to_kafka: for i in range(0, 3): logger.info( f'Trying ({i}) to initializing DataManager with Kafka servers: {kafka_servers}' ) try: self.kafka_producer = KafkaProducer( bootstrap_servers=kafka_servers, value_serializer=lambda v: json.dumps(v).encode( 'utf-8' ), # Serialize JSON messages key_serializer=lambda k: str(k).encode('utf-8') if k else None, ) # Kafka connected metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(1) break except NoBrokersAvailable: logger.error(f'Kafka servers {kafka_servers} are not available. Retrying...') sleep(5) else: # Kafka not connected metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0) logger.error( f'Failed to connect to Kafka servers {kafka_servers} after 3 attempts.' ) raise NoBrokersAvailable( f'Failed to connect to Kafka servers {kafka_servers} after 3 attempts.' ) logger.info(f'DataManager initialized with Kafka servers: {kafka_servers}') logger.info( f'Trying to initializing DataManager with MongoDB servers: {mongo_connection_string}' ) self.connection_string = mongo_connection_string self.database = mongo_database self.mongo_repository = MongoDBRepository( connection_string=self.connection_string, database_name=self.database, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) self.metadata = metadata logger.info(f'DataManager initialized with MongoDB servers: {self.connection_string}') def shutdown(self): """ Gracefully shuts down the DataManager and closes all connections. This method ensures proper cleanup of: - Kafka producer connection with message flushing - MongoDB client connection - Metrics recording for connection status The method handles connection closure gracefully, logging any errors that occur during shutdown while ensuring all resources are properly released. """ if self.kafka_producer: try: self.kafka_producer.flush(timeout=10) self.kafka_producer.close() # Mark as disconnected metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0) except Exception as e: self.logger.error(f'Error closing Kafka producer: {e}') else: self.logger.warning('Kafka producer is already closed or not initialized.') try: self.mongo_repository.close() except Exception as e: self.logger.error(f'Error closing MongoDB client: {e}') def __del__(self): self.shutdown() def delivery_report(self, msg): """ Callback for successful Kafka message delivery reports. This method is called by the Kafka producer when a message is successfully delivered to a topic. It logs the delivery details including topic, partition, and offset information for debugging and monitoring purposes. Args: msg: Kafka message object containing delivery details """ self.logger.debug( f'Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}' ) def delivery_error(self, err): """ Callback for Kafka message delivery error reports. This method is called by the Kafka producer when a message delivery fails. It logs the error details for debugging and monitoring purposes. Args: err: Error information from the failed delivery attempt """ self.logger.error(f'Delivery failed for record : {err}') async def publish(self, topic: str, data: dict) -> None: """ Publishes a message to a specified Kafka topic. Args: topic (str): The name of the Kafka topic to which the message will be published. data (dict): The message data to be sent to the Kafka topic. Returns: None Raises: Exception: If there is an error during message delivery, it will be handled by the `delivery_error` callback. """ if self.export_to_kafka and self.kafka_producer: try: self.logger.debug(f'Publishing message to topic {topic}: {data}') self.kafka_producer.send(topic=topic, value=data).add_callback( self.delivery_report ).add_errback(self.delivery_error) self.kafka_producer.flush(timeout=10) await self.emit_metric( metric_object=metrics.KAFKA_MESSAGES_SENT, tags={ 'pod_id': self.pod_id, 'topic': topic, }, ) except Exception as e: await self.emit_metric( metric_object=metrics.KAFKA_MESSAGES_ERRORS, tags={ 'pod_id': self.pod_id, 'topic': topic, }, ) trace = traceback.format_exc() await self.send_notification_async( metadata=self.metadata, notification_id=f'KAFKA_PRODUCER_ERROR_{topic}', message=f'Error publishing message to topic {topic}: {e}', block='kafka_producer', level=NotificationLevel.ERROR, attachment_content=trace, ) self.logger.error(trace) try: await self.mongo_repository.insert( collection_name=topic, document=data, metadata=self.metadata, ) self.logger.debug(f'Message inserted into MongoDB collection {topic}: {data}') await self.emit_metric( metric_object=metrics.TAG_WRITTEN_COUNT, tags={ 'pod_id': self.pod_id, 'tag_name': data['name'], 'collection_name': topic, }, ) except Exception as e: trace = traceback.format_exc() await self.send_notification_async( metadata=self.metadata, notification_id=f'MONGO_PRODUCER_ERROR_{topic}', message=f'Error inserting message to MongoDB: {e}', block='mongo_producer', level=NotificationLevel.ERROR, attachment_content=trace, ) self.logger.error(trace)