From a4e4d3b0851ce741da387f638b6a2aef17a71777 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 16 Jun 2025 12:05:56 -0300 Subject: [PATCH] SIENTIAPDE-1107 feat: add MongoDB integration and update configuration for activities --- docker-compose.yml | 15 ++ input_sample.json | 34 +++- orchestrator/activities/activities.py | 26 +++- orchestrator/activities/mongo_db.py | 198 ++++++++++++++++++++++++ orchestrator/utils/connectors_config.py | 15 ++ orchestrator/worker/worker.py | 11 +- orchestrator/workflows/orchestrator.py | 22 ++- requirements.txt | 1 + 8 files changed, 308 insertions(+), 14 deletions(-) create mode 100644 orchestrator/activities/mongo_db.py diff --git a/docker-compose.yml b/docker-compose.yml index baab896..e3aba82 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -117,6 +117,19 @@ services: networks: - sientia-network + mongodb: + image: mongo:7.0 + container_name: mongodb + ports: + - "27017:27017" + environment: + MONGO_INITDB_ROOT_USERNAME: sientia + MONGO_INITDB_ROOT_PASSWORD: sientia + volumes: + - mongodb_data:/data/db + networks: + - sientia-network + networks: sientia-network: @@ -130,4 +143,6 @@ volumes: redis_data: driver: local kafka_data: + driver: local + mongodb_data: driver: local \ No newline at end of file diff --git a/input_sample.json b/input_sample.json index 4a89997..e8d3641 100644 --- a/input_sample.json +++ b/input_sample.json @@ -1,5 +1,37 @@ +[ { "schedule_name": "orchestrator-test", "pipelines_query": "SELECT pipelines.*, models FROM `pipelines` JOIN `models` ON KEYS pipelines.model_id;", "opc_servers_query": "select META().id, opc_servers.* from `opc_servers`;" -} \ No newline at end of file +}, +{ + "schedule_name": "orchestrator-test", + "pipelines_query": { + "collection": "pipelines", + "aggregation": [ + { + "$lookup": { + "from": "models", + "localField": "model_id", + "foreignField": "id", + "as": "model_docs" + } + }, + { + "$addFields": { + "models": { "$arrayElemAt": ["$model_docs", 0] } + } + }, + { + "$project": { + "model_docs": 0 + } + } + ] + }, + "opc_servers_query": { + "collection": "opc-servers", + "filters": {} + } +} +] \ No newline at end of file diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py index 1265539..5049e3f 100644 --- a/orchestrator/activities/activities.py +++ b/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) diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py new file mode 100644 index 0000000..e969e43 --- /dev/null +++ b/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 diff --git a/orchestrator/utils/connectors_config.py b/orchestrator/utils/connectors_config.py index 2e5e80f..008dc97 100644 --- a/orchestrator/utils/connectors_config.py +++ b/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'), diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index 0814b25..383c1be 100644 --- a/orchestrator/worker/worker.py +++ b/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, diff --git a/orchestrator/workflows/orchestrator.py b/orchestrator/workflows/orchestrator.py index 12550d5..b731eeb 100644 --- a/orchestrator/workflows/orchestrator.py +++ b/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'] }, diff --git a/requirements.txt b/requirements.txt index 07d06ae..25595c6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ psycopg2-binary sqlalchemy redis couchbase +pymongo git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14