SIENTIAPDE-1005

Add initial implementation of data processing activities and filters

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

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']

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()

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()

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)

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

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()]

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)