SIENTIAPDE-1325
Refactor MongoDB, Redis, and Gates classes to extend SientiaMonitoring instead of BaseActivity. Update requirements to point to local dataops library path. Implement repository pattern for MongoDB and Redis operations, enhancing code organization and maintainability.
This commit is contained in:
@@ -5,6 +5,6 @@ asyncua
|
|||||||
redis
|
redis
|
||||||
aiokafka
|
aiokafka
|
||||||
pymongo
|
pymongo
|
||||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.7
|
/home/grezewave/Documents/projects/sientia/sientia-dataops-library
|
||||||
pydruid[pandas]
|
pydruid[pandas]
|
||||||
prometheus-client
|
prometheus-client
|
||||||
@@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through():
|
|||||||
from sientia_do.notifications.handlers import NotificationHandler
|
from sientia_do.notifications.handlers import NotificationHandler
|
||||||
from sientia_do.notifications.models import NotificationLevel
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
from sientia_do.observability.logger import Logger
|
from sientia_do.observability.logger import Logger
|
||||||
from sientia_do.temporal.activities.base import BaseActivity
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||||
|
|
||||||
from scouter import metrics
|
from scouter import metrics
|
||||||
from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter
|
from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter
|
||||||
@@ -21,7 +21,7 @@ quality_gate_filters = {
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
class Gates(BaseActivity):
|
class Gates(SientiaMonitoring):
|
||||||
"""
|
"""
|
||||||
Data quality gates and filtering operations.
|
Data quality gates and filtering operations.
|
||||||
|
|
||||||
@@ -44,7 +44,13 @@ class Gates(BaseActivity):
|
|||||||
logger (Logger): Logger instance for operation logging
|
logger (Logger): Logger instance for operation logging
|
||||||
notification_handler (NotificationHandler): Handler for system notifications
|
notification_handler (NotificationHandler): Handler for system notifications
|
||||||
"""
|
"""
|
||||||
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
|
SientiaMonitoring.__init__(self, logger, notification_handler)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
"""
|
||||||
|
Close the Gates class.
|
||||||
|
"""
|
||||||
|
SientiaMonitoring.close(self)
|
||||||
|
|
||||||
def apply_aggregation(
|
def apply_aggregation(
|
||||||
self, values: DataFrame, aggr_function: str, metadata: dict[str, Any]
|
self, values: DataFrame, aggr_function: str, metadata: dict[str, Any]
|
||||||
|
|||||||
@@ -13,45 +13,13 @@ with workflow.unsafe.imports_passed_through():
|
|||||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||||
from sientia_do.notifications.models import NotificationLevel
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
from sientia_do.observability.logger import Logger
|
from sientia_do.observability.logger import Logger
|
||||||
from sientia_do.temporal.activities.base import BaseActivity
|
from sientia_do.repository.mongodb_repository import MongoDBRepository
|
||||||
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||||
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
|
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
|
||||||
|
|
||||||
|
|
||||||
def clear_mongo_id(docs: list) -> list:
|
|
||||||
"""
|
|
||||||
Remove MongoDB internal `_id` fields from documents.
|
|
||||||
|
|
||||||
This utility function recursively removes the MongoDB `_id` field from
|
class MongoDB(SientiaMonitoring):
|
||||||
documents and nested structures. It's used to clean data before
|
|
||||||
processing or export operations.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
docs (list): List of documents to clean
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
list: Documents with `_id` fields removed
|
|
||||||
|
|
||||||
Note:
|
|
||||||
This function modifies the input list in-place and returns the same reference
|
|
||||||
"""
|
|
||||||
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):
|
|
||||||
"""
|
"""
|
||||||
MongoDB operations for data retrieval and storage.
|
MongoDB operations for data retrieval and storage.
|
||||||
|
|
||||||
@@ -85,37 +53,22 @@ class MongoDB(BaseActivity):
|
|||||||
Raises:
|
Raises:
|
||||||
ConnectionError: If MongoDB connection fails
|
ConnectionError: If MongoDB connection fails
|
||||||
"""
|
"""
|
||||||
self.connection_string = connection_string
|
|
||||||
self.database_name = database_name
|
|
||||||
|
|
||||||
self.client: MongoClient = MongoClient(
|
self.mongodb_repository = MongoDBRepository(
|
||||||
self.connection_string, serverSelectionTimeoutMS=5000
|
connection_string=connection_string,
|
||||||
)
|
database_name=database_name,
|
||||||
self.client.server_info() # Trigger an exception if connection fails
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
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):
|
SientiaMonitoring.__init__(self, logger=logger, notification_handler=notification_handler)
|
||||||
"""
|
|
||||||
Gracefully close MongoDB client connection.
|
|
||||||
|
|
||||||
This method ensures proper cleanup of MongoDB connections to prevent
|
def close(self):
|
||||||
connection leaks and ensure graceful application termination.
|
|
||||||
"""
|
"""
|
||||||
try:
|
Close the MongoDB connection.
|
||||||
if self.client:
|
"""
|
||||||
self.logger.info('Closing MongoDB connection...')
|
self.mongodb_repository.close()
|
||||||
self.client.close()
|
SientiaMonitoring.shutdown(self)
|
||||||
self.logger.info('MongoDB connection closed successfully')
|
|
||||||
except Exception as e:
|
|
||||||
self.logger.error(f'Failed to close MongoDB connection: {e}')
|
|
||||||
|
|
||||||
def __del__(self):
|
def __del__(self):
|
||||||
"""
|
"""
|
||||||
@@ -124,7 +77,7 @@ class MongoDB(BaseActivity):
|
|||||||
This destructor ensures that MongoDB connections are properly closed
|
This destructor ensures that MongoDB connections are properly closed
|
||||||
when the object is garbage collected, preventing resource leaks.
|
when the object is garbage collected, preventing resource leaks.
|
||||||
"""
|
"""
|
||||||
self.shutdown()
|
self.close()
|
||||||
|
|
||||||
@activity.defn(name='load_latest_data')
|
@activity.defn(name='load_latest_data')
|
||||||
async def load_latest_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
|
async def load_latest_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
|
||||||
@@ -166,9 +119,11 @@ class MongoDB(BaseActivity):
|
|||||||
|
|
||||||
self.debug(f'Data filter: {data_filter}', metadata=metadata)
|
self.debug(f'Data filter: {data_filter}', metadata=metadata)
|
||||||
|
|
||||||
data = list(self.database[collection_name].find(data_filter, {'_id': 0}))
|
data = self.mongodb_repository.find(
|
||||||
|
collection_name=collection_name,
|
||||||
data = clear_mongo_id(data)
|
filters=data_filter,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
self.debug(f'Collected: {data}', metadata=metadata)
|
self.debug(f'Collected: {data}', metadata=metadata)
|
||||||
|
|
||||||
|
|||||||
@@ -10,11 +10,12 @@ with workflow.unsafe.imports_passed_through():
|
|||||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||||
from sientia_do.notifications.models import NotificationLevel
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
from sientia_do.observability.logger import Logger
|
from sientia_do.observability.logger import Logger
|
||||||
from sientia_do.temporal.activities.redis_base import Redis as RedisBase
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||||
|
from sientia_do.repository.redis_repository import RedisRepository
|
||||||
from sientia_do.temporal.constants import DATETIME_FORMAT, now
|
from sientia_do.temporal.constants import DATETIME_FORMAT, now
|
||||||
|
|
||||||
|
|
||||||
class Redis(RedisBase):
|
class Redis(SientiaMonitoring):
|
||||||
"""
|
"""
|
||||||
Redis operations for data caching and temporary storage.
|
Redis operations for data caching and temporary storage.
|
||||||
|
|
||||||
@@ -49,7 +50,22 @@ class Redis(RedisBase):
|
|||||||
logger (Logger): Logger instance for operation logging
|
logger (Logger): Logger instance for operation logging
|
||||||
notification_handler (NotificationHandler): Handler for system notifications
|
notification_handler (NotificationHandler): Handler for system notifications
|
||||||
"""
|
"""
|
||||||
RedisBase.__init__(self, host, port, username, password, logger, notification_handler)
|
SientiaMonitoring.__init__(self, logger, notification_handler)
|
||||||
|
self.redis_repository = RedisRepository(
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
username=username,
|
||||||
|
password=password,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
"""
|
||||||
|
Close the Redis connection.
|
||||||
|
"""
|
||||||
|
self.redis_repository.close()
|
||||||
|
SientiaMonitoring.shutdown(self)
|
||||||
|
|
||||||
@activity.defn(name='get_last_data_timestamp')
|
@activity.defn(name='get_last_data_timestamp')
|
||||||
async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
|
async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
|
||||||
@@ -79,7 +95,7 @@ class Redis(RedisBase):
|
|||||||
self.info(f'Getting last data timestamp for {key}')
|
self.info(f'Getting last data timestamp for {key}')
|
||||||
|
|
||||||
try:
|
try:
|
||||||
data_hold = self.get(key)
|
data_hold = self.redis_repository.get(key)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.send_notification(
|
self.send_notification(
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
@@ -137,7 +153,7 @@ class Redis(RedisBase):
|
|||||||
self.info(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
|
self.info(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
self.set(key, last_data_timestamp, ttl=60 * 60 * 5)
|
self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.send_notification(
|
self.send_notification(
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
@@ -190,7 +206,7 @@ class Redis(RedisBase):
|
|||||||
self.info(f'Getting held data for {key}')
|
self.info(f'Getting held data for {key}')
|
||||||
|
|
||||||
try:
|
try:
|
||||||
data_hold = self.get(key)
|
data_hold = self.redis_repository.get(key)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.send_notification(
|
self.send_notification(
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
@@ -234,7 +250,7 @@ class Redis(RedisBase):
|
|||||||
data['timestamp'].max() if not data.empty else data_hold['timestamp']
|
data['timestamp'].max() if not data.empty else data_hold['timestamp']
|
||||||
)
|
)
|
||||||
|
|
||||||
self.set(key, data_hold, ttl=retention_time)
|
self.redis_repository.set(key, data_hold, ttl=retention_time)
|
||||||
|
|
||||||
data_hold_df = DataFrame(data_hold, index=[0])
|
data_hold_df = DataFrame(data_hold, index=[0])
|
||||||
data_hold_melted = data_hold_df.melt(
|
data_hold_melted = data_hold_df.melt(
|
||||||
@@ -280,7 +296,7 @@ class Redis(RedisBase):
|
|||||||
cache = {'data': data.to_dict(), 'held_data': held_data.to_dict()}
|
cache = {'data': data.to_dict(), 'held_data': held_data.to_dict()}
|
||||||
|
|
||||||
try:
|
try:
|
||||||
self.set(key, cache, ttl=120)
|
self.redis_repository.set(key, cache, ttl=120)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.send_notification(
|
self.send_notification(
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
|
|||||||
Reference in New Issue
Block a user