diff --git a/requirements.txt b/requirements.txt index c29f059..3e098f3 100644 --- a/requirements.txt +++ b/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 diff --git a/scouter/activities/activities.py b/scouter/activities/activities.py index 8fc7efb..ddeae2e 100644 --- a/scouter/activities/activities.py +++ b/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) diff --git a/scouter/activities/mongodb.py b/scouter/activities/mongodb.py new file mode 100644 index 0000000..81370db --- /dev/null +++ b/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 diff --git a/scouter/activities/redis.py b/scouter/activities/redis.py index 09c3f40..1083d27 100644 --- a/scouter/activities/redis.py +++ b/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]): """ diff --git a/scouter/utils/connectors_config.py b/scouter/utils/connectors_config.py index 6acc4ae..64a7feb 100644 --- a/scouter/utils/connectors_config.py +++ b/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') + } diff --git a/scouter/workflow/scouter.py b/scouter/workflow/scouter.py index ecf0716..cf36c49 100644 --- a/scouter/workflow/scouter.py +++ b/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 == {}: