import json from time import sleep from pymongo import MongoClient 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.temporal.constants import now from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.logger import Logger import traceback import ingestor.metrics as metrics import os class DataManager(BaseActivity): """ 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, ) -> 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 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_client = MongoClient(self.connection_string) self.mongo_client.server_info() self.metadata = metadata self.mongo_db = self.mongo_client[self.database] logger.info( f"DataManager initialized with MongoDB servers: {self.connection_string}" ) BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler, set_error_counter=True) def shutdown(self): """Closes the Kafka producer connection.""" 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.") if self.mongo_client: try: self.mongo_client.close() except Exception as e: self.logger.error(f"Error closing MongoDB client: {e}") else: self.logger.warning( "MongoDB client is already closed or not initialized.") def __del__(self): self.shutdown() def delivery_report(self, msg: str): """Callback for delivery reports from Kafka.""" self.logger.debug( f"Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}" ) def delivery_error(self, err: str): """Callback for delivery reports from Kafka.""" self.logger.error(f"Delivery failed for record : {err}") 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: 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) metrics.KAFKA_MESSAGES_SENT.labels( pod_id=self.pod_id, topic=topic).inc() except Exception as e: metrics.KAFKA_MESSAGES_ERRORS.labels( pod_id=self.pod_id, topic=topic).inc() trace = traceback.format_exc() self.send_notification( 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: collection = self.mongo_db[topic] collection.insert_one( { **data, "inserted_at": now(), } ) self.logger.debug( f"Message inserted into MongoDB collection {topic}: {data}") metrics.TAG_WRITTEN_COUNT.labels( pod_id=self.pod_id, tag_name=data["name"], collection_name=topic ).inc() except Exception as e: trace = traceback.format_exc() self.send_notification( 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)