c9b9974f9d
Refactor Redis activity to remove unused 'ttl' parameter in set method, streamlining data handling.
119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
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.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]) -> dict[str, Any]:
|
|
"""
|
|
Gets the last data timestamp from redis.
|
|
"""
|
|
metadata = input_data['metadata']
|
|
key = f"{input_data['workflow_name']}_{input_data['schedule_name']}"
|
|
|
|
data_hold = self.get(key)
|
|
|
|
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"{input_data['workflow_name']}_{input_data['schedule_name']}"
|
|
|
|
data = DataFrame(input_data['data'])
|
|
|
|
last_data_timestamp = data['inserted_at'].max()
|
|
|
|
self.debug(
|
|
f"Last collected timestamp to insert: {last_data_timestamp}",
|
|
metadata=metadata
|
|
)
|
|
|
|
self.set(key, last_data_timestamp)
|
|
|
|
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"{input_data['workflow_name']}_{input_data['schedule_name']}"
|
|
|
|
data_hold = self.get(key)
|
|
|
|
if not data_hold:
|
|
data_hold = {}
|
|
if data.empty:
|
|
self.warning("No data to export",
|
|
metadata=metadata
|
|
)
|
|
return data_hold
|
|
|
|
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)
|
|
|
|
self.debug(
|
|
f"Data grouped and held successfully:\n {data_hold_melted.to_string()}",
|
|
metadata=metadata
|
|
)
|
|
|
|
return data_hold_melted.to_dict()
|