0e915f5491
Update Redis activity to include specific prefixes for keys in last data timestamp and held data, enhancing clarity in data management.
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"last_data_timestamp_{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"held_data_{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()
|