kaxil commented on code in PR #69575: URL: https://github.com/apache/airflow/pull/69575#discussion_r3574638256
########## providers/common/dataquality/src/airflow/providers/common/dataquality/engines/sql.py: ########## @@ -0,0 +1,137 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""SQL execution engine: compiles a ruleset into check queries run through a DB-API hook.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.dataquality.rules import CUSTOM_SQL_CHECK +from airflow.providers.common.dataquality.rules.checks import CHECK_SPECS + +if TYPE_CHECKING: + from airflow.providers.common.dataquality.rules import DQRule, RuleSet + from airflow.providers.common.sql.hooks.sql import DbApiHook + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Observation: + """Raw value a rule observed, before its condition is evaluated.""" + + rule: DQRule + observed_value: Any = None + duration_ms: float | None = None + error_message: str | None = None + sql: str | None = None + + +class SQLDQEngine: + """ + Runs built-in checks as a single UNION ALL query and custom SQL rules individually. + + Table, column, and partition-clause values come from the Dag author and are interpolated + into SQL the same way the ``common.sql`` check operators do — they are trusted input. + """ + + check_sql_template = "SELECT '{rule_uid}' AS rule_uid, {expression} AS observed FROM {table}{where}" + + def __init__(self, hook: DbApiHook) -> None: + self.hook = hook + + 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)) + for rule in custom_rules: + observations.append(self._measure_custom(rule, 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_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] + where = f" WHERE {' AND '.join(predicates)}" if predicates else "" + return self.check_sql_template.format( + rule_uid=rule.rule_uid, expression=expression, table=table, where=where + ) + + def _measure_builtin( + self, rules: list[DQRule], table: str, partition_clause: str | None + ) -> list[Observation]: + sql = self.build_batch_sql(rules, table, partition_clause) + log.info("Running %d built-in checks against %s", len(rules), table) + started = time.monotonic() + try: + records = self.hook.get_records(sql) + except Exception as e: Review Comment: Since all built-ins run as one UNION ALL, a single bad rule (say an LLM-generated `MIN(no_such_column)`) fails the whole batch here, and every other rule gets recorded as ERROR with that unrelated database message. `_raise_for_failures` then fails the task on ERROR regardless of `fail_on`, so one bad generated rule takes down rules that would have passed, and their persisted history shows an error that isn't theirs -- which cuts against "Every rule is evaluated and recorded regardless of the task outcome" in the operator docstring. A per-rule fallback when the batch errors would keep the batching win on the happy path and isolate the blast radius. Related: `duration_ms` on each observation is the whole batch's elapsed time, which will mislead once the UI charts per-rule durations. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/backends/object_storage.py: ########## @@ -0,0 +1,252 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Object-storage results backend. + +Each DQ check writes a keyed JSON document plus read indexes optimized for the UI: + + runs/by_task/dag_id=<dag>/task_id=<task>/date=<2026-07-04>/<run_uid>.json + Canonical run record: ``{"run": ..., "results": [...], "summary": ...}``. + + runs/by_task_instance/dag_id=<dag>/task_id=<task>/<safe_run_id>__<map_index>.json + Latest result for a task-instance page. Last write wins across retries. + + rules/by_task_rule/dag_id=<dag>/task_id=<task>/rule_uid=<uid>/<started_at>__<run_uid>.json + One rule result plus run context: ``{"run": ..., "result": ...}``. + +The duplicate files are intentional read indexes: DQ tasks write once, while the UI reads +many times. Keeping these indexes avoids scanning all task runs for common UI views. +""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime, timezone +from typing import Any + +from airflow.providers.common.dataquality.results import DQRun, RuleResult, build_summary +from airflow.sdk import ObjectStoragePath + +log = logging.getLogger(__name__) + + +class ObjectStorageResultsBackend: + """Persist DQ results as JSON files via ``ObjectStoragePath``.""" + + 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: + timestamp = run.started_at or datetime.now(tz=timezone.utc).isoformat() + payload = self._build_run_payload(run, results) + + self._write_run_file(run, timestamp[:10], payload) + self._write_task_instance_index(run, payload) + self._write_rule_indexes(run, results, 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]: + """Return recent results for one rule produced by one task, newest first.""" + rule_dir = ( + self.root + / "rules" + / "by_task_rule" + / f"dag_id={dag_id}" + / f"task_id={task_id}" + / f"rule_uid={rule_uid}" + ) + return self._read_rule_history_dir(rule_dir, limit, before) + + def read_task_runs( + self, dag_id: str, task_id: str, limit: int = 50, before: str | None = None + ) -> dict[str, Any]: + """ + Return recent data quality runs for one task, newest first. + + Returns ``{"items": [...], "next_cursor": ...}``. ``before`` is the opaque + ``next_cursor`` from the previous page, so "load more" only reads runs it hasn't + shown yet. ``next_cursor`` is set by reading one extra entry past ``limit``: cheap + on every page except the last one, where there's no way to confirm history is + exhausted without walking to the end. + + ``date=`` partition names sort correctly as plain strings, so directories are walked + newest-first and scanning stops as soon as ``limit + 1`` matching runs have been + collected — a task with years of history doesn't pay for a full scan on every page. + """ + task_dir = self.root / "runs" / "by_task" / f"dag_id={dag_id}" / f"task_id={task_id}" + if not task_dir.exists(): + return {"items": [], "next_cursor": None} + + date_dirs = sorted((path for path in task_dir.iterdir() if path.is_dir()), reverse=True) + runs = [] + for date_dir in date_dirs: + for path in date_dir.iterdir(): Review Comment: The early exit here breaks newest-first within a date partition: `date_dir.iterdir()` has no ordering and the filenames are random `run_uid`s, so once a single day holds more than `limit` runs (a mapped check writes one file per map_index into the same partition, and retries add more), this collects an arbitrary `limit+1` subset, sorts only that subset, and sets `next_cursor` past runs it never read -- the genuinely newest runs become permanently unreachable on later pages. The rule-history path avoids this only because its filenames are timestamp-prefixed, so its sort-then-scan is chronological. I'd prefix the by_task filename with `started_at` the same way (`{compact_ts}__{run_uid}.json`) and sort each partition before the early exit. Worth doing before the first release since this layout is the persisted format -- changing filenames later strands existing history. The same-day test at `test_object_storage.py:294` only asserts count and cursor presence; once fixed it should asse rt *which* run comes back, like the cross-date test above it does. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/rules/rule.py: ########## @@ -0,0 +1,289 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Declarative data quality rule model: rules are data, not code.""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from typing import Any, cast + +import yaml +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from airflow.providers.common.dataquality.exceptions import DQRuleValidationError +from airflow.providers.common.dataquality.rules.checks import CHECK_SPECS, Dimension + +CUSTOM_SQL_CHECK = "custom_sql" + + +class Severity(str, Enum): + """Severity used to decide whether a failing rule fails the task.""" + + WARN = "warn" + ERROR = "error" + + +class Condition(BaseModel): + """ + Pass/fail condition evaluated against a rule's observed value. + + Uses the same grammar as the ``common.sql`` check operators: ``equal_to``, + ``greater_than``, ``less_than``, ``geq_to``, ``leq_to``, plus a percentage + ``tolerance`` that widens comparisons. + + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + equal_to: float | None = None + greater_than: float | None = None + less_than: float | None = None + geq_to: float | None = None + leq_to: float | None = None + tolerance: float | None = None + + @model_validator(mode="after") + def _validate_comparisons(self) -> Condition: + comparisons = { + "equal_to": self.equal_to, + "greater_than": self.greater_than, + "less_than": self.less_than, + "geq_to": self.geq_to, + "leq_to": self.leq_to, + } + set_comparisons = {name for name, value in comparisons.items() if value is not None} + if not set_comparisons: + raise ValueError(f"Condition needs at least one comparison out of: {', '.join(comparisons)}") + if self.equal_to is not None and len(set_comparisons) > 1: + raise ValueError("equal_to cannot be combined with other comparisons") + return self + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Condition: + return cls(**data) + + def to_dict(self) -> dict[str, float]: + data = { + "equal_to": self.equal_to, + "greater_than": self.greater_than, + "less_than": self.less_than, + "geq_to": self.geq_to, + "leq_to": self.leq_to, + "tolerance": self.tolerance, + } + return {k: v for k, v in data.items() if v is not None} + + def evaluate(self, observed: Any) -> bool: + if observed is None: + return False + value = float(observed) + if self.equal_to is not None: + if self.tolerance is not None: + delta = abs(self.equal_to) * self.tolerance + return self.equal_to - delta <= value <= self.equal_to + delta + return value == self.equal_to + geq_to = self.geq_to + greater_than = self.greater_than + leq_to = self.leq_to + less_than = self.less_than + if self.tolerance is not None: + if geq_to is not None: + geq_to -= abs(geq_to) * self.tolerance + if greater_than is not None: + greater_than -= abs(greater_than) * self.tolerance + if leq_to is not None: + leq_to += abs(leq_to) * self.tolerance + if less_than is not None: + less_than += abs(less_than) * self.tolerance + return ( + (greater_than is None or value > greater_than) + and (less_than is None or value < less_than) + and (geq_to is None or value >= geq_to) + and (leq_to is None or value <= leq_to) + ) + + +class DQRule(BaseModel): + """ + A single, named data quality rule. + + :param name: Rule name, unique within its ruleset. + :param check: One of the built-in checks (see ``CHECK_SPECS``) or ``custom_sql``. Built-in + checks are plain ANSI SQL and are not guaranteed to work against every ``DbApiHook``; + see "Supported checks and databases" in the provider docs. Use ``custom_sql`` if a + built-in check doesn't fit your database's dialect. + :param condition: Pass condition for the observed value (``Condition`` or its dict form). + Optional only for checks whose ``CheckSpec.default_condition`` is set; every current + built-in check requires one explicitly. + :param column: Target column; required for column-level built-in checks. + :param sql: SQL statement returning a single scalar; required for ``custom_sql``. + May reference the target table as ``{table}``. + :param severity: ``error`` (fails the task by default) or ``warn`` (recorded only). + :param partition_clause: Extra predicate ANDed into the check's WHERE clause. + :param previous_name: Set when renaming a rule, to keep its history continuous. + :param description: Human-readable description shown in results and the UI. When omitted, + the provider generates a short default description from the rule and condition. + + Invalid input raises pydantic's own :class:`~pydantic.ValidationError`. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + check: str = Field(description="One of the built-in checks, or custom_sql.") + condition: Condition | dict[str, Any] | None = None Review Comment: Probed this: `condition={"equal_to": 0}` coerces to `Condition` as intended, but `{"nonsense": 1}` or `{"tolerance": 0.1}` fail Condition validation and fall into the `dict[str, Any]` branch, surviving model construction as raw dicts -- then `rule_uid`/`to_dict()` crash with AttributeError instead of the ValidationError the docstring promises. Exactly the shape of mistake a generated ruleset makes. Typing this `Condition | None` gives you the dict coercion for free and rejects garbage at construction; the generated schema also loses its `additionalProperties: true` branch. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/assets.py: ########## @@ -0,0 +1,172 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Asset-level data quality declarations. + +Quality configuration lives inside ``Asset.extra`` under the ``airflow.dataquality`` key, so it is +serialized with the Dag and needs no Airflow core changes. The rules travel with the asset +definition instead of being scattered across the Dags that check it. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.dataquality.exceptions import DQRuleValidationError +from airflow.providers.common.dataquality.rules import RuleSet +from airflow.sdk import task + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from airflow.sdk import Asset + from airflow.sdk.execution_time.comms import AssetEventDagRunReferenceResult + +log = logging.getLogger(__name__) + +DQ_EXTRA_KEY = "airflow.dataquality" +DQ_RESULT_EXTRA_KEY = "airflow.dataquality.result" + + +def asset_quality( + asset: Asset, + *, + ruleset: RuleSet | dict[str, Any] | str, + conn_id: str | None = None, + table: str | None = None, +) -> Asset: + """ + Attach data quality configuration to an asset, returning the same asset. + + :param asset: The asset the rules describe. + :param ruleset: A :class:`~airflow.providers.common.dataquality.rules.RuleSet`, its dict form, or a path + to a YAML ruleset file (resolved eagerly, at Dag-parse time). + :param conn_id: Default connection a check operator should use for this asset. + :param table: Default table to check; falls back to the asset name when unset. + + Usage:: + + orders = asset_quality( + Asset("orders", uri="postgres://warehouse/analytics/orders"), + ruleset=rules, + conn_id="warehouse", + table="analytics.orders", + ) + check = DQCheckOperator(task_id="dq", asset=orders) + """ + if isinstance(ruleset, str): + ruleset = RuleSet.from_file(ruleset) + elif isinstance(ruleset, dict): + ruleset = RuleSet.from_dict(ruleset) + config: dict[str, Any] = {"ruleset": ruleset.to_dict()} + if conn_id: + config["conn_id"] = conn_id + if table: + config["table"] = table + asset.extra[DQ_EXTRA_KEY] = config + return asset + + +def get_asset_quality_config(asset: Asset) -> dict[str, Any] | None: + """Return the raw ``airflow.dataquality`` config attached to an asset, if any.""" + config = asset.extra.get(DQ_EXTRA_KEY) + if not isinstance(config, dict): + return None + return config + + +def get_asset_ruleset(asset: Asset) -> RuleSet: + """Return the ruleset attached to an asset, raising when none is attached.""" + config = get_asset_quality_config(asset) + if not config or "ruleset" not in config: + raise DQRuleValidationError( + f"Asset {asset.name!r} has no data quality config; attach one with asset_quality()" + ) + return RuleSet.from_dict(config["ruleset"]) + + +def _quality_score_passes( + asset: Asset, + min_score: float, + triggering_asset_events: Mapping[Asset, Sequence[AssetEventDagRunReferenceResult]], +) -> bool: + """Pure decision logic behind :func:`require_quality`, kept separate so it is testable without a Dag.""" + events = triggering_asset_events.get(asset, []) + if not events: + log.warning( + "require_quality(%s): no triggering event for this run; skipping downstream tasks", + asset.name, + ) + return False + + summary = events[-1].extra.get(DQ_RESULT_EXTRA_KEY) Review Comment: `consumed_asset_events` is a plain M2M backref with no `order_by`, and the execution API loads it with a bare `joinedload`, so the list order here is whatever the database returns -- `events[-1]` isn't guaranteed to be the newest event. And when several producer runs coalesce into one consumer run, gating only on one event means a low-scoring event slips through if a higher-scoring one sits after it in the list. I'd pick the event explicitly by `timestamp`, and consider whether the safer default is requiring every triggering event's score to pass, with latest-only as an opt-in. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/skills/dataquality-rule-authoring/SKILL.md: ########## @@ -0,0 +1,173 @@ +--- +name: dataquality-rule-authoring +description: "Generate a RuleSet/DQRule JSON payload for Apache Airflow's dataquality provider from a table's column definitions. Use this whenever asked to write, generate, or suggest data quality rules, checks, or a ruleset for a table or dataset -- especially when the output is structured JSON handed to DQCheckOperator, @task.dq_check, asset_quality(), or RuleSet.from_file()." +--- + +<!-- SPDX-License-Identifier: Apache-2.0 + https://www.apache.org/licenses/LICENSE-2.0 --> + +# dataquality rule authoring + +You are producing a **RuleSet**: a JSON object naming a table's data quality checks. This +document is the schema and the full list of valid values -- do not guess field names or check +names, and do not invent checks that aren't in the catalog below. + +## Output shape + +```json +{ + "name": "orders_quality", + "rules": [ + { + "name": "order_id_not_null", + "check": "null_count", + "column": "order_id", + "condition": {"equal_to": 0} + } + ] +} +``` + +`name` (ruleset name) and `rules` (array) are the only top-level keys. Every rule name must be +unique within the ruleset. + +## `DQRule` fields + +| Field | Required | Notes | +|---|---|---| +| `name` | yes | Unique within the ruleset. Shows up in data quality results and logs. | +| `check` | yes | One of the catalog names below, or `"custom_sql"`. Exact string match -- there is no fuzzy matching. | +| `condition` | yes* | See `Condition` grammar below. *Every catalog check currently requires one explicitly (no defaults) -- always include it. | +| `column` | check-dependent | Required for every catalog check except `row_count`. Not used with `custom_sql`. | +| `sql` | `custom_sql` only | A SQL statement returning a single scalar. May reference the checked table as `{table}`. Invalid together with any catalog check. | +| `severity` | no | `"error"` (default) or `"warn"`. `"error"` fails the task on a failing rule (subject to the operator's `fail_on`); `"warn"` only records it. | +| `partition_clause` | no | Extra SQL predicate ANDed into this rule's `WHERE` clause, e.g. `"region = 'EU'"`. | +| `previous_name` | no | Only set when told a rule is being renamed, to keep its history continuous. | +| `description` | no | Human-readable text shown in data quality results. Use it when a clear business meaning is known; otherwise omit it and let Airflow generate a default description. | +| `dimension` | no | One of `completeness`, `uniqueness`, `validity`, `freshness`, `volume`, `consistency`. Defaults to the check's catalog dimension (`validity` for `custom_sql`) -- only set this explicitly when a `custom_sql` rule measures something the default doesn't capture (e.g. a freshness check written as `custom_sql` should set `"dimension": "freshness"`). Leave it unset for every built-in check. | + +**Do not add any key not listed above.** The schema rejects unrecognized fields. + +## Built-in check catalog + +| `check` | SQL expression | needs `column` | typical use | +|---|---|---|---| +| `null_count` | `SUM(CASE WHEN {column} IS NULL THEN 1 ELSE 0 END)` | yes | count of nulls in a column | +| `null_ratio` | `SUM(CASE WHEN {column} IS NULL THEN 1.0 ELSE 0.0 END) / COUNT(*)` | yes | fraction of nulls, 0.0-1.0 | +| `distinct_count` | `COUNT(DISTINCT {column})` | yes | number of distinct values | +| `unique_violations` | `COUNT({column}) - COUNT(DISTINCT {column})` | yes | 0 means the column is fully unique | +| `min` | `MIN({column})` | yes | minimum value | +| `max` | `MAX({column})` | yes | maximum value | +| `mean` | `AVG({column})` | yes | average value | +| `row_count` | `COUNT(*)` | no | total row count -- omit `column` entirely | + +These are the *only* built-in check names. Anything else -- comparing two columns, joining +another table, checking a format/regex, freshness against `now()`, referential integrity -- must +be written as `custom_sql`. + +## `Condition` grammar + +`condition` is an object with these optional numeric keys; at least one is required: + +- `equal_to` -- exact match. Cannot be combined with any other key below. +- `greater_than` / `geq_to` -- lower bound, exclusive / inclusive. +- `less_than` / `leq_to` -- upper bound, exclusive / inclusive. +- `tolerance` -- a fraction (e.g. `0.1` for 10%) that widens the comparison bounds. + +Combine `greater_than`/`less_than`/`geq_to`/`leq_to` freely to express a range, e.g. +`{"geq_to": 0, "leq_to": 100}`. Never combine `equal_to` with another comparison key (other than +`tolerance`). + +## `custom_sql`: when a catalog check doesn't fit Review Comment: The skill steers models toward emitting `custom_sql`, and the engine executes that string verbatim through the connection -- a generated `DELETE FROM {table} RETURNING 1` satisfies the scalar contract while mutating data, so a model error or prompt injection inherits the connection's write privileges. For Dag-author-written rules this matches the common.sql trust model, but the skill is specifically the path where the author isn't writing the SQL. I'd have the skill instruct models to prefer built-in checks and treat `custom_sql` as requiring human review, and recommend read-only credentials for DQ connections in the docs. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/rules/rule.py: ########## @@ -0,0 +1,289 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Declarative data quality rule model: rules are data, not code.""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from typing import Any, cast + +import yaml +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from airflow.providers.common.dataquality.exceptions import DQRuleValidationError +from airflow.providers.common.dataquality.rules.checks import CHECK_SPECS, Dimension + +CUSTOM_SQL_CHECK = "custom_sql" + + +class Severity(str, Enum): + """Severity used to decide whether a failing rule fails the task.""" + + WARN = "warn" + ERROR = "error" + + +class Condition(BaseModel): + """ + Pass/fail condition evaluated against a rule's observed value. + + Uses the same grammar as the ``common.sql`` check operators: ``equal_to``, + ``greater_than``, ``less_than``, ``geq_to``, ``leq_to``, plus a percentage + ``tolerance`` that widens comparisons. + + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + equal_to: float | None = None + greater_than: float | None = None + less_than: float | None = None + geq_to: float | None = None + leq_to: float | None = None + tolerance: float | None = None + + @model_validator(mode="after") + def _validate_comparisons(self) -> Condition: + comparisons = { + "equal_to": self.equal_to, + "greater_than": self.greater_than, + "less_than": self.less_than, + "geq_to": self.geq_to, + "leq_to": self.leq_to, + } + set_comparisons = {name for name, value in comparisons.items() if value is not None} + if not set_comparisons: + raise ValueError(f"Condition needs at least one comparison out of: {', '.join(comparisons)}") + if self.equal_to is not None and len(set_comparisons) > 1: + raise ValueError("equal_to cannot be combined with other comparisons") + return self + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Condition: + return cls(**data) + + def to_dict(self) -> dict[str, float]: + data = { + "equal_to": self.equal_to, + "greater_than": self.greater_than, + "less_than": self.less_than, + "geq_to": self.geq_to, + "leq_to": self.leq_to, + "tolerance": self.tolerance, + } + return {k: v for k, v in data.items() if v is not None} + + def evaluate(self, observed: Any) -> bool: + if observed is None: + return False + value = float(observed) + if self.equal_to is not None: + if self.tolerance is not None: + delta = abs(self.equal_to) * self.tolerance + return self.equal_to - delta <= value <= self.equal_to + delta + return value == self.equal_to + geq_to = self.geq_to + greater_than = self.greater_than + leq_to = self.leq_to + less_than = self.less_than + if self.tolerance is not None: + if geq_to is not None: + geq_to -= abs(geq_to) * self.tolerance + if greater_than is not None: + greater_than -= abs(greater_than) * self.tolerance + if leq_to is not None: + leq_to += abs(leq_to) * self.tolerance + if less_than is not None: + less_than += abs(less_than) * self.tolerance + return ( + (greater_than is None or value > greater_than) + and (less_than is None or value < less_than) + and (geq_to is None or value >= geq_to) + and (leq_to is None or value <= leq_to) + ) + + +class DQRule(BaseModel): + """ + A single, named data quality rule. + + :param name: Rule name, unique within its ruleset. + :param check: One of the built-in checks (see ``CHECK_SPECS``) or ``custom_sql``. Built-in + checks are plain ANSI SQL and are not guaranteed to work against every ``DbApiHook``; + see "Supported checks and databases" in the provider docs. Use ``custom_sql`` if a + built-in check doesn't fit your database's dialect. + :param condition: Pass condition for the observed value (``Condition`` or its dict form). + Optional only for checks whose ``CheckSpec.default_condition`` is set; every current + built-in check requires one explicitly. + :param column: Target column; required for column-level built-in checks. + :param sql: SQL statement returning a single scalar; required for ``custom_sql``. + May reference the target table as ``{table}``. + :param severity: ``error`` (fails the task by default) or ``warn`` (recorded only). + :param partition_clause: Extra predicate ANDed into the check's WHERE clause. + :param previous_name: Set when renaming a rule, to keep its history continuous. + :param description: Human-readable description shown in results and the UI. When omitted, + the provider generates a short default description from the rule and condition. + + Invalid input raises pydantic's own :class:`~pydantic.ValidationError`. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str + check: str = Field(description="One of the built-in checks, or custom_sql.") + condition: Condition | dict[str, Any] | None = None + column: str | None = None + sql: str | None = None + severity: Severity = Severity.ERROR + partition_clause: str | None = None + previous_name: str | None = None + description: str | None = Field( + default=None, + description=( + "Human-readable description shown in results and the UI. When omitted, the provider " + "generates a short default description from the rule and condition." + ), + ) + dimension: Dimension | None = Field( + default=None, + description=( + "Data quality dimension recorded for this rule's results. Defaults to the check's " + "catalog dimension (validity for custom_sql) when not set explicitly." + ), + ) + + @model_validator(mode="after") + def _validate(self) -> DQRule: + if not self.name: + raise ValueError("Rule name cannot be empty") + catalog_spec = None + if self.check == CUSTOM_SQL_CHECK: + if not self.sql: + raise ValueError(f"Rule {self.name!r}: custom_sql check requires 'sql'") + else: + catalog_spec = CHECK_SPECS.get(self.check) + if catalog_spec is None: + supported = sorted([*CHECK_SPECS, CUSTOM_SQL_CHECK]) + raise ValueError(f"Rule {self.name!r}: unknown check {self.check!r}; supported: {supported}") + if self.sql: + raise ValueError(f"Rule {self.name!r}: 'sql' is only valid with custom_sql") + if catalog_spec.requires_column and not self.column: + raise ValueError(f"Rule {self.name!r}: check {self.check!r} requires 'column'") + if self.condition is None and catalog_spec.default_condition is not None: + object.__setattr__(self, "condition", Condition.from_dict(catalog_spec.default_condition)) + if self.condition is None: + raise ValueError(f"Rule {self.name!r}: condition is required for check {self.check!r}") + if self.dimension is None: + object.__setattr__( + self, "dimension", catalog_spec.dimension if catalog_spec else Dimension.VALIDITY + ) + return self + + @property + def rule_uid(self) -> str: + """Stable identity across runs: survives severity/dimension tweaks and Dag refactors.""" + condition = cast("Condition", self.condition) + identity = { + "name": self.previous_name or self.name, + "check": self.check, + "column": self.column, + "sql": self.sql, + "condition": condition.to_dict(), + } + digest = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() Review Comment: The `previous_name or name` fallback can collide: a rule whose `previous_name` matches another rule's name (with the same check/column/condition) hashes to the same uid, and since RuleSet only rejects duplicate *names*, both pass validation -- then `observed_by_uid` keeps only one UNION ALL row and the two rules overwrite each other's files under `rule_uid=<uid>/`. I'd also question hashing the condition itself. The docstring promises identity that survives severity/dimension tweaks, but tightening a threshold is the edit you actually make while watching a rule's history, and it orphans everything recorded so far. The rename mechanism has a similar wrinkle: continuity only holds while `previous_name` stays in the Dag forever, so deleting it once the rename settles flips the uid back to a fresh history. An explicit optional `id` field would sidestep all of this; short of that, hashing the name-chain/check/column/sql without the condition plus a duplicate-uid check in RuleSet would close the collision. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/version_compat.py: ########## @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# NOTE! THIS FILE IS COPIED MANUALLY IN OTHER PROVIDERS DELIBERATELY TO AVOID ADDING UNNECESSARY +# DEPENDENCIES BETWEEN PROVIDERS. IF YOU WANT TO ADD CONDITIONAL CODE IN YOUR PROVIDER THAT DEPENDS +# ON AIRFLOW VERSION, PLEASE COPY THIS FILE TO THE ROOT PACKAGE OF YOUR PROVIDER AND IMPORT +# THOSE CONSTANTS FROM IT RATHER THAN IMPORTING THEM FROM ANOTHER PROVIDER OR TEST CODE +# +from __future__ import annotations + + +def get_base_airflow_version_tuple() -> tuple[int, int, int]: + from packaging.version import Version + + from airflow import __version__ + + airflow_version = Version(__version__) + return airflow_version.major, airflow_version.minor, airflow_version.micro + + +AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) Review Comment: `AIRFLOW_V_3_1_PLUS` isn't imported anywhere in the provider. Everything used here exists at the 3.0 floor, so it can just be dropped. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/assets.py: ########## @@ -0,0 +1,172 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Asset-level data quality declarations. + +Quality configuration lives inside ``Asset.extra`` under the ``airflow.dataquality`` key, so it is +serialized with the Dag and needs no Airflow core changes. The rules travel with the asset +definition instead of being scattered across the Dags that check it. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.dataquality.exceptions import DQRuleValidationError +from airflow.providers.common.dataquality.rules import RuleSet +from airflow.sdk import task + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from airflow.sdk import Asset + from airflow.sdk.execution_time.comms import AssetEventDagRunReferenceResult + +log = logging.getLogger(__name__) + +DQ_EXTRA_KEY = "airflow.dataquality" +DQ_RESULT_EXTRA_KEY = "airflow.dataquality.result" + + +def asset_quality( + asset: Asset, + *, + ruleset: RuleSet | dict[str, Any] | str, + conn_id: str | None = None, + table: str | None = None, +) -> Asset: + """ + Attach data quality configuration to an asset, returning the same asset. + + :param asset: The asset the rules describe. + :param ruleset: A :class:`~airflow.providers.common.dataquality.rules.RuleSet`, its dict form, or a path + to a YAML ruleset file (resolved eagerly, at Dag-parse time). + :param conn_id: Default connection a check operator should use for this asset. + :param table: Default table to check; falls back to the asset name when unset. + + Usage:: + + orders = asset_quality( + Asset("orders", uri="postgres://warehouse/analytics/orders"), + ruleset=rules, + conn_id="warehouse", + table="analytics.orders", + ) + check = DQCheckOperator(task_id="dq", asset=orders) + """ + if isinstance(ruleset, str): + ruleset = RuleSet.from_file(ruleset) + elif isinstance(ruleset, dict): + ruleset = RuleSet.from_dict(ruleset) + config: dict[str, Any] = {"ruleset": ruleset.to_dict()} + if conn_id: + config["conn_id"] = conn_id + if table: + config["table"] = table + asset.extra[DQ_EXTRA_KEY] = config Review Comment: Design question, non-blocking: core dedupes assets by (name, uri) and `sync_assets` overwrites `model.extra` on every parse, last writer wins -- so a second Dag referencing `Asset("orders")` without quality config wipes the registry copy of the ruleset, and two Dags with different rulesets flip-flop. This slice is unaffected (the operator reads the in-memory object, the gate reads event extra), but the #69413 UI presumably reads the registry copy. How is the attached config expected to survive there? ########## providers/common/dataquality/src/airflow/providers/common/dataquality/backends/__init__.py: ########## @@ -0,0 +1,32 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.providers.common.dataquality.backends.object_storage import ObjectStorageResultsBackend Review Comment: Importing the operator pulls this in eagerly, and `ObjectStoragePath` drags fsspec + upath at import -- paid at Dag-parse time in every file that imports `DQCheckOperator`, even with no `results_path` configured. Since the backend is only built inside `get_backend_from_config()` at execute time, importing `ObjectStorageResultsBackend` lazily there (with a comment) keeps the parse path light. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
