fa88572126
SIENTIAPDE-1712 Enhance logging in Gates, MLFlow, and Storage classes by integrating logger parameter for improved traceability. This update allows for better monitoring of operations and data handling across these components.
210 lines
7.5 KiB
Python
210 lines
7.5 KiB
Python
from temporalio import activity, workflow
|
|
|
|
from laborious.utils.repository.minio_manager import MinioManager
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
import traceback
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
import pandas as pd
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.temporal.activities.postgres import Postgres
|
|
from sientia_do.temporal.constants import now
|
|
|
|
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
|
|
|
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
|
|
|
|
|
|
class Storage(Postgres, MinioManager):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name.
|
|
"""
|
|
|
|
minio_repository: MinioRepository | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
min_connections: int,
|
|
max_connections: int,
|
|
retention_hours: int = 24,
|
|
minio_repository: MinioRepository | None = None,
|
|
logger: Logger | None = None,
|
|
notification_handler: NotificationHandler | None = None,
|
|
metrics_controller: MetricsController | None = None,
|
|
):
|
|
self.retention_hours = retention_hours
|
|
Postgres.__init__(
|
|
self,
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
dbname=dbname,
|
|
min_connections=min_connections,
|
|
max_connections=max_connections,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
MinioManager.__init__(
|
|
self, minio_repository, logger, notification_handler, metrics_controller
|
|
)
|
|
|
|
@activity.defn(name='load_query_with_minio_offload')
|
|
async def load_query_with_minio_offload(
|
|
self, input_data: dict[str, Any]
|
|
) -> MinioDataFramePayload:
|
|
"""
|
|
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata (same as load_custom_query).
|
|
query (str): SQL query.
|
|
datetime_columns (list[str], optional): Datetime column names.
|
|
model_name (str): Model name for object key basename.
|
|
key_prefix (str, optional): Directory prefix inside the bucket.
|
|
size_threshold_bytes (int, optional): Override env offload threshold.
|
|
|
|
Returns:
|
|
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
|
|
"""
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata: dict = input_data.get('metadata', {})
|
|
model_name = input_data['model_name']
|
|
|
|
rows = await self.load_custom_query(
|
|
input_data,
|
|
)
|
|
if not rows:
|
|
self.error(
|
|
'load_query_with_minio_offload failed: No data returned from query', metadata
|
|
)
|
|
dataframe = None
|
|
else:
|
|
dataframe = pd.DataFrame(rows)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe,
|
|
minio_repo=self.minio_repository,
|
|
workflow_metadata=metadata,
|
|
model_name=model_name,
|
|
operation='initial',
|
|
logger=self.logger,
|
|
)
|
|
|
|
@activity.defn(name='export_payload_to_postgres')
|
|
async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
|
"""
|
|
Export a payload to PostgreSQL.
|
|
"""
|
|
metadata = input_data.get('metadata')
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
return await self.export_data_to_postgres(
|
|
{
|
|
**input_data,
|
|
'data': data,
|
|
}
|
|
)
|
|
|
|
@activity.defn(name='cleanup_minio_objects_expired')
|
|
async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Delete objects under the given prefixes that are older than the retention window.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata for logging and metrics.
|
|
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
|
|
|
|
Returns:
|
|
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
|
|
"""
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
prefix = payload.cleanup_prefix()
|
|
base = now()
|
|
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
|
|
hours=self.retention_hours
|
|
)
|
|
|
|
report: dict[str, Any] = {
|
|
'failed': {},
|
|
'deleted': {},
|
|
'failed_count': 0,
|
|
'deleted_count': 0,
|
|
}
|
|
try:
|
|
keys = await self.minio_repository.list_objects(
|
|
prefix=prefix,
|
|
recursive=True,
|
|
metadata=metadata,
|
|
)
|
|
for key in keys:
|
|
try:
|
|
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
|
if ts is None:
|
|
continue
|
|
if ts >= cutoff:
|
|
continue
|
|
await self.minio_repository.delete_file(
|
|
object_name=key,
|
|
metadata=metadata,
|
|
)
|
|
except Exception as e:
|
|
report['failed'][key] = {
|
|
'success': False,
|
|
'message': str(e),
|
|
}
|
|
report['failed_count'] += 1
|
|
continue
|
|
report['deleted'][key] = {
|
|
'success': True,
|
|
'message': 'Deleted',
|
|
}
|
|
report['deleted_count'] += 1
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
|
|
message=f'Error cleaning up MinIO objects: {e}',
|
|
block='cleanup_minio_objects_expired',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata)
|
|
else:
|
|
# Cleanup success is expected in normal flow; avoid noisy INFO notifications
|
|
# that do not impact behavior and can flood observability in test runs.
|
|
self.info('MinIO objects cleaned up successfully', metadata)
|
|
|
|
return report
|
|
|
|
def close(self) -> None:
|
|
"""Close Storage resources (MinIO client and Postgres engine)."""
|
|
Postgres.close(self)
|
|
MinioManager.close(self)
|
|
|
|
def __del__(self):
|
|
self.close()
|