76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
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()
|