o-nikolas commented on code in PR #69575: URL: https://github.com/apache/airflow/pull/69575#discussion_r4557964029
########## providers/common/dataquality/src/airflow/providers/common/dataquality/assets.py: ########## @@ -0,0 +1,190 @@ +# 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]], + *, + require_all: bool = True, +) -> bool: + """ + Pure decision logic behind :func:`require_quality`, kept separate so it is testable without a Dag. + + ``consumed_asset_events`` carries no ordering guarantee, so the triggering event is never + picked positionally. By default every triggering event must pass, since several producer + runs can coalesce into one consumer run; pass ``require_all=False`` to gate on the single + most recent event (by ``timestamp``) instead. + """ + 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 + + events_to_check = events if require_all else [max(events, key=lambda event: event.timestamp)] + + for event in events_to_check: + summary = event.extra.get(DQ_RESULT_EXTRA_KEY) + score = summary.get("score") if isinstance(summary, dict) else None + if not isinstance(score, (int, float)) or isinstance(score, bool): + log.warning( + "require_quality(%s): triggering event has no data quality summary; skipping downstream tasks", + asset.name, + ) + return False + + if score < min_score: + log.warning( + "require_quality(%s): score %s below required minimum %s; skipping downstream tasks", + asset.name, + score, + min_score, + ) + return False + + return True + + +def require_quality( + asset: Asset, + *, + min_score: float, + task_id: str | None = None, + require_all: bool = True, +) -> Any: + """ + Gate a Dag run on the data quality score attached to its triggering asset event(s). + + Reads the summary a :class:`~airflow.providers.common.dataquality.operators.dq_check.DQCheckOperator` + attaches to ``asset``'s outlet event, under ``asset_event.extra["airflow.dataquality.result"]`` + (see :func:`asset_quality`), and short-circuits the run — skipping every downstream task + — when that summary is missing or its score is below ``min_score``. Call it inside a Dag + scheduled by ``asset``:: + + with DAG("orders_consumer", schedule=orders) as consumer: + start = require_quality(orders, min_score=0.95) + start >> process_orders() + + :param asset: The asset whose triggering event(s) carry the quality summary. + :param min_score: Minimum required score in ``[0, 1]``; the run proceeds only when the + checked event(s)' score is at least this value. + :param task_id: Task id for the generated gate task. Defaults to + ``f"require_quality_{asset.name}"`` so gating on several assets in one Dag doesn't + collide on task id. + :param require_all: When ``True`` (default), every asset event that triggered this run must + pass ``min_score`` — the safe default when several producer runs coalesce into one + consumer run. Set to ``False`` to gate on only the most recent triggering event instead. + """ + if not 0 <= min_score <= 1: + raise ValueError(f"min_score must be between 0 and 1, got {min_score!r}") + gate_task_id = task_id or f"require_quality_{asset.name}" Review Comment: The default `require_quality_{asset.name}` task id won't pass task id validation for URI-style asset names, e.g. `Asset("s3://bucket/orders")`, where `:` and `/` aren't allowed. Maybe replace disallowed characters with `_` when building the default. ########## providers/common/dataquality/src/airflow/providers/common/dataquality/rules/rule.py: ########## @@ -0,0 +1,314 @@ +# 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. A rule whose + ``previous_name`` happens to match another rule's identity (same check/column/condition) + collides on ``rule_uid``; set an explicit ``id`` on one of them to avoid this. + :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. + :param id: Explicit, stable identity for this rule's history. When set, it is used directly + as the ``rule_uid`` instead of the derived hash, so it survives ``previous_name`` chains + and sidesteps any collision between rules that would otherwise hash the same. + + 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 | None = None + column: str | None = None + sql: str | None = None + severity: Severity = Severity.ERROR + partition_clause: str | None = None + previous_name: str | None = None + id: str | None = Field( + default=None, + description=( + "Explicit, stable identity for this rule's history, used directly as rule_uid " + "instead of a derived hash." + ), + ) Review Comment: `id` is used as-is for `rule_uid`, and `rule_uid` gets quoted straight into the check SQL (`engines/sql.py` L55/L86) and into the history path (`backends/object_storage.py` L80/L175). So an id with a `'` makes that rule's SQL invalid, and a `/` splits its history across directories. That's also one more path for generated rules to put text into the query. Could we constrain `id` to something like `^[A-Za-z0-9_.-]+$` in the model and the schema before 0.1.0 ships? Tightening it after release would break existing rulesets. -- 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]
