sientia-dataops-scouter_temporal/scouter/activities/redis.py

import traceback
from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.redis_base import Redis as RedisBase
    from sientia_do.temporal.utils.logger import Logger
    from typing import Any
    from pandas import DataFrame
    from datetime import datetime


class Redis(RedisBase):
    def __init__(self, host: str, port: int,
                 username: str, password: str,
                 logger: Logger, notification_handler: NotificationHandler):

        RedisBase.__init__(self, host, port, username,
                           password, logger, notification_handler)

    @activity.defn(name="get_last_data_timestamp")
    async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
        """
        Gets the last data timestamp from redis.
        """
        metadata = input_data['metadata']
        key = f"last_data_timestamp_{input_data['workflow_name']}_{input_data['schedule_name']}"

        try:
            data_hold = self.get(key)
        except Exception as e:
            self.send_notification(
                metadata=metadata,
                notification_id="REDIS_GET_ERROR",
                message=f"Error getting last data timestamp: {e}",
                block="get_last_data_timestamp",
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc()
            )
            raise e

        self.debug(
            f"Last collected timestamp: {data_hold}",
            metadata=metadata
        )

        if not data_hold:
            return None

        return data_hold

    @activity.defn(name="put_last_data_timestamp")
    async def put_last_data_timestamp(self, input_data: dict[str, Any]):
        """
        Puts the last data timestamp into redis.
        """
        metadata = input_data['metadata']
        key = f"last_data_timestamp_{input_data['workflow_name']}_{input_data['schedule_name']}"

        data = DataFrame(input_data['data'])

        if data.empty:
            self.warning("No data to insert",
                         metadata=metadata
                         )
            return None

        last_data_timestamp = data['inserted_at'].max()

        self.debug(
            f"Last collected timestamp to insert: {last_data_timestamp}",
            metadata=metadata
        )

        try:
            self.set(key, last_data_timestamp, ttl=None)
        except Exception as e:
            self.send_notification(
                metadata=metadata,
                notification_id="REDIS_SET_ERROR",
                message=f"Error setting last data timestamp: {e}",

                block="put_last_data_timestamp",
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc()
            )
            raise e

        return last_data_timestamp

    @activity.defn(name="group_and_hold_data")
    async def group_and_hold_data(self, input_data: dict[str, Any]):
        """
        Groups and holds data in redis. Keep a copy of the most recent 
        received data for a given pipeline and schedule. This activity updates
        the data in redis and return the full keeped data.

        Args:
            input_data (dict[str, Any]): The data to group and hold.
                workflow_name (str): The name of the workflow.
                schedule_name (str): The name of the schedule.
                data (dict[str, Any]): The data to group and hold.
                retention_time (int): The retention time for data in redis in seconds.
        """
        metadata = input_data['metadata']

        self.debug("Grouping and holding data...",
                   metadata=metadata
                   )
        data = DataFrame(input_data['data'])
        retention_time = input_data['retention_time']

        key = f"held_data_{input_data['workflow_name']}_{input_data['schedule_name']}"

        try:
            data_hold = self.get(key)
        except Exception as e:
            self.send_notification(
                metadata=metadata,
                notification_id="REDIS_GET_ERROR",
                message=f"Error getting held data: {e}",
                block="group_and_hold_data",
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc()
            )
            raise e

        if not data_hold:
            data_hold = {}
            if data.empty:
                self.warning("No data to export",
                             metadata=metadata
                             )
                return data_hold

        try:
            for _, row in data.iterrows():
                value = row['value']

                data_hold[row['name']] = value

            data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            self.set(key, data_hold, ttl=retention_time)

            data_hold_df = DataFrame(data_hold, index=[0])
            data_hold_melted = data_hold_df.melt(
                id_vars='timestamp', var_name='variable', value_name='value')
            data_hold_melted['model_id'] = input_data['model_id']

            data_hold_melted.reset_index(drop=True, inplace=True)
        except Exception as e:
            self.send_notification(
                metadata=metadata,
                notification_id="REDIS_SET_ERROR",
                message=f"Error setting held data: {e}",
                block="group_and_hold_data",
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc()
            )
            raise e

        self.debug(
            f"Data grouped and held successfully:\n {data_hold_melted.to_string()}",
            metadata=metadata
        )

        return data_hold_melted.to_dict()