From 4e579dd5bd62cad75835aae298065d6003a2055c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 13 May 2025 16:58:03 -0300 Subject: [PATCH] SIENTIAPDE-1005 Add initial implementation of data processing activities and filters --- requirements.txt | 5 ++ scouter/activities/base.py | 26 +++++++++++ scouter/activities/gates.py | 75 ++++++++++++++++++++++++++++++ scouter/activities/kafka.py | 58 ++++++++++++++++++++++++ scouter/activities/postgres.py | 78 ++++++++++++++++++++++++++++++++ scouter/activities/redis.py | 18 ++++++++ scouter/utils/quality/filters.py | 21 +++++++++ tests/__init__.py | 0 tests/activities/__init__.py | 0 tests/activities/test_kafka.py | 64 ++++++++++++++++++++++++++ tests/utils/__init__.py | 0 tests/worker/__init__.py | 0 tests/workflow/__init__.py | 0 13 files changed, 345 insertions(+) create mode 100644 requirements.txt create mode 100644 scouter/activities/base.py create mode 100644 scouter/activities/gates.py create mode 100644 scouter/activities/kafka.py create mode 100644 scouter/activities/postgres.py create mode 100644 scouter/activities/redis.py create mode 100644 scouter/utils/quality/filters.py create mode 100644 tests/__init__.py create mode 100644 tests/activities/__init__.py create mode 100644 tests/activities/test_kafka.py create mode 100644 tests/utils/__init__.py create mode 100644 tests/worker/__init__.py create mode 100644 tests/workflow/__init__.py diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..1ccf40c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +temporalio +psycopg2-binary +asyncua +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git diff --git a/scouter/activities/base.py b/scouter/activities/base.py new file mode 100644 index 0000000..7269d41 --- /dev/null +++ b/scouter/activities/base.py @@ -0,0 +1,26 @@ +from typing import Any +from logging import Logger +from temporalio import activity +from sientia_do.notifications.handlers import NotificationHandler + + +class BaseActivity: + def __init__(self, logger: Logger, notification_handler: NotificationHandler): + self.logger = logger + self.notification_handler = notification_handler + + @activity.defn(name="prepare_activity") + def prepare_activity(self, input_data: dict[str, Any]): + """ + Prepare the activity for the notification handler. + + Args: + workflow_name (str): The name of the workflow. + schedule_name (str): The name of the schedule. + model_name (str): The name of the model. + model_id (str): The id of the model. + """ + self.notification_handler.base_notification.pipeline_name = input_data['workflow_name'] + self.notification_handler.base_notification.schedule_name = input_data['schedule_name'] + self.notification_handler.base_notification.model_name = input_data['model_name'] + self.notification_handler.base_notification.model_id = input_data['model_id'] diff --git a/scouter/activities/gates.py b/scouter/activities/gates.py new file mode 100644 index 0000000..d28e750 --- /dev/null +++ b/scouter/activities/gates.py @@ -0,0 +1,75 @@ +from temporalio import workflow, activity + +with workflow.unsafe.imports_passed_through(): + from sientia_do.notifications.models import NotificationLevel + from pandas import DataFrame + from scouter.activities.base import BaseActivity + from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter + from typing import Any + import traceback + +quality_gate_filters = { + 'NULL_VALUES_FILTER': null_values_filter, + 'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter +} + + +class Gates(BaseActivity): + @activity.defn(name="data_quality_gate") + async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Data quality gate activity. for each selected filter, + extracts filtered data, discards or keeps filtered data + based on the filter. + + Args: + input_data (dict[str, Any]): The data to validate. Contains: + filters (dict[str, str]): The filters to apply. In format: + {filter_name: policy}. + filter_name: The name of the filter. + policy: The policy to apply. Can be "DISCARD" or "KEEP". + data (dict[str, Any]): The data to validate. + + Returns: + dict[str, Any]: The data validated. + """ + + filters = input_data['filters'] + data = DataFrame(input_data['data']) + + for filter_name, policy in filters.items(): + if filter_name not in quality_gate_filters: + self.logger.warning(f"Filter {filter_name} not found") + continue + + try: + filtered_data = quality_gate_filters[filter_name](data) + + except Exception as e: + self.notification_handler.build_and_send_notification( + notification_id="DATA_QUALITY_GATE_ISSUES", + message=f"Error applying filter {filter_name}: {e}", + block="data_quality_gate", + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc() + ) + + else: + if filtered_data.empty: + continue + + message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}" + attachment = filtered_data.to_string() + + self.notification_handler.build_and_send_notification( + notification_id="DATA_QUALITY_GATE_ISSUES", + message=message, + block="data_quality_gate", + level=NotificationLevel.WARNING, + attachment_content=attachment + ) + + if policy == "DISCARD": + data = data[not data.isin(filtered_data).all(axis=1)] + + return data.to_dict() diff --git a/scouter/activities/kafka.py b/scouter/activities/kafka.py new file mode 100644 index 0000000..916bb4d --- /dev/null +++ b/scouter/activities/kafka.py @@ -0,0 +1,58 @@ +from temporalio import workflow, activity + +with workflow.unsafe.imports_passed_through(): + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + from scouter.activities.base import BaseActivity + from typing import Any + from kafka import KafkaConsumer + from pandas import DataFrame + + +class Kafka(BaseActivity): + def __init__(self, bootstrap_servers: str, polling_time: int, + group_id: str, logger: Logger, notification_handler: NotificationHandler): + self.polling_time = polling_time + + self.kafka_connector = KafkaConsumer( + bootstrap_servers=bootstrap_servers, + auto_offset_reset="earliest", + enable_auto_commit=True, + group_id=group_id, + value_deserializer=lambda x: x.decode("utf-8") + ) + + super().__init__(logger, notification_handler) + + @activity.defn(name="load_from_kafka") + def load_from_kafka(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Loads data from a kafka topic. Polls the topic for a given time and returns the data. + + Args: + input_data (dict[str, Any]): The data to load. Contains: + topic (str): The topic to load data from. + Returns: + dict[str, Any]: The data loaded from the topic. + """ + topic = input_data["topic"] + + # Subscribe to the specified topic + self.kafka_connector.subscribe([topic]) + + # List to store message values + message_values = [] + + # Poll for messages + records = self.kafka_connector.poll(timeout_ms=self.polling_time) + + # Process the polled records + for _topic_partition, msgs in records.items(): + for msg in msgs: + message_values.append(msg.value) + + # Return empty dict if no messages were received + if not message_values: + return {} + + return DataFrame(message_values).to_dict() diff --git a/scouter/activities/postgres.py b/scouter/activities/postgres.py new file mode 100644 index 0000000..b6b7d94 --- /dev/null +++ b/scouter/activities/postgres.py @@ -0,0 +1,78 @@ +import traceback +from temporalio import workflow, activity + +with workflow.unsafe.imports_passed_through(): + from psycopg2.pool import ThreadedConnectionPool + from pandas import DataFrame + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from scouter.activities.base import BaseActivity + 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 + + self.pool = ThreadedConnectionPool( + minconn=min_connections, + maxconn=max_connections, + host=self.host, + port=self.port, + user=self.user, + password=self.password, + dbname=self.dbname) + + super().__init__(logger, notification_handler) + + def close(self): + self.pool.closeall() + + def __del__(self): + self.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. + """ + + schema = input_data["schema"] + table_name = input_data["table_name"] + data = DataFrame(input_data["data"]) + + conn = self.pool.getconn() + + try: + data.to_sql(table_name, conn, schema=schema, + if_exists="append", index=False) + conn.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) + + finally: + self.pool.putconn(conn) diff --git a/scouter/activities/redis.py b/scouter/activities/redis.py new file mode 100644 index 0000000..8a2afdf --- /dev/null +++ b/scouter/activities/redis.py @@ -0,0 +1,18 @@ +from temporalio import workflow, activity + +with workflow.unsafe.imports_passed_through(): + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + from scouter.activities.base import BaseActivity + from typing import Any + from kafka import KafkaConsumer + from pandas import DataFrame + + +class Redis(BaseActivity): + def __init__(self, host: str, port: int, db: int, logger: Logger, notification_handler: NotificationHandler): + self.host = host + self.port = port + self.db = db + + super().__init__(logger, notification_handler)f diff --git a/scouter/utils/quality/filters.py b/scouter/utils/quality/filters.py new file mode 100644 index 0000000..84ebd10 --- /dev/null +++ b/scouter/utils/quality/filters.py @@ -0,0 +1,21 @@ +from pandas import DataFrame + + +def check_data_range(value, val_range: list) -> bool: + if not value: + return True + + bottom = val_range[0] + up = val_range[-1] + + return value < bottom or value > up + + +def out_of_bounds_filter(df: DataFrame, nodes_data_range: dict): + return df[df.apply(lambda x: check_data_range( + x['value'], nodes_data_range[x['tag']]), + axis=1)] + + +def null_values_filter(df: DataFrame): + return df[df['value'].isnull()] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/activities/__init__.py b/tests/activities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/activities/test_kafka.py b/tests/activities/test_kafka.py new file mode 100644 index 0000000..e5d382a --- /dev/null +++ b/tests/activities/test_kafka.py @@ -0,0 +1,64 @@ +from unittest.mock import MagicMock, patch, ANY +from pytest import fixture +from pandas import DataFrame +from scouter.activities.kafka import Kafka + + +@fixture +@patch("scouter.activities.kafka.KafkaConsumer") +def kafka(_kafka_consumer): + return Kafka( + bootstrap_servers="localhost:9092", + polling_time=1000, + group_id="test-group", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + +@patch("scouter.activities.kafka.KafkaConsumer") +def test___init__(kafka_consumer): + kafka = Kafka( + bootstrap_servers="localhost:9092", + polling_time=1000, + group_id="test-group", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert kafka.polling_time == 1000 + assert kafka.kafka_connector == kafka_consumer.return_value + + kafka_consumer.assert_called_once_with( + bootstrap_servers="localhost:9092", + auto_offset_reset="earliest", + enable_auto_commit=True, + group_id="test-group", + value_deserializer=ANY + ) + + +def test_load_from_kafka(kafka): + input_data = {"topic": "test-topic"} + + data = [ + ("test-topic", [ + MagicMock( + value=f"test-value-{i}" + ) for i in range(10) + ]) + ] + + kafka.kafka_connector.poll.return_value = MagicMock( + items=MagicMock(return_value=data) + ) + + expected = DataFrame([d.value for d in data[0][1]]).to_dict() + + result = kafka.load_from_kafka(input_data) + + assert result == expected + + kafka.kafka_connector.subscribe.assert_called_once_with(["test-topic"]) + + kafka.kafka_connector.poll.assert_called_once_with(timeout_ms=1000) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/worker/__init__.py b/tests/worker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/workflow/__init__.py b/tests/workflow/__init__.py new file mode 100644 index 0000000..e69de29