SIENTIAPDE-1107

feat: add MongoDB integration and update configuration for activities
This commit is contained in:
vitor-aignosi
2025-06-16 12:05:56 -03:00
parent 2db8435a42
commit a4e4d3b085
8 changed files with 308 additions and 14 deletions

View File

@@ -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:
@@ -131,3 +144,5 @@ volumes:
driver: local
kafka_data:
driver: local
mongodb_data:
driver: local

View File

@@ -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`;"
},
{
"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": {}
}
}
]

View File

@@ -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)

View File

@@ -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

View File

@@ -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'),

View File

@@ -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,

View File

@@ -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']
},

View File

@@ -3,4 +3,5 @@ psycopg2-binary
sqlalchemy
redis
couchbase
pymongo
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14