Resolve comments from wei

Author: gopidesupavan

providers/common/dataquality/docs/rules.rst

@@ -26,26 +26,10 @@ their own. They describe *what* to check; :class:`~airflow.providers.common.data
 (see :doc:`operators`) is what actually runs them. Rules are usually written by hand, but they
 can also be proposed by an LLM -- see :doc:`agents`.
 
-.. code-block:: python
-
-    from airflow.providers.common.dataquality.rules import DQRule, RuleSet
-
-    orders_ruleset = RuleSet(
-        name="orders_quality",
-        rules=(
-            DQRule(
-                name="order_id_not_null",
-                check="null_count",
-                column="order_id",
-                condition={"equal_to": 0},
-            ),
-            DQRule(
-                name="row_count_present",
-                check="row_count",
-                condition={"greater_than": 0},
-            ),
-        ),
-    )
+.. exampleinclude:: /../src/airflow/providers/common/dataquality/example_dags/example_dq_in_taskflow.py
+    :language: python
+    :start-after: [START howto_rules_dq_ruleset]
+    :end-before: [END howto_rules_dq_ruleset]
 
 ``DQRule`` fields
 ------------------

providers/common/dataquality/src/airflow/providers/common/dataquality/backends/object_storage.py

@@ -37,32 +37,39 @@
 import json
 import logging
 from datetime import datetime, timezone
-from typing import Any
+from typing import Any, TypedDict
 
 from airflow.providers.common.dataquality.results import DQRun, RuleResult, build_summary
 from airflow.sdk import ObjectStoragePath
 
 log = logging.getLogger(__name__)
 
 
+class DQPageResult(TypedDict):
+    """Paginated DQ read result."""
+
+    items: list[dict[str, Any]]
+    next_cursor: str | None
+
+
 class ObjectStorageResultsBackend:
     """Persist DQ results as JSON files via ``ObjectStoragePath``."""
 
-    def __init__(self, results_path: str, conn_id: str | None = None) -> None:
+    def __init__(self, *, results_path: str, conn_id: str | None = None) -> None:
         self.root = ObjectStoragePath(results_path, conn_id=conn_id)
 
-    def write_run(self, run: DQRun, results: list[RuleResult]) -> None:
+    def write_run(self, *, run: DQRun, results: list[RuleResult]) -> None:
         timestamp = run.started_at or datetime.now(tz=timezone.utc).isoformat()
         compact_ts = self._get_safe_key(timestamp)
-        payload = self._build_run_payload(run, results)
+        payload = self._build_run_payload(run=run, results=results)
 
-        self._write_run_file(run, timestamp[:10], compact_ts, payload)
-        self._write_task_instance_index(run, payload)
-        self._write_rule_indexes(run, results, timestamp)
+        self._write_run_file(run=run, date_part=timestamp[:10], compact_ts=compact_ts, payload=payload)
+        self._write_task_instance_index(run=run, payload=payload)
+        self._write_rule_indexes(run=run, results=results, timestamp=timestamp)
 
     def read_task_rule_history(
-        self, dag_id: str, task_id: str, rule_uid: str, limit: int = 100, before: str | None = None
-    ) -> dict[str, Any]:
+        self, *, dag_id: str, task_id: str, rule_uid: str, limit: int = 100, before: str | None = None
+    ) -> DQPageResult:
         """Return recent results for one rule produced by one task, newest first."""
         rule_dir = (
             self.root
@@ -72,11 +79,11 @@ def read_task_rule_history(
             / f"task_id={task_id}"
             / f"rule_uid={rule_uid}"
         )
-        return self._read_rule_history_dir(rule_dir, limit, before)
+        return self._read_rule_history_dir(rule_dir=rule_dir, limit=limit, before=before)
 
     def read_task_runs(
-        self, dag_id: str, task_id: str, limit: int = 50, before: str | None = None
-    ) -> dict[str, Any]:
+        self, *, dag_id: str, task_id: str, limit: int = 50, before: str | None = None
+    ) -> DQPageResult:
         """
         Return recent data quality runs for one task, newest first.
 
@@ -103,11 +110,9 @@ def read_task_runs(
             for path in sorted(date_dir.iterdir(), key=lambda p: p.name, reverse=True):
                 if not path.name.endswith(".json"):
                     continue
-                payload = self._read_json(path)
-                if payload is None:
+                if (payload := self._read_json(path)) is None:
                     continue
-                cursor = self._get_run_payload_cursor(payload)
-                if before is not None and cursor >= before:
+                if before is not None and self._get_run_payload_cursor(payload) >= before:
                     continue
                 runs.append(payload)
                 if len(runs) > limit:
@@ -121,7 +126,7 @@ def read_task_runs(
         return {"items": page, "next_cursor": next_cursor}
 
     def read_by_task_instance(
-        self, dag_id: str, task_id: str, run_id: str, map_index: int = -1
+        self, *, dag_id: str, task_id: str, run_id: str, map_index: int = -1
     ) -> dict[str, Any]:
         """Read the latest run for one task instance as ``{"run": ..., "results": ..., "summary": ...}``."""
         path = (
@@ -134,7 +139,9 @@ def read_by_task_instance(
         )
         return self._read_json_or_raise(path)
 
-    def _write_run_file(self, run: DQRun, date_part: str, compact_ts: str, payload: dict[str, Any]) -> None:
+    def _write_run_file(
+        self, *, run: DQRun, date_part: str, compact_ts: str, payload: dict[str, Any]
+    ) -> None:
         run_dir = (
             self.root
             / "runs"
@@ -146,39 +153,41 @@ def _write_run_file(self, run: DQRun, date_part: str, compact_ts: str, payload:
         run_dir.mkdir(parents=True, exist_ok=True)
         (run_dir / f"{compact_ts}__{run.run_uid}.json").write_text(json.dumps(payload, default=str))
 
-    def _write_task_instance_index(self, run: DQRun, payload: dict[str, Any]) -> None:
+    def _write_task_instance_index(self, *, run: DQRun, payload: dict[str, Any]) -> None:
         ti_dir = self.root / "runs" / "by_task_instance" / f"dag_id={run.dag_id}" / f"task_id={run.task_id}"
         ti_dir.mkdir(parents=True, exist_ok=True)
         (ti_dir / f"{self._get_safe_key(run.run_id)}__{run.map_index}.json").write_text(
             json.dumps(payload, default=str)
         )
 
-    def _write_rule_indexes(self, run: DQRun, results: list[RuleResult], timestamp: str) -> None:
+    def _write_rule_indexes(self, *, run: DQRun, results: list[RuleResult], timestamp: str) -> None:
         run_context = self._build_run_context(run)
         compact_ts = self._get_safe_key(timestamp)
         for result in results:
             payload = {"run": run_context, "result": result.to_dict()}
             self._write_rule_index(
-                self.root
-                / "rules"
-                / "by_task_rule"
-                / f"dag_id={run.dag_id}"
-                / f"task_id={run.task_id}"
-                / f"rule_uid={result.rule_uid}",
-                compact_ts,
-                run.run_uid,
-                payload,
+                rule_dir=(
+                    self.root
+                    / "rules"
+                    / "by_task_rule"
+                    / f"dag_id={run.dag_id}"
+                    / f"task_id={run.task_id}"
+                    / f"rule_uid={result.rule_uid}"
+                ),
+                compact_ts=compact_ts,
+                run_uid=run.run_uid,
+                payload=payload,
             )
 
     def _write_rule_index(
-        self, rule_dir: ObjectStoragePath, compact_ts: str, run_uid: str, payload: dict[str, Any]
+        self, *, rule_dir: ObjectStoragePath, compact_ts: str, run_uid: str, payload: dict[str, Any]
     ) -> None:
         rule_dir.mkdir(parents=True, exist_ok=True)
         (rule_dir / f"{compact_ts}__{run_uid}.json").write_text(json.dumps(payload, default=str))
 
     def _read_rule_history_dir(
-        self, rule_dir: ObjectStoragePath, limit: int, before: str | None = None
-    ) -> dict[str, Any]:
+        self, *, rule_dir: ObjectStoragePath, limit: int, before: str | None = None
+    ) -> DQPageResult:
         """
         Read rule-result records newest-first, as ``{"items": [...], "next_cursor": ...}``.
 
@@ -211,12 +220,12 @@ def _get_safe_key(value: str) -> str:
         return value.replace("/", "_").replace(":", "_").replace("+", "_")
 
     @staticmethod
-    def _build_run_payload(run: DQRun, results: list[RuleResult]) -> dict[str, Any]:
+    def _build_run_payload(*, run: DQRun, results: list[RuleResult]) -> dict[str, Any]:
         result_records = [result.to_dict() for result in results]
         return {
             "run": run.to_dict(),
             "results": result_records,
-            "summary": build_summary(run, results),
+            "summary": build_summary(run=run, results=results),
         }
 
     @staticmethod

providers/common/dataquality/src/airflow/providers/common/dataquality/engines/sql.py

@@ -57,21 +57,27 @@ class SQLDQEngine:
     def __init__(self, hook: DbApiHook) -> None:
         self.hook = hook
 
-    def measure(self, ruleset: RuleSet, table: str, partition_clause: str | None = None) -> list[Observation]:
+    def measure(
+        self, *, ruleset: RuleSet, table: str, partition_clause: str | None = None
+    ) -> list[Observation]:
         builtin_rules = [rule for rule in ruleset.rules if rule.check != CUSTOM_SQL_CHECK]
         custom_rules = [rule for rule in ruleset.rules if rule.check == CUSTOM_SQL_CHECK]
 
         observations = []
         if builtin_rules:
-            observations.extend(self._measure_builtin(builtin_rules, table, partition_clause))
+            observations.extend(
+                self._measure_builtin(rules=builtin_rules, table=table, partition_clause=partition_clause)
+            )
         for rule in custom_rules:
-            observations.append(self._measure_custom(rule, table))
+            observations.append(self._measure_custom(rule=rule, table=table))
         return observations
 
-    def build_batch_sql(self, rules: list[DQRule], table: str, partition_clause: str | None) -> str:
-        return " UNION ALL ".join(self.build_rule_sql(rule, table, partition_clause) for rule in rules)
+    def build_batch_sql(self, *, rules: list[DQRule], table: str, partition_clause: str | None) -> str:
+        return " UNION ALL ".join(
+            self.build_rule_sql(rule=rule, table=table, partition_clause=partition_clause) for rule in rules
+        )
 
-    def build_rule_sql(self, rule: DQRule, table: str, partition_clause: str | None = None) -> str:
+    def build_rule_sql(self, *, rule: DQRule, table: str, partition_clause: str | None = None) -> str:
         """Build the SQL used to measure one built-in rule."""
         expression = CHECK_SPECS[rule.check].expression.format(column=rule.column)
         predicates = [p for p in (partition_clause, rule.partition_clause) if p]
@@ -81,11 +87,14 @@ def build_rule_sql(self, rule: DQRule, table: str, partition_clause: str | None
         )
 
     def _measure_builtin(
-        self, rules: list[DQRule], table: str, partition_clause: str | None
+        self, *, rules: list[DQRule], table: str, partition_clause: str | None
     ) -> list[Observation]:
         # Built once per rule and reused below, instead of re-deriving each rule's SQL for the
         # batch, for the returned Observation, and again in the per-rule fallback.
-        rule_sql = {rule.rule_uid: self.build_rule_sql(rule, table, partition_clause) for rule in rules}
+        rule_sql = {
+            rule.rule_uid: self.build_rule_sql(rule=rule, table=table, partition_clause=partition_clause)
+            for rule in rules
+        }
         sql = " UNION ALL ".join(rule_sql.values())
         log.info("Running %d built-in checks against %s", len(rules), table)
         started = time.monotonic()
@@ -97,7 +106,7 @@ def _measure_builtin(
                 table,
             )
 
-            return [self._measure_builtin_single(rule, rule_sql[rule.rule_uid]) for rule in rules]
+            return [self._measure_builtin_single(rule=rule, sql=rule_sql[rule.rule_uid]) for rule in rules]
         elapsed_ms = (time.monotonic() - started) * 1000
         observed_by_uid = {str(row[0]): row[1] for row in records or []}
         return [
@@ -111,7 +120,7 @@ def _measure_builtin(
             for rule in rules
         ]
 
-    def _measure_builtin_single(self, rule: DQRule, sql: str) -> Observation:
+    def _measure_builtin_single(self, *, rule: DQRule, sql: str) -> Observation:
         started = time.monotonic()
         try:
             row = self.hook.get_first(sql)
@@ -126,7 +135,7 @@ def _measure_builtin_single(self, rule: DQRule, sql: str) -> Observation:
             )
         return Observation(rule=rule, observed_value=row[1], duration_ms=elapsed_ms, sql=sql)
 
-    def _measure_custom(self, rule: DQRule, table: str) -> Observation:
+    def _measure_custom(self, *, rule: DQRule, table: str) -> Observation:
         if rule.sql is None:
             raise ValueError(f"Rule {rule.name!r} has no SQL to execute")
         sql = rule.sql.replace("{table}", table)

providers/common/dataquality/src/airflow/providers/common/dataquality/example_dags/example_dq_in_taskflow.py

@@ -39,8 +39,7 @@
 READY_TABLE = "dq_taskflow_ready_orders"
 RESULTS_PATH = Path("/tmp/airflow_dq_example/results")
 
-os.environ.setdefault("AIRFLOW__COMMON_DATAQUALITY__RESULTS_PATH", f"file://{RESULTS_PATH}")
-
+# [START howto_rules_dq_ruleset]
 orders_ruleset = RuleSet(
     name="orders_taskflow_quality",
     rules=(
@@ -63,6 +62,9 @@
         ),
     ),
 )
+# [END howto_rules_dq_ruleset]
+
+os.environ.setdefault("AIRFLOW__COMMON_DATAQUALITY__RESULTS_PATH", f"file://{RESULTS_PATH}")
 
 with DAG(
     dag_id=DAG_ID,

providers/common/dataquality/src/airflow/providers/common/dataquality/execution.py

@@ -90,7 +90,11 @@ def run_quality_checks(
     resolved_ruleset = _resolve_ruleset(ruleset)
     resolved_hook = hook if hook is not None else _get_hook(conn_id, hook_params)
     started_at = datetime.now(tz=timezone.utc).isoformat()
-    observations = SQLDQEngine(resolved_hook).measure(resolved_ruleset, table, partition_clause)
+    observations = SQLDQEngine(resolved_hook).measure(
+        ruleset=resolved_ruleset,
+        table=table,
+        partition_clause=partition_clause,
+    )
     finished_at = datetime.now(tz=timezone.utc).isoformat()
     return DataQualityResult(
         ruleset=resolved_ruleset,
@@ -123,7 +127,7 @@ def persist_quality_results(
         finished_at=result.finished_at,
     )
     results = list(result.results)
-    summary = build_summary(run, results)
+    summary = build_summary(run=run, results=results)
 
     backend = get_backend_from_config()
     if backend is None:
@@ -132,7 +136,7 @@ def persist_quality_results(
         # Persistence is best-effort: an unreachable results store leaves a gap in
         # history but must not change the outcome of the check itself.
         try:
-            backend.write_run(run, results)
+            backend.write_run(run=run, results=results)
         except Exception:
             log.exception("Failed to persist data quality results; continuing")
 

providers/common/dataquality/src/airflow/providers/common/dataquality/results.py

@@ -100,7 +100,7 @@ def compute_score(results: list[RuleResult]) -> float | None:
     return round(1.0 - penalty / len(results), 4)
 
 
-def build_summary(run: DQRun, results: list[RuleResult]) -> dict[str, Any]:
+def build_summary(*, run: DQRun, results: list[RuleResult]) -> dict[str, Any]:
     """Compact run summary attached to XCom and outlet asset events."""
     return {
         "run_uid": run.run_uid,