sientia-dataops-laborious_temporal/laborious/activities/storage.py

import json

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 io import BytesIO
    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 DATETIME_FORMAT_FILENAME, 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',
        )

    @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 = 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', {})
        prefix = input_data['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:
            await self.send_notification_async(
                metadata=metadata,
                notification_id='CLEANUP_MINIO_OBJECTS_EXPIRED',
                message='MinIO objects cleaned up successfully',
                block='cleanup_minio_objects_expired',
                level=NotificationLevel.INFO,
                attachment_content=json.dumps(report),
            )

        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()