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

orchestrator/utils/connectors_config.py

@@ -1,5 +1,8 @@
 from os import getenv
 
+from debugpy import connect
+from matplotlib.pylab import f
+
 
 def build_redis_config():
     return {
@@ -10,6 +13,18 @@ def build_redis_config():
     }
 
 
+def build_mongodb_config():
+    username = getenv('MONGODB_USERNAME', 'sientia')
+    password = getenv('MONGODB_PASSWORD', 'sientia')
+    uri = getenv('MONGODB_URL', 'localhost:27017')
+
+    connection_string = f'mongodb://{username}:{password}@{uri}'
+    return {
+        'connection_string': connection_string,
+        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia')
+    }
+
+
 def build_couchbase_config():
     return {
         'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'),

orchestrator/worker/worker.py

@@ -8,8 +8,9 @@ with workflow.unsafe.imports_passed_through():
     from orchestrator.workflows.orchestrator import Orchestrator
     from orchestrator.activities.activities import Activities
     from orchestrator.utils.connectors_config import (
-        build_couchbase_config,
+        # build_couchbase_config,
         build_redis_config,
+        build_mongodb_config
     )
     from sientia_do.notifications.handlers import NotificationHandler
     from sientia_do.temporal.utils.logger import get_logger
@@ -40,8 +41,9 @@ async def main():
 
     activities = Activities(
         temporal_client=temporal_client,
-        couchbase_config=build_couchbase_config(),
+        # couchbase_config=build_couchbase_config(),
         redis_config=build_redis_config(),
+        mongodb_config=build_mongodb_config(),
         logger=logger,
         notification_handler=notification_handler
     )
@@ -62,7 +64,10 @@ async def main():
                 activities.update_slots,
                 activities.delete_slots,
                 # Couchbase
-                activities.load_query_from_couchbase,
+                # activities.load_query_from_couchbase,
+                # MongoDB
+                activities.aggregate_documents_in_mongodb,
+                activities.find_documents_in_mongodb,
                 # Temporal
                 activities.load_schedule,
                 activities.create_schedules,

orchestrator/workflows/orchestrator.py

@@ -26,8 +26,17 @@ class Orchestrator:
             start_to_close_timeout=timedelta(seconds=60)
         )
 
+        # pipeline_config_handler = workflow.execute_local_activity_method(
+        #     Activities.load_query_from_couchbase,
+        #     {
+        #         'query': input_data['pipelines_query']
+        #     },
+        #     retry_policy=retry_policy,
+        #     start_to_close_timeout=timedelta(seconds=60)
+        # )
+
         pipeline_config_handler = workflow.execute_local_activity_method(
-            Activities.load_query_from_couchbase,
+            Activities.aggregate_documents_in_mongodb,
             {
                 'query': input_data['pipelines_query']
             },
@@ -35,8 +44,17 @@ class Orchestrator:
             start_to_close_timeout=timedelta(seconds=60)
         )
 
+        # opc_servers_handler = workflow.execute_local_activity_method(
+        #     Activities.load_query_from_couchbase,
+        #     {
+        #         'query': input_data['opc_servers_query']
+        #     },
+        #     retry_policy=retry_policy,
+        #     start_to_close_timeout=timedelta(seconds=60)
+        # )
+
         opc_servers_handler = workflow.execute_local_activity_method(
-            Activities.load_query_from_couchbase,
+            Activities.find_documents_in_mongodb,
             {
                 'query': input_data['opc_servers_query']
             },