467 lines
21 KiB
Python
467 lines
21 KiB
Python
"""
|
|
Native Great Expectations Declarative Lakehouse Operator
|
|
Provides standardized data quality gates for Bronze, Silver, and Gold layers.
|
|
Supports:
|
|
1. Declarative GX rules (not_null, unique, in_set, row_count)
|
|
2. Native SQL assertions (Unexpected Rows queries)
|
|
3. Automated connection resolution via Airflow Connection Hook
|
|
4. OpenMetadata 2.0 Data Quality REST emission & XCom reporting
|
|
"""
|
|
import os
|
|
import re
|
|
import json
|
|
import time
|
|
import logging
|
|
from typing import Any, Dict, List, Optional, Sequence, Union
|
|
|
|
from sqlalchemy import create_engine, text
|
|
|
|
try:
|
|
from airflow.exceptions import AirflowException
|
|
from airflow.hooks.base import BaseHook
|
|
from airflow.sdk.bases.operator import BaseOperator
|
|
from airflow.sdk.definitions.context import Context
|
|
except ImportError:
|
|
try:
|
|
from airflow.exceptions import AirflowException
|
|
from airflow.hooks.base import BaseHook
|
|
from airflow.models.baseoperator import BaseOperator
|
|
from airflow.utils.context import Context
|
|
except ImportError:
|
|
class AirflowException(Exception):
|
|
pass
|
|
class BaseHook:
|
|
@classmethod
|
|
def get_connection(cls, conn_id):
|
|
return None
|
|
class BaseOperator:
|
|
def __init__(self, *args, **kwargs):
|
|
self.task_id = kwargs.get("task_id", "ge_quality_task")
|
|
self.log = logging.getLogger(self.__class__.__name__)
|
|
class Context(dict):
|
|
pass
|
|
|
|
from helpers.base_lakehouse_operator import BaseLakehouseOperator
|
|
try:
|
|
from helpers.merica_alert_hook import MericaAlertHook
|
|
except ImportError:
|
|
MericaAlertHook = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GreatExpectationsLakehouseOperator(BaseLakehouseOperator):
|
|
"""
|
|
Operator untuk validasi kualitas data secara deklaratif (Low-Code)
|
|
mengikuti standar Great Expectations 1.x dan SQL-centric Unexpected Rows.
|
|
|
|
:param expectation_suite: Nama suite pengujian (e.g. 'permits_bronze_suite')
|
|
:param table_name: Nama tabel target (e.g. 'silver.silver_environmental_permits_clean')
|
|
:param connection_id: Airflow Connection ID (Default: 'conn_postgres_source_db')
|
|
:param rules: Daftar aturan deklaratif (not_null, unique, in_set, row_count_between)
|
|
:param queries: Daftar query SQL kustom (unexpected rows assertion)
|
|
:param fail_on_error: Jika True, lempar AirflowException untuk menghentikan DAG jika ada pelanggaran
|
|
:param database_url: Opsi override connection string langsung
|
|
"""
|
|
|
|
template_fields: Sequence[str] = (
|
|
"expectation_suite",
|
|
"table_name",
|
|
"connection_id",
|
|
"rules",
|
|
"queries",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
expectation_suite: str = "default_quality_suite",
|
|
table_name: Optional[str] = None,
|
|
connection_id: str = "conn_postgres_source_db",
|
|
rules: Optional[List[Dict[str, Any]]] = None,
|
|
queries: Optional[List[Dict[str, Any]]] = None,
|
|
fail_on_error: bool = True,
|
|
database_url: Optional[str] = None,
|
|
layer: str = "silver",
|
|
**kwargs
|
|
):
|
|
super().__init__(layer=layer, engine="great_expectations", **kwargs)
|
|
self.expectation_suite = expectation_suite
|
|
self.table_name = table_name
|
|
self.connection_id = connection_id
|
|
self.rules = rules or []
|
|
self.queries = queries or []
|
|
self.fail_on_error = fail_on_error
|
|
self.database_url = database_url
|
|
|
|
def _resolve_connection_uri(self) -> str:
|
|
"""Resolve database URI via explicit param, Airflow Connection Hook, or fallback env."""
|
|
if self.database_url:
|
|
return self.database_url
|
|
|
|
try:
|
|
conn = BaseHook.get_connection(self.connection_id)
|
|
if conn:
|
|
uri = conn.get_uri()
|
|
# Standardize postgres:// to postgresql:// for SQLAlchemy
|
|
if uri.startswith("postgres://"):
|
|
uri = uri.replace("postgres://", "postgresql://", 1)
|
|
return uri
|
|
except Exception as e:
|
|
logger.debug("Airflow BaseHook lookup note: %s", e)
|
|
|
|
# Fallback to standard environment variables
|
|
fallback = os.getenv("AIRFLOW_DATABASE_URL") or os.getenv("DATABASE_URL")
|
|
if fallback:
|
|
if fallback.startswith("postgres://"):
|
|
fallback = fallback.replace("postgres://", "postgresql://", 1)
|
|
return fallback
|
|
|
|
# Default local docker postgres
|
|
return "postgresql://dlh_user:dlh_password@postgres:5432/dlh_lakehouse"
|
|
|
|
def execute(self, context: Optional[Context] = None) -> Dict[str, Any]:
|
|
"""
|
|
Mengeksekusi seluruh ekspektasi deklaratif dan query SQL asersi.
|
|
Menghasilkan laporan evaluasi standar Great Expectations.
|
|
"""
|
|
logger.info(
|
|
"🚀 Memulai Great Expectations Quality Gate: suite='%s', table='%s'",
|
|
self.expectation_suite, self.table_name or "N/A"
|
|
)
|
|
uri = self._resolve_connection_uri()
|
|
engine = create_engine(uri)
|
|
|
|
results: List[Dict[str, Any]] = []
|
|
overall_success = True
|
|
total_evaluations = 0
|
|
total_unexpected = 0
|
|
|
|
with engine.connect() as conn:
|
|
# 1. Evaluasi Aturan Deklaratif (Rules)
|
|
for idx, rule in enumerate(self.rules):
|
|
total_evaluations += 1
|
|
rule_type = (rule.get("expect") or rule.get("type") or "not_null").lower()
|
|
col = rule.get("column")
|
|
tbl = rule.get("table") or self.table_name
|
|
|
|
if not tbl:
|
|
logger.warning("Rule #%d dilewati karena 'table_name' tidak ditentukan.", idx)
|
|
continue
|
|
|
|
eval_res = self._evaluate_declarative_rule(conn, tbl, col, rule_type, rule)
|
|
results.append(eval_res)
|
|
if not eval_res["success"]:
|
|
overall_success = False
|
|
total_unexpected += eval_res.get("unexpected_count", 1)
|
|
|
|
# 2. Evaluasi Kueri SQL Kustom (Queries / Unexpected Rows)
|
|
for idx, q_item in enumerate(self.queries):
|
|
total_evaluations += 1
|
|
q_name = q_item.get("name", f"query_assertion_{idx + 1}")
|
|
sql = q_item.get("sql") or q_item.get("query")
|
|
expected = q_item.get("expected", 0)
|
|
err_msg = q_item.get("error_message") or f"Query assertion '{q_name}' failed."
|
|
|
|
if not sql or not str(sql).strip() or "silver.table" in str(sql).lower():
|
|
continue
|
|
|
|
eval_res = self._evaluate_sql_query(conn, q_name, sql, expected, err_msg)
|
|
results.append(eval_res)
|
|
if not eval_res["success"]:
|
|
overall_success = False
|
|
total_unexpected += eval_res.get("unexpected_count", 1)
|
|
|
|
success_count = sum(1 for r in results if r["success"])
|
|
success_percent = (success_count / total_evaluations * 100.0) if total_evaluations > 0 else 100.0
|
|
|
|
summary_payload = {
|
|
"status": "PASSED" if overall_success else "FAILED",
|
|
"expectation_suite": self.expectation_suite,
|
|
"table_name": self.table_name,
|
|
"overall_success": overall_success,
|
|
"total_evaluations": total_evaluations,
|
|
"successful_evaluations": success_count,
|
|
"unsuccessful_evaluations": total_evaluations - success_count,
|
|
"success_percent": round(success_percent, 2),
|
|
"total_unexpected_records": total_unexpected,
|
|
"evaluated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
"details": results
|
|
}
|
|
|
|
# Log hasil ke terminal / log task Airflow
|
|
if overall_success:
|
|
logger.info(
|
|
"✅ Great Expectations Suite '%s' LULUS 100%% (%d/%d aturan berhasil).",
|
|
self.expectation_suite, success_count, total_evaluations
|
|
)
|
|
else:
|
|
logger.error(
|
|
"❌ Great Expectations Suite '%s' GAGAL: %d aturan melanggar integritas (Total anomali: %d).",
|
|
self.expectation_suite, total_evaluations - success_count, total_unexpected
|
|
)
|
|
|
|
# 3. Publikasi Metrik Kualitas Data ke OpenMetadata REST API
|
|
self._publish_openmetadata_data_quality(summary_payload)
|
|
|
|
# 3.1 Pemicu Notifikasi Alerting Multi-Saluran jika Breach Kualitas Terjadi
|
|
if not overall_success and MericaAlertHook:
|
|
try:
|
|
failed_rules = [{"rule": r["expectation_name"], "column": r.get("column", "table")} for r in results if not r.get("success")]
|
|
dag = context.get('dag') if isinstance(context, dict) else None
|
|
dag_params = getattr(dag, 'params', {}) if dag else {}
|
|
alert_cfg = dag_params.get('alerting') if isinstance(dag_params, dict) else None
|
|
MericaAlertHook.on_quality_gate_breach(context, failed_rules, success_percent, alert_config=alert_cfg)
|
|
except Exception as alert_err:
|
|
logger.warning("MericaAlertHook quality breach notification note: %s", alert_err)
|
|
|
|
# 4. Ambil Tindakan Kegagalan (Fail-Fast Policy)
|
|
if not overall_success and self.fail_on_error:
|
|
failed_details = [r["expectation_name"] for r in results if not r["success"]]
|
|
raise AirflowException(
|
|
f"Great Expectations Quality Gate GAGAL pada suite '{self.expectation_suite}'! "
|
|
f"Aturan gagal: {failed_details}. Pipeline dihentikan untuk menjaga integritas data."
|
|
)
|
|
|
|
return summary_payload
|
|
|
|
def _evaluate_declarative_rule(
|
|
self, conn, table: str, column: Optional[str], rule_type: str, rule_cfg: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
"""Eksekusi aturan deklaratif standar Great Expectations."""
|
|
rule_name = f"expect_{rule_type}" if not rule_type.startswith("expect_") else rule_type
|
|
|
|
try:
|
|
if any(k in rule_type for k in ("not_null", "not_be_null", "non_null")):
|
|
sql = f"SELECT count(*) FROM {table} WHERE {column} IS NULL;"
|
|
res = conn.execute(text(sql)).scalar() or 0
|
|
success = (res == 0)
|
|
return {
|
|
"expectation_name": f"{rule_name}({column})",
|
|
"type": "column_values_to_not_be_null",
|
|
"column": column,
|
|
"table": table,
|
|
"success": success,
|
|
"unexpected_count": res,
|
|
"message": "Passed: No NULL values detected." if success else f"Failed: Found {res} NULL values in column '{column}'."
|
|
}
|
|
|
|
elif any(k in rule_type for k in ("unique", "distinct", "be_unique")):
|
|
sql = f"SELECT count(*) FROM (SELECT {column}, count(*) FROM {table} GROUP BY {column} HAVING count(*) > 1) dups;"
|
|
res = conn.execute(text(sql)).scalar() or 0
|
|
success = (res == 0)
|
|
return {
|
|
"expectation_name": f"{rule_name}({column})",
|
|
"type": "column_values_to_be_unique",
|
|
"column": column,
|
|
"table": table,
|
|
"success": success,
|
|
"unexpected_count": res,
|
|
"message": "Passed: All values are unique." if success else f"Failed: Found {res} duplicate entries for column '{column}'."
|
|
}
|
|
|
|
elif any(k in rule_type for k in ("in_set", "be_in_set", "allowed_values")):
|
|
allowed_vals = rule_cfg.get("allowed_values") or rule_cfg.get("value_set") or []
|
|
formatted_set = ", ".join([f"'{v}'" for v in allowed_vals])
|
|
sql = f"SELECT count(*) FROM {table} WHERE {column} NOT IN ({formatted_set}) OR {column} IS NULL;"
|
|
res = conn.execute(text(sql)).scalar() or 0
|
|
success = (res == 0)
|
|
return {
|
|
"expectation_name": f"{rule_name}({column})",
|
|
"type": "column_values_to_be_in_set",
|
|
"column": column,
|
|
"table": table,
|
|
"allowed_set": allowed_vals,
|
|
"success": success,
|
|
"unexpected_count": res,
|
|
"message": f"Passed: All values in allowed set {allowed_vals}." if success else f"Failed: Found {res} records outside allowed domain."
|
|
}
|
|
|
|
elif any(k in rule_type for k in ("row_count", "rowcount", "table_row_count")):
|
|
min_v = rule_cfg.get("min_value", 1)
|
|
max_v = rule_cfg.get("max_value")
|
|
sql = f"SELECT count(*) FROM {table};"
|
|
cnt = conn.execute(text(sql)).scalar() or 0
|
|
success = (cnt >= min_v) and (max_v is None or cnt <= max_v)
|
|
return {
|
|
"expectation_name": f"{rule_name}({table})",
|
|
"type": "table_row_count_to_be_between",
|
|
"table": table,
|
|
"observed_row_count": cnt,
|
|
"min_value": min_v,
|
|
"max_value": max_v,
|
|
"success": success,
|
|
"unexpected_count": 0 if success else 1,
|
|
"message": f"Passed: Row count is {cnt}." if success else f"Failed: Observed row count {cnt} out of bounds [{min_v}, {max_v}]."
|
|
}
|
|
|
|
else:
|
|
return {
|
|
"expectation_name": rule_name,
|
|
"success": True,
|
|
"message": f"Info: Custom rule type '{rule_type}' passed by default."
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error("Error evaluating rule '%s': %s", rule_name, e)
|
|
return {
|
|
"expectation_name": rule_name,
|
|
"success": False,
|
|
"error": str(e),
|
|
"unexpected_count": 1,
|
|
"message": f"Execution error: {e}"
|
|
}
|
|
|
|
def _evaluate_sql_query(
|
|
self, conn, query_name: str, sql_query: str, expected_val: Any, err_msg: str
|
|
) -> Dict[str, Any]:
|
|
"""Eksekusi query SQL asersi kustom (Unexpected Rows / Count evaluation)."""
|
|
try:
|
|
res = conn.execute(text(sql_query)).scalar()
|
|
res_val = int(res) if res is not None else 0
|
|
exp_val = int(expected_val)
|
|
success = (res_val == exp_val)
|
|
|
|
return {
|
|
"expectation_name": query_name,
|
|
"type": "custom_sql_query",
|
|
"sql_query": sql_query,
|
|
"observed_value": res_val,
|
|
"expected_value": exp_val,
|
|
"success": success,
|
|
"unexpected_count": abs(res_val - exp_val) if not success else 0,
|
|
"message": f"Passed: Query '{query_name}' matched expected value ({exp_val})." if success else f"{err_msg} (Observed: {res_val}, Expected: {exp_val})"
|
|
}
|
|
except Exception as e:
|
|
logger.error("Error executing query assertion '%s': %s", query_name, e)
|
|
return {
|
|
"expectation_name": query_name,
|
|
"type": "custom_sql_query",
|
|
"sql_query": sql_query,
|
|
"success": False,
|
|
"error": str(e),
|
|
"unexpected_count": 1,
|
|
"message": f"SQL Execution error: {e}"
|
|
}
|
|
|
|
def _publish_openmetadata_data_quality(self, summary: Dict[str, Any]):
|
|
"""Publish test results to OpenMetadata REST API (/api/v1/dataQuality/testCases)."""
|
|
import urllib.request
|
|
import urllib.parse
|
|
om_host = os.getenv("OPENMETADATA_HOST", "http://openmetadata:8585").rstrip("/")
|
|
|
|
try:
|
|
# Login to get token
|
|
om_user = os.getenv("OPENMETADATA_USERNAME", "admin@open-metadata.org")
|
|
om_pass = os.getenv("OPENMETADATA_PASSWORD", "admin")
|
|
if om_pass == "admin":
|
|
b64_pass = "YWRtaW4="
|
|
else:
|
|
import base64
|
|
b64_pass = base64.b64encode(om_pass.encode("utf-8")).decode("utf-8")
|
|
|
|
login_req = urllib.request.Request(
|
|
f"{om_host}/api/v1/users/login",
|
|
data=json.dumps({"email": om_user, "password": b64_pass}).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(login_req, timeout=3) as resp:
|
|
token = json.loads(resp.read().decode("utf-8")).get("accessToken")
|
|
|
|
if not token:
|
|
return
|
|
|
|
db_svc = os.getenv("OPENMETADATA_DB_SERVICE_NAME", "dlh_postgres_service")
|
|
db_name = os.getenv("POSTGRES_DB", "dlh_lakehouse")
|
|
raw_tbl = str(self.table_name or getattr(self, "ge_table_name", None) or "table")
|
|
parts = raw_tbl.split(".")
|
|
if len(parts) == 2:
|
|
schema = parts[0]
|
|
tbl_name = parts[1]
|
|
else:
|
|
tbl_name = parts[-1]
|
|
if tbl_name.startswith("gold_"):
|
|
schema = "gold"
|
|
elif tbl_name.startswith("silver_"):
|
|
schema = "silver"
|
|
elif tbl_name.startswith("bronze_"):
|
|
schema = "bronze"
|
|
elif "permit" in tbl_name.lower():
|
|
schema = "bronze_lakehouse"
|
|
else:
|
|
schema = "bronze"
|
|
table_fqn = f"{db_svc}.{db_name}.{schema}.{tbl_name}"
|
|
|
|
# 1. Ensure Table Test Suite
|
|
suite_payload = {
|
|
"name": f"{table_fqn}.testSuite",
|
|
"displayName": f"Quality Suite for {tbl_name}",
|
|
"basicEntityReference": table_fqn
|
|
}
|
|
suite_req = urllib.request.Request(
|
|
f"{om_host}/api/v1/dataQuality/testSuites/basic",
|
|
data=json.dumps(suite_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(suite_req, timeout=3) as _:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# 2. Ensure Test Case on Table EntityLink
|
|
tc_name = re.sub(r'[^a-zA-Z0-9_.-]', '_', f"{self.task_id}_quality").lower()
|
|
case_payload = {
|
|
"name": tc_name,
|
|
"displayName": f"GE Gate {self.task_id}",
|
|
"description": f"Great Expectations suite '{self.expectation_suite}' on {table_fqn}",
|
|
"entityLink": f"<#E::table::{table_fqn}>",
|
|
"testDefinition": "tableCustomSQLQuery",
|
|
"parameterValues": [
|
|
{"name": "sqlExpression", "value": str(self.sql_assertion or getattr(self, "sql", "SELECT 1"))}
|
|
]
|
|
}
|
|
tc_req = urllib.request.Request(
|
|
f"{om_host}/api/v1/dataQuality/testCases",
|
|
data=json.dumps(case_payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(tc_req, timeout=3) as _:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# 3. Post Test Case Result to OpenMetadata Table Test Case
|
|
tc_fqn = f"{table_fqn}.{tc_name}"
|
|
payload = {
|
|
"timestamp": int(time.time() * 1000),
|
|
"testCaseStatus": "Success" if summary["overall_success"] else "Failed",
|
|
"result": f"Suite '{self.expectation_suite}': {summary['successful_evaluations']}/{summary['total_evaluations']} rules passed.",
|
|
"passedRows": summary.get("successful_evaluations", 1),
|
|
"failedRows": summary.get("total_unexpected_records", 0),
|
|
"passedRowsPercentage": float(summary.get("success_percent", 100.0)),
|
|
"testResultValue": [
|
|
{"name": "evaluated_count", "value": str(summary["total_evaluations"])},
|
|
{"name": "unexpected_count", "value": str(summary["total_unexpected_records"])},
|
|
{"name": "success_percent", "value": str(summary["success_percent"])}
|
|
]
|
|
}
|
|
|
|
req = urllib.request.Request(
|
|
f"{om_host}/api/v1/dataQuality/testCases/testCaseResults/{urllib.parse.quote(tc_fqn, safe='')}",
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=4) as resp:
|
|
logger.info("✅ Table Data Quality metric emitted to OpenMetadata: %s (%s)", tc_fqn, payload["testCaseStatus"])
|
|
except Exception:
|
|
pass
|
|
except Exception as e:
|
|
logger.debug("OpenMetadata DQ emission note: %s", e)
|