SIENTIAPDE-1110

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

requirements.txt

@@ -4,5 +4,6 @@ sqlalchemy
 asyncua
 redis
 aiokafka
+pymongo
 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.0
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.1

scouter/activities/activities.py

@@ -7,16 +7,18 @@ with workflow.unsafe.imports_passed_through():
     from scouter.activities.redis import Redis
     from scouter.activities.kafka import Kafka
     from scouter.activities.gates import Gates
+    from scouter.activities.mongodb import MongoDB
     from typing import Any
 
 
-class Activities(Postgres, Redis, Kafka, Gates):
+class Activities(Postgres, Redis, Kafka, Gates, MongoDB):
     """Activities class that combines multiple services with proper initialization."""
 
     def __init__(self,
                  postgres_config: dict[str, Any],
                  redis_config: dict[str, Any],
                  kafka_config: dict[str, Any],
+                 mongodb_config: dict[str, Any],
                  logger: Logger,
                  notification_handler: NotificationHandler):
 
@@ -62,6 +64,16 @@ class Activities(Postgres, Redis, Kafka, Gates):
             notification_handler=notification_handler
         )
 
+        # Initialize MongoDB
+        MongoDB.__init__(
+            self,
+            connection_string=mongodb_config['connection_string'],
+            database_name=mongodb_config['database_name'],
+            logger=logger,
+            notification_handler=notification_handler
+        )
+
     def shutdown(self):
         Postgres.close(self)
         Kafka.close(self)
+        MongoDB.close(self)

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

scouter/activities/redis.py

@@ -18,6 +18,47 @@ class Redis(RedisBase):
         RedisBase.__init__(self, host, port, username,
                            password, logger, notification_handler)
 
+    @activity.defn(name="get_last_data_timestamp")
+    async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Gets the last data timestamp from redis.
+        """
+        metadata = input_data['metadata']
+        key = f"{input_data['workflow_name']}_{input_data['schedule_name']}"
+
+        data_hold = self.get(key)
+
+        self.debug(
+            f"Last collected timestamp: {data_hold}",
+            metadata=metadata
+        )
+
+        if not data_hold:
+            return None
+
+        return data_hold
+
+    @activity.defn(name="put_last_data_timestamp")
+    async def put_last_data_timestamp(self, input_data: dict[str, Any]):
+        """
+        Puts the last data timestamp into redis.
+        """
+        metadata = input_data['metadata']
+        key = f"{input_data['workflow_name']}_{input_data['schedule_name']}"
+
+        data = DataFrame(input_data['data'])
+
+        last_data_timestamp = data['inserted_at'].max()
+
+        self.debug(
+            f"Last collected timestamp to insert: {last_data_timestamp}",
+            metadata=metadata
+        )
+
+        self.set(key, last_data_timestamp, ttl=input_data['retention_time'])
+
+        return last_data_timestamp
+
     @activity.defn(name="group_and_hold_data")
     async def group_and_hold_data(self, input_data: dict[str, Any]):
         """

scouter/utils/connectors_config.py

@@ -28,3 +28,15 @@ def build_redis_config():
         'username': getenv('REDIS_USERNAME', None),
         'password': getenv('REDIS_PASSWORD', None)
     }
+
+
+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')
+    }

scouter/workflow/scouter.py

@@ -41,14 +41,48 @@ class Scouter:
             }
         }
 
-        data = await workflow.execute_activity_method(
-            Activities.load_from_kafka,
+        # data = await workflow.execute_activity_method(
+        #     Activities.load_from_kafka,
+        #     {
+        #         **metadata,
+        #         'topic': input_data['topic']
+        #     },
+        #     retry_policy=retry_policy,
+        #     start_to_close_timeout=timedelta(seconds=60)
+        # )
+
+        last_data_timestamp = await workflow.execute_local_activity_method(
+            Activities.get_last_data_timestamp,
             {
                 **metadata,
-                'topic': input_data['topic']
+                'workflow_name': input_data['workflow_name'],
+                'schedule_name': input_data['schedule_name']
             },
-            retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
+            retry_policy=retry_policy
+        )
+
+        data = await workflow.execute_local_activity_method(
+            Activities.load_latest_data,
+            {
+                **metadata,
+                'collection_name': f"raw_{input_data['schedule_name']}",
+                'last_data_timestamp': last_data_timestamp
+            },
+            start_to_close_timeout=timedelta(seconds=60),
+            retry_policy=retry_policy
+        )
+
+        await workflow.execute_activity_method(
+            Activities.put_last_data_timestamp,
+            {
+                **metadata,
+                'data': data,
+                'workflow_name': input_data['workflow_name'],
+                'schedule_name': input_data['schedule_name']
+            },
+            start_to_close_timeout=timedelta(seconds=60),
+            retry_policy=retry_policy
         )
 
         if data == {}: