SIENTIAPDE-1110

SIENTIAPDE-1110 Add MongoDB integration and enhance Redis activity with timestamp management functions.

scouter/activities/mongodb.py

@@ -0,0 +1,121 @@
+from temporalio import workflow, activity
+
+with workflow.unsafe.imports_passed_through():
+    from typing import Any
+    import traceback
+    from logging import Logger
+    import datetime
+    from pymongo import MongoClient
+    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.models import NotificationLevel
+    from sientia_do.temporal.activities.base import BaseActivity
+
+
+def clear_mongo_id(docs: list) -> list:
+    """
+    Remove the MongoDB internal `_id` field from the document.
+
+    Args:
+        docs (list): The document to clear.
+
+    Returns:
+        list: The documents without the `_id` field.
+    """
+    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):
+    def __init__(self, connection_string: str, database_name: str,
+                 logger: Logger,
+                 notification_handler: NotificationHandler):
+        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)
+
+    def shutdown(self):
+        """
+        Close the MongoDB client connection.
+        """
+        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 when the object is deleted.
+        """
+        self.shutdown()
+
+    @activity.defn(name="load_latest_data")
+    async def load_latest_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Loads the latest data from MongoDB.
+        """
+        metadata = input_data['metadata']
+        collection_name = input_data['collection_name']
+        last_data_timestamp = input_data['last_data_timestamp']
+
+        try:
+
+            data_filter = {
+                "inserted_at": {
+                    "$gt": datetime.strptime(last_data_timestamp, "%Y-%m-%d %H:%M:%S.%f")
+                }
+            }
+
+            data = self.database[collection_name].find(data_filter, {"_id": 0})
+
+            data = clear_mongo_id(data)
+
+            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()
+            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