Add MongoDB integration to DataManager, switch Python language server to Jedi, and include pymongo in requirements.txt

ingestor/managers/data_manager.py

@@ -1,6 +1,8 @@
 import json
 from logging import Logger
 from time import sleep
+from datetime import datetime, timezone
+from pymongo import MongoClient
 from kafka import KafkaProducer
 from kafka.errors import NoBrokersAvailable
 from sientia_do.notifications.handlers import NotificationHandler
@@ -14,6 +16,8 @@ class DataManager:
     def __init__(
         self,
         kafka_servers: str,
+        mongo_connection_string: str,
+        mongo_database: str,
         logger: Logger,
         notification_handler: NotificationHandler,
     ) -> None:
@@ -41,10 +45,12 @@ class DataManager:
                     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,
+                    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)
+                metrics.KAFKA_CONNECTION_STATUS.labels(
+                    pod_id=self.pod_id).set(1)
                 break
             except NoBrokersAvailable:
                 logger.error(
@@ -61,7 +67,34 @@ class DataManager:
                 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"DataManager initialized with Kafka servers: {kafka_servers}")
+
+        self.connection_string = mongo_connection_string
+        self.database = mongo_database
+
+        for i in range(0, 3):
+            logger.info(
+                f"Trying ({i}) to initializing DataManager with MongoDB servers: {self.connection_string}"
+            )
+            try:
+                self.mongo_client = MongoClient(self.connection_string)
+                self.mongo_client.server_info()
+
+                self.mongo_db = self.mongo_client[self.database]
+                break
+            except Exception as e:
+                logger.error(f"Error connecting to MongoDB: {e}")
+                sleep(5)
+        else:
+            raise ValueError(
+                f"Failed to connect to MongoDB servers {self.connection_string} after 3 attempts."
+            )
+
+        logger.info(
+            f"DataManager initialized with MongoDB servers: {self.connection_string}"
+        )
+
         self.logger = logger
         self.notification_handler = notification_handler
 
@@ -72,11 +105,19 @@ class DataManager:
                 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)
+                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.")
+            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}")
 
     def __del__(self):
         self.shutdown()
@@ -114,10 +155,12 @@ class DataManager:
             ).add_errback(self.delivery_error)
 
             self.kafka_producer.flush(timeout=10)
-            metrics.KAFKA_MESSAGES_SENT.labels(pod_id=self.pod_id, topic=topic).inc()
+            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()
+            metrics.KAFKA_MESSAGES_ERRORS.labels(
+                pod_id=self.pod_id, topic=topic).inc()
             trace = traceback.format_exc()
             self.notification_handler.build_and_send_notification(
                 notification_id=f"KAFKA_PRODUCER_ERROR_{topic}",
@@ -127,3 +170,24 @@ class DataManager:
                 attachment_content=trace,
             )
             self.logger.error(trace)
+
+        try:
+            collection = self.mongo_db[topic]
+
+            collection.insert_one(
+                {
+                    **data,
+                    "inserted_at": datetime.now(timezone.utc),
+                }
+            )
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.notification_handler.build_and_send_notification(
+                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)