f441873969
Refactor MongoDB data retrieval to convert cursor to list, improving data handling and compatibility.
134 lines
4.1 KiB
Python
134 lines
4.1 KiB
Python
from temporalio import workflow, activity
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from typing import Any
|
|
import traceback
|
|
from logging import Logger
|
|
from datetime 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']
|
|
|
|
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")
|
|
}
|
|
}
|
|
|
|
data = list(self.database[collection_name].find(
|
|
data_filter, {"_id": 0}))
|
|
|
|
data = clear_mongo_id(data)
|
|
|
|
for item in data:
|
|
item['inserted_at'] = item['inserted_at'].isoformat()
|
|
|
|
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
|