SIENTIAPDE-1110

Enhance Gates and Redis activities by adding metadata parameter to apply_aggregation and notification methods. Refactor notification handling to use send_notification for improved consistency. Update tests to reflect changes in notification method calls and ensure proper functionality with new metadata integration.
2025-07-04 11:01:20 -03:00
parent d08d1b1337
commit 08240efc8b
6 changed files with 425 additions and 34 deletions
--- a/scouter/activities/gates.py
+++ b/scouter/activities/gates.py
@@ -16,7 +16,8 @@ quality_gate_filters = {

 class Gates(BaseActivity):

-    def apply_aggregation(self, group: DataFrame, aggr_function: str) -> float | None | str:
+    def apply_aggregation(self, group: DataFrame, aggr_function: str,
+                          metadata: dict[str, Any]) -> float | None | str:
        """
        Apply aggregation function to a group of data.

@@ -48,7 +49,8 @@ class Gates(BaseActivity):
            elif aggr_function == 'min':
                return group['value'].min()
            else:
-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                    notification_id="AGGREGATION_ISSUES",
                    message=f"Invalid aggregation function: {aggr_function}",
                    block="aggregate_data",
@@ -100,7 +102,8 @@ class Gates(BaseActivity):
                # Get the latest timestamp
                latest_timestamp = group['timestamp'].max()

-                aggr_value = self.apply_aggregation(group, aggr_function)
+                aggr_value = self.apply_aggregation(
+                    group, aggr_function, metadata)

                if aggr_value == 'continue':
                    continue
@@ -145,7 +148,8 @@ class Gates(BaseActivity):
        except Exception as e:
            trace = traceback.format_exc()

-            self.notification_handler.build_and_send_notification(
+            self.send_notification(
+                metadata=metadata,
                notification_id="AGGREGATION_ISSUES",
                message=f"Error aggregating data: {e}",
                block="aggregate_data",
@@ -202,7 +206,8 @@ class Gates(BaseActivity):

            except Exception as e:
                trace = traceback.format_exc()
-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                    notification_id="DATA_QUALITY_GATE_ISSUES",
                    message=f"Error applying filter {filter_name}: {e}",
                    block="data_quality_gate",
@@ -219,7 +224,8 @@ class Gates(BaseActivity):
                message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}"
                attachment = filtered_data.to_string()

-                self.notification_handler.build_and_send_notification(
+                self.send_notification(
+                    metadata=metadata,
                    notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}",
                    message=message,
                    block="data_quality_gate",
--- a/scouter/activities/redis.py
+++ b/scouter/activities/redis.py
@@ -1,8 +1,10 @@
+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
@@ -26,7 +28,18 @@ class Redis(RedisBase):
        metadata = input_data['metadata']
        key = f"last_data_timestamp_{input_data['workflow_name']}_{input_data['schedule_name']}"

-        data_hold = self.get(key)
+        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}",
@@ -48,6 +61,12 @@ class Redis(RedisBase):

        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(
@@ -55,7 +74,19 @@ class Redis(RedisBase):
            metadata=metadata
        )

-        self.set(key, last_data_timestamp, ttl=None)
+        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

@@ -83,7 +114,18 @@ class Redis(RedisBase):

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

-        data_hold = self.get(key)
+        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 = {}
@@ -93,22 +135,33 @@ class Redis(RedisBase):
                             )
                return data_hold

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

-            data_hold[row['name']] = 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")
+            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)
+            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_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)
+            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()}",