sientia-dataops-orchestrator_temporal/laborious/activities/postgres.py

import traceback
from temporalio import workflow, activity

from laborious.activities.base import BaseActivity
with workflow.unsafe.imports_passed_through():
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.pool import QueuePool
    from psycopg2.pool import ThreadedConnectionPool
    from pandas import read_sql_query, DataFrame
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from typing import Any


class Postgres(BaseActivity):
    def __init__(self, host: str, port: int,
                 user: str, password: str, dbname: str,
                 min_connections: int, max_connections: int,
                 logger: Logger, notification_handler: NotificationHandler):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.dbname = dbname

        # Create SQLAlchemy engine with connection pooling
        self.engine = create_engine(
            f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
            poolclass=QueuePool,
            pool_size=min_connections,
            max_overflow=max_connections - min_connections,
            pool_pre_ping=True
        )
        self.session_factory = sessionmaker(bind=self.engine)

        BaseActivity.__init__(self, logger, notification_handler)

    def close(self):
        self.engine.dispose()

    def __del__(self):
        self.close()

    @activity.defn(name="load_custom_query")
    async def load_custom_query(self, query: str) -> dict[str, Any]:
        """
        Loads data from a custom query.

        Args:
            query (str): The query to load data from.

        Returns:
            dict[str, dict]: The data from the query.
        """
        self.logger.info(f"Fetching data from query: {query}")

        data = None
        with self.session_factory() as session:
            try:
                data = read_sql_query(query, self.engine)

            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    notification_id="ERROR_LOADING_CUSTOM_QUERY",
                    message=f"Error fetching data from query: {e}",
                    block="load_custom_query",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

                self.logger.error(trace)

                return {}
            finally:
                session.close()

        if data is None:
            return {}

        # Converts any datetime datatype columns to string
        for col in data.select_dtypes(include=['datetime64']).columns:
            data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S')

        self.logger.info(f"Fetched {len(data)} rows")
        self.logger.debug(f"Data: \n{data.to_string()}")

        return data.to_dict()

    @activity.defn(name="repeat_last_prediction")
    async def repeat_last_prediction(self, query_items: dict[str, str]):
        """
        Repeats the last prediction for a given model.

        Args:
            query_items (dict[str, str]): The query items. Contains:
                schema (str): The schema of the table.
                table_name (str): The name of the table.
                model (int): The model to repeat the prediction for.

        Returns:
            None
        """
        schema = query_items["schema"]
        table_name = query_items["table_name"]
        model = query_items["model"]

        repeat_query = f"""
        INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
        SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
        FROM \"{schema}\".{table_name}
        WHERE model_id = {model}
        ORDER BY timestamp DESC
        LIMIT 1;
        """
        self.logger.info(f"Repeating last prediction for model {model}")
        self.logger.debug(f"Query: {repeat_query}")

        with self.session_factory() as session:
            try:
                session.execute(repeat_query)
                session.commit()

            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    notification_id="ERROR_REPEATING_LAST_PREDICTION",
                    message=f"Error repeating last prediction: {e}",
                    block="repeat_last_prediction",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

                self.logger.error(trace)

            finally:
                session.close()

    @activity.defn(name="export_data_to_postgres")
    async def export_data_to_postgres(self, input_data: dict[str, Any]):
        """
        Exports data to a postgres table.

        Args:
            input_data (dict[str, Any]): The data to export. Contains:
                schema (str): The schema of the table.
                table_name (str): The name of the table.
                data (DataFrame): The data to export.
        """

        self.logger.debug(
            f"Exporting data to postgres: {input_data['data']}")

        schema = input_data["schema"]
        table_name = input_data["table_name"]
        data = DataFrame(input_data["data"])

        with self.session_factory() as session:
            try:
                data.to_sql(table_name, self.engine, schema=schema,
                            if_exists="append", index=False)
                session.commit()

            except Exception as e:
                trace = traceback.format_exc()
                self.notification_handler.build_and_send_notification(
                    notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
                    message=f"Error exporting data to postgres: {e}",
                    block="export_data_to_postgres",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

                self.logger.error(trace)

            else:
                self.logger.debug("Data exported to postgres")
            finally:
                session.close()