SIENTIAPDE-1107

feat: add MongoDB integration and update configuration for activities

orchestrator/activities/activities.py

@@ -1,6 +1,8 @@
 from temporalio import activity, workflow
 from temporalio.client import Client
 
+from orchestrator.activities.mongo_db import MongoDB
+
 with workflow.unsafe.imports_passed_through():
     from orchestrator.activities.couchbase import Couchbase
     from orchestrator.activities.temporal_manager import TemporalManager
@@ -11,21 +13,23 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.notifications.handlers import NotificationHandler
 
 
-class Activities(Couchbase, TemporalManager, SlotManager, Formatters):
+class Activities(  # Couchbase,
+        TemporalManager, SlotManager, Formatters, MongoDB):
 
     def __init__(self,
                  temporal_client: Client,
-                 couchbase_config: dict[str, Any],
+                 # couchbase_config: dict[str, Any],
                  redis_config: dict[str, Any],
+                 mongodb_config: dict[str, Any],
                  logger: Logger,
                  notification_handler: NotificationHandler):
 
         # Initialize parent classes
-        Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
-                           username=couchbase_config['username'],
-                           password=couchbase_config['password'],
-                           logger=logger,
-                           notification_handler=notification_handler)
+        # Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
+        #                    username=couchbase_config['username'],
+        #                    password=couchbase_config['password'],
+        #                    logger=logger,
+        #                    notification_handler=notification_handler)
 
         TemporalManager.__init__(self,
                                  temporal_client=temporal_client,
@@ -44,9 +48,15 @@ class Activities(Couchbase, TemporalManager, SlotManager, Formatters):
                             logger=logger,
                             notification_handler=notification_handler)
 
+        MongoDB.__init__(self,
+                         connection_string=mongodb_config['connection_string'],
+                         database_name=mongodb_config['database_name'],
+                         logger=logger,
+                         notification_handler=notification_handler)
+
     @activity.defn(name="prepare_activity")
     async def prepare_activity(self, input_data: dict[str, Any]):
         await super().prepare_activity(input_data)
 
     def shutdown(self):
-        Couchbase.shutdown(self)
+        MongoDB.shutdown(self)

orchestrator/activities/mongo_db.py

@@ -0,0 +1,198 @@
+from ast import Not
+from pydoc import doc
+import trace
+from grpc import server
+import pymongo
+from temporalio import workflow, activity
+
+with workflow.unsafe.imports_passed_through():
+    from typing import Any
+    import traceback
+    import json
+    from logging import Logger
+    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 not isinstance(doc, dict) and not isinstance(doc, list):
+            continue
+
+        if "_id" in doc:
+            del doc["_id"]
+
+        if isinstance(doc, dict):
+            for _key, value in doc.items():
+                if isinstance(value, dict):
+                    clear_mongo_id([value])
+                elif isinstance(value, list):
+                    clear_mongo_id(value)
+
+        if isinstance(doc, list):
+            for item in doc:
+                if isinstance(item, dict):
+                    clear_mongo_id([item])
+                elif isinstance(item, list):
+                    clear_mongo_id(item)
+
+    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 = pymongo.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)
+        self.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="find_documents_in_mongodb",)
+    async def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Find documents in a MongoDB collection based on the provided query parameters.
+
+        Args:
+            - input_data (dict): Input data containing query parameters. Contains:
+                - query (dict): Query parameters to filter documents.
+
+        Returns:
+            list[dict]: List of documents matching the query.
+        """
+
+        query = input_data.get("query", {})
+
+        collection_name = query.get("collection")
+        if not collection_name:
+            raise ValueError("Collection name must be provided in the query.")
+
+        filters = query.get("filters", {})
+
+        self.logger.info(
+            f"Loading documents from collection '{collection_name}' with filters: {filters}")
+
+        try:
+            collection = self.database[collection_name]
+
+            documents = list(collection.find(filters, {"_id": 0}))
+
+            documents = clear_mongo_id(documents)
+
+            self.logger.info(
+                f"Loaded {len(documents)} documents from collection '{collection_name}'")
+
+            self.logger.debug(
+                f"Documents loaded: {documents}"
+            )
+
+            return documents
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.notification_handler.build_and_send_notification(
+                notification_id="MONGODB_QUERY_ERROR",
+                message=f"Failed to execute MongoDB query: {e}",
+                block="load_query_from_mongodb",
+                level=NotificationLevel.ERROR,
+                attachment_content=trace
+            )
+            self.logger.error(trace)
+
+            raise e
+
+    @activity.defn(name="aggregate_documents_in_mongodb")
+    async def aggregate_documents_in_mongodb(self,
+                                             input_data: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Aggregate documents in a MongoDB collection based on the provided aggregation pipeline.
+
+        Args:
+            - input_data (dict): Input data containing aggregation parameters. Contains:
+                - query (dict): Query parameters to filter documents.
+
+        Returns:
+            list[dict]: List of aggregated documents.
+        """
+
+        query = input_data.get("query", {})
+
+        collection_name = query.get("collection")
+        if not collection_name:
+            raise ValueError("Collection name must be provided in the query.")
+        aggregation = query.get("aggregation")
+        if not aggregation:
+            raise ValueError("Aggregation must be provided.")
+        aggregation.append({"$project": {"_id": 0}})
+
+        self.logger.info(
+            f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}")
+
+        try:
+            collection = self.database[collection_name]
+
+            aggregated_documents = list(
+                collection.aggregate(aggregation))
+
+            aggregated_documents = clear_mongo_id(aggregated_documents)
+
+            self.logger.info(
+                f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'")
+
+            self.logger.debug(
+                f"Aggregation result: {aggregated_documents}"
+            )
+
+            return aggregated_documents
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.notification_handler.build_and_send_notification(
+                notification_id="MONGODB_AGGREGATION_ERROR",
+                message=f"Failed to execute MongoDB aggregation: {e}",
+                block="aggregate_documents_in_mongodb",
+                level=NotificationLevel.ERROR,
+                attachment_content=trace
+            )
+            self.logger.error(trace)
+
+            raise e