Copilot commented on code in PR #70225: URL: https://github.com/apache/airflow/pull/70225#discussion_r4330787082
########## providers/standard/src/airflow/providers/standard/sensors/asset.py: ########## @@ -0,0 +1,201 @@ +# 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 typing import TYPE_CHECKING, Any, Literal + +from airflow.providers.common.compat.module_loading import import_string +from airflow.providers.common.compat.sdk import ( + Asset, + AssetAlias, + BaseSensorOperator, + PokeReturnValue, +) +from airflow.providers.standard.version_compat import AIRFLOW_V_3_4_PLUS + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from datetime import datetime + + from airflow.providers.common.compat.sdk import Context + from airflow.sdk.definitions.asset import AssetRef + + +def _serialize_events(events: list[Any]) -> list[Any]: + """Serialize a list of (processed) asset events to JSON-safe values.""" + serialized: list[Any] = [] + for event in events: + model_dump = getattr(event, "model_dump", None) + if callable(model_dump): + serialized.append(model_dump(mode="json")) + else: + serialized.append(event) + return serialized + + +class AssetEventSensor(BaseSensorOperator): + """ + Wait for asset events matching the given filters to reach an expected count. + + The sensor fetches asset events (by asset or asset alias) matching the supplied filters, + optionally applies a ``process_result`` callable to transform, deduplicate or filter them, + and succeeds once the resulting number of events satisfies ``expected_count``. + + This sensor requires Apache Airflow 3.4+ because the ``partition_key``, + ``partition_key_regexp_pattern`` and ``extra`` asset-event filters are only available there. + + :param obj: The :class:`~airflow.sdk.Asset` or :class:`~airflow.sdk.AssetAlias` to wait on. + The target is declared as an inlet. As an alternative, pass ``name``/``uri``/``alias_name`` + directly. These alternatives cannot be combined with ``obj`` or with each other, except + for ``name`` and ``uri`` together. Target identifiers are static; event filters can be templated. + :param name: The asset name to fetch events for. + :param uri: The asset uri to fetch events for. + :param alias_name: The asset alias name to fetch events for. + :param after: Only include events at or after this timestamp. + :param before: Only include events at or before this timestamp. + :param ascending: Whether events are returned in ascending timestamp order. + :param limit: Positive maximum number of events to fetch, before applying ``process_result``. + :param partition_key: Filter by exact partition key match. + :param partition_key_regexp_pattern: Filter by partition key regexp pattern. + :param extra: Filter by key/value pairs contained in the event ``extra`` field. + :param expected_count: The non-negative number of processed events required to succeed. + Defaults to one. If ``process_result`` is omitted, ``limit`` must not be smaller. + :param count_policy: ``"minimum"`` (the default) waits for at least ``expected_count`` events; + ``"exact"`` requires exactly that number. To check for no events, use ``"exact"`` and zero. + :param process_result: A callable (or a dotted import path to one) applied to the fetched + events before the count check, to transform, deduplicate or filter them. It receives the + list of asset events and must return a list. It runs on every poke, so it should be + **idempotent / side-effect free**, and its return value must be JSON-serializable to be + pushed to XCom. + """ + + template_fields: Sequence[str] = ( + "partition_key", + "partition_key_regexp_pattern", + "extra", + "after", + "before", + ) + + def __init__( + self, + *, + obj: Asset | AssetAlias | None = None, + name: str | None = None, + uri: str | None = None, + alias_name: str | None = None, + after: datetime | str | None = None, + before: datetime | str | None = None, + ascending: bool = True, + limit: int | None = None, + partition_key: str | None = None, + partition_key_regexp_pattern: str | None = None, + extra: dict[str, str] | None = None, + expected_count: int = 1, + count_policy: Literal["minimum", "exact"] = "minimum", + process_result: Callable[[list[Any]], list[Any]] | str | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if not AIRFLOW_V_3_4_PLUS: + raise RuntimeError( + "AssetEventSensor requires Apache Airflow 3.4+ because the asset event filters " + "it relies on are only available from 3.4 onwards." + ) + if obj is not None and any(value is not None for value in (name, uri, alias_name)): + raise ValueError("`obj` cannot be combined with `name`, `uri`, or `alias_name`.") + if alias_name is not None and (name is not None or uri is not None): + raise ValueError("`alias_name` cannot be combined with `name` or `uri`.") + self.obj: Asset | AssetAlias | AssetRef + if obj is not None: + if not isinstance(obj, (Asset, AssetAlias)): + raise TypeError(f"`obj` must be an Asset or AssetAlias, got {type(obj).__name__}") + self.obj = obj + elif alias_name is not None: + self.obj = AssetAlias(alias_name) + elif name is not None and uri is not None: + self.obj = Asset(name=name, uri=uri) + elif name is not None: + self.obj = Asset.ref(name=name) + elif uri is not None: + self.obj = Asset.ref(uri=uri) + else: + raise ValueError("One of `obj`, `name`, `uri`, or `alias_name` must be provided.") + if not isinstance(expected_count, int) or isinstance(expected_count, bool): + raise TypeError("`expected_count` must be an integer.") + if expected_count < 0: + raise ValueError(f"`expected_count` must be a non-negative integer, got {expected_count}.") + if limit is not None: + if not isinstance(limit, int) or isinstance(limit, bool): + raise TypeError("`limit` must be an integer.") + if limit <= 0: + raise ValueError("`limit` must be positive.") + if count_policy not in {"minimum", "exact"}: + raise ValueError(f"`count_policy` must be 'minimum' or 'exact', got {count_policy!r}.") + if process_result is None and limit is not None and limit < expected_count: + raise ValueError( + "`limit` must be at least `expected_count` when `process_result` is not provided." + ) + if self.obj not in self.inlets: + self.inlets.append(self.obj) + + self.after = after + self.before = before + self.ascending = ascending + self.limit = limit + self.partition_key = partition_key + self.partition_key_regexp_pattern = partition_key_regexp_pattern + self.extra = extra + self.expected_count = expected_count + self.count_policy = count_policy + self.process_result = process_result + + def _apply_process_result(self, events: list[Any]) -> list[Any]: + if self.process_result is None: + return events + func = self.process_result if callable(self.process_result) else import_string(self.process_result) + return func(events) + + def poke(self, context: Context) -> PokeReturnValue: + accessor = context["inlet_events"][self.obj] Review Comment: A mapped sensor such as `AssetEventSensor.partial(task_id="s", obj=asset).expand(partition_key=["a", "b"])` fails here with `KeyError` unless the target is also explicitly supplied in `inlets`. The runtime constructs `context["inlet_events"]` before unmapping, using `MappedOperator.inlets`; the target is only appended when the sensor's constructor later runs. `context_update_for_unmapped()` updates the task and params, but not the inlet accessor. Make target registration work before context creation, or refresh the accessor from the resolved task's inlets through an appropriate SDK interface. Add a regression test that constructs the context before unmapping, as real execution does. ########## providers/standard/docs/sensors/asset.rst: ########## @@ -0,0 +1,112 @@ + .. 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. + + + +.. _howto/sensor:AssetEventSensor: + +AssetEventSensor +================ + +Use the :class:`~airflow.providers.standard.sensors.asset.AssetEventSensor` to wait for +asset events matching a set of filters to reach an expected count. + +Use this sensor when an already-running Dag needs an asset between tasks. For example, a +daily consumer can run its preparation tasks immediately, then wait for another Dag to +produce the partition for its own data interval. Moving that dependency to an asset schedule +would change when the entire Dag starts and the data interval used by the consumer. + +The sensor declares its target as an inlet, so the dependency is visible in asset lineage +without changing the Dag's schedule. It reads events through ``context["inlet_events"]``. + +.. note:: + + This sensor requires **Apache Airflow 3.4+**, because the ``partition_key``, + ``partition_key_regexp_pattern`` and ``extra`` asset-event filters it relies on are only + available in Airflow 3.4 and later. + +Basic usage +----------- + +Pass an :class:`~airflow.sdk.Asset` or :class:`~airflow.sdk.AssetAlias` as ``obj``, or use the raw +``name`` / ``uri`` / ``alias_name`` arguments. A name-only or URI-only target is an asset reference +and must resolve to an existing asset when the task context is created. Target identifiers are +static so Airflow can record the inlet; use templated event filters to select a run's events. +By default the sensor succeeds once **at least one** matching event exists. + +Choose one target form: ``obj``, ``alias_name``, or ``name`` / ``uri``. The name and URI may be +supplied together, but other combinations are rejected to avoid silently selecting a different asset. + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py + :language: python + :dedent: 4 + :start-after: [START example_asset_event_sensor] + :end-before: [END example_asset_event_sensor] + +Filtering events and expected count +----------------------------------- + +Events can be narrowed with ``after`` / ``before`` (time range), ``ascending`` / ``limit`` +(ordering and cap), ``partition_key`` / ``partition_key_regexp_pattern`` and ``extra`` (key/value +pairs contained in the event ``extra`` field). Prefer a partition key for the event's partition +identity; ``extra`` can further filter metadata, such as a validation status. Review Comment: Even on Airflow 3.4+, the execution API rejects `partition_key_regexp_pattern` with HTTP 400 unless `[api] regexp_query_timeout` is positive. With the default configuration this filter fails the sensor instead of polling. Document how to enable it and clarify that exact-match `partition_key` needs no opt-in. -- 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]
