461a885eeb
Refactor OPC handling by removing pod_id from initialization and updating logging format - Removed pod_id parameter from OPC class and repository initialization to streamline connection management. - Updated logging statements for improved readability during disconnection attempts and error handling.
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
import traceback
|
|
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.temporal.activities.postgres import Postgres
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now
|
|
|
|
from laborious.utils.repository.minio_repository import MinioRepository
|
|
|
|
|
|
class Storage(Postgres):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
min_connections: int,
|
|
max_connections: int,
|
|
minio_config: dict[str, Any],
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler,
|
|
):
|
|
super().__init__(
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
dbname=dbname,
|
|
min_connections=min_connections,
|
|
max_connections=max_connections,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
if not hasattr(self, 'minio_repository'):
|
|
self.minio_repository: MinioRepository | None = None
|
|
|
|
if self.minio_repository is None:
|
|
self.minio_repository = MinioRepository(
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
minio_endpoint_url=minio_config['endpoint_url'],
|
|
minio_access_key=minio_config['access_key'],
|
|
minio_secret_key=minio_config['secret_key'],
|
|
minio_region_name=minio_config['region_name'],
|
|
minio_default_bucket=minio_config['default_bucket'],
|
|
)
|
|
|
|
@activity.defn(name='query_to_minio')
|
|
async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Execute SQL query, write result as Parquet to MinIO, and return object name.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata
|
|
query (str): SQL query
|
|
model_name (str): Model name for object naming
|
|
object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain)
|
|
|
|
Returns:
|
|
dict: { success: bool, object_name: str, uri: str }
|
|
"""
|
|
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
object_prefix = input_data.get('object_prefix', 'datasets/retrain')
|
|
|
|
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
|
|
object_name = f'{object_prefix}_{timestamp}.parquet'
|
|
uri = f's3://{self.minio_repository.minio_bucket}/{object_name}'
|
|
|
|
try:
|
|
data = await self.load_custom_query(input_data)
|
|
if not data:
|
|
self.error('query_to_minio failed: No data returned from query', metadata)
|
|
return {'success': False, 'message': 'No data returned from query'}
|
|
|
|
# Ensure we have a DataFrame
|
|
data = pd.DataFrame(data)
|
|
|
|
# Write parquet to memory and upload via persistent client
|
|
self.minio_repository.store_dataframe_as_parquet(
|
|
dataframe=data, uri=uri, object_name=object_name, metadata=metadata
|
|
)
|
|
|
|
return {'success': True, 'object_key': object_name, 'uri': uri}
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='ERROR_STORING_QUERY_TO_MINIO',
|
|
message=f'Error storing query to MinIO: {e}',
|
|
block='query_to_minio',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
self.error(trace, metadata)
|
|
|
|
return {'success': False, 'message': str(e)}
|
|
|
|
def close(self) -> None:
|
|
"""Close Storage resources (MinIO client and Postgres engine)."""
|
|
try:
|
|
if hasattr(self, 'minio_repository') and self.minio_repository is not None:
|
|
try:
|
|
self.minio_repository.close()
|
|
finally:
|
|
self.minio_repository = None
|
|
finally:
|
|
# Ensure Postgres resources are disposed as well
|
|
try:
|
|
super().close()
|
|
except Exception:
|
|
self.logger.error('Error closing Postgres resources')
|
|
|
|
def __del__(self):
|
|
self.close()
|