2efad30304
SIENTIAPDE-1174 Update dependencies, modify image tag, and enhance metrics tracking in various activities for improved observability and performance.
147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
from temporalio import workflow, activity
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from typing import Any
|
|
import traceback
|
|
from datetime import datetime
|
|
from pymongo import MongoClient
|
|
from pandas import DataFrame
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.activities.base import BaseActivity
|
|
from sientia_do.temporal.utils.logger import Logger
|
|
|
|
|
|
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,
|
|
set_error_counter=True)
|
|
|
|
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']
|
|
|
|
self.debug(
|
|
f"Loading data from MongoDB: {input_data}",
|
|
metadata=metadata
|
|
)
|
|
|
|
try:
|
|
|
|
if last_data_timestamp is None:
|
|
data_filter = {}
|
|
else:
|
|
data_filter = {
|
|
"inserted_at": {
|
|
"$gt": datetime.strptime(last_data_timestamp, "%Y-%m-%d %H:%M:%S.%f")
|
|
}
|
|
}
|
|
|
|
self.debug(
|
|
f"Data filter: {data_filter}",
|
|
metadata=metadata
|
|
)
|
|
|
|
data = list(self.database[collection_name].find(
|
|
data_filter, {"_id": 0}))
|
|
|
|
data = clear_mongo_id(data)
|
|
|
|
self.debug(
|
|
f"Collected: {data}",
|
|
metadata=metadata
|
|
)
|
|
|
|
for item in data:
|
|
item['inserted_at'] = item['inserted_at'].strftime(
|
|
"%Y-%m-%d %H:%M:%S.%f")
|
|
|
|
self.info(
|
|
f"Loaded {len(data)} documents from MongoDB",
|
|
metadata=metadata
|
|
)
|
|
|
|
self.debug(
|
|
f"Loaded data: {data}",
|
|
metadata=metadata
|
|
)
|
|
|
|
return DataFrame(data).to_dict()
|
|
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
|