diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 7eb25dec2b3cb..e1b51ebe0347c 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,6 +42,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_2_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 2) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) if AIRFLOW_V_3_1_PLUS: from airflow.sdk import PokeReturnValue, timezone @@ -64,6 +65,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "AIRFLOW_V_3_3_PLUS", + "AIRFLOW_V_3_4_PLUS", "NOTSET", "XCOM_RETURN_KEY", "ArgNotSet", diff --git a/providers/standard/docs/sensors/asset.rst b/providers/standard/docs/sensors/asset.rst new file mode 100644 index 0000000000000..c2e373cc80759 --- /dev/null +++ b/providers/standard/docs/sensors/asset.rst @@ -0,0 +1,87 @@ + .. 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. + +.. 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 +----------- + +Point the sensor at an :class:`~airflow.sdk.Asset` (or :class:`~airflow.sdk.AssetAlias`, or the raw +``name`` / ``uri`` / ``alias_name``). By default the sensor succeeds once **at least one** matching +event exists. + +.. 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). Use ``expected_count`` to control how many +(processed) events are required to succeed: ``-1`` (the default) means "at least one", ``0`` means +"exactly zero", and any other positive value requires an exact match. + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py + :language: python + :dedent: 4 + :start-after: [START example_asset_event_sensor_filtered] + :end-before: [END example_asset_event_sensor_filtered] + +Post-processing results +----------------------- + +Pass a ``process_result`` callable (or a dotted import path to one) to transform, deduplicate or +filter the fetched events **before** the count check. It receives the list of asset events and must +return a list; the processed events are pushed to XCom. + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py + :language: python + :dedent: 0 + :start-after: [START example_asset_event_process_result] + :end-before: [END example_asset_event_process_result] + +Deferrable mode +--------------- + +Set ``deferrable=True`` to run the sensor in :ref:`deferrable mode `. In this +mode ``process_result`` runs in the triggerer, so it must be a **top-level importable function** +(defined in an installed module, not a lambda, nested function, or bound method) and its return +value must be JSON-serializable. + +.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py + :language: python + :dedent: 4 + :start-after: [START example_asset_event_sensor_async] + :end-before: [END example_asset_event_sensor_async] diff --git a/providers/standard/provider.yaml b/providers/standard/provider.yaml index cb493caa69d09..3f2b64e1a91a4 100644 --- a/providers/standard/provider.yaml +++ b/providers/standard/provider.yaml @@ -81,6 +81,7 @@ integrations: - /docs/apache-airflow-providers-standard/sensors/datetime.rst - /docs/apache-airflow-providers-standard/sensors/file.rst - /docs/apache-airflow-providers-standard/sensors/external_task_sensor.rst + - /docs/apache-airflow-providers-standard/sensors/asset.rst operators: - integration-name: Standard @@ -98,6 +99,7 @@ operators: sensors: - integration-name: Standard python-modules: + - airflow.providers.standard.sensors.asset - airflow.providers.standard.sensors.date_time - airflow.providers.standard.sensors.time_delta - airflow.providers.standard.sensors.time @@ -116,6 +118,7 @@ hooks: triggers: - integration-name: Standard python-modules: + - airflow.providers.standard.triggers.asset - airflow.providers.standard.triggers.external_task - airflow.providers.standard.triggers.file - airflow.providers.standard.triggers.temporal diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_asset_sensor.py b/providers/standard/src/airflow/providers/standard/example_dags/example_asset_sensor.py new file mode 100644 index 0000000000000..175c9266d8f86 --- /dev/null +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_asset_sensor.py @@ -0,0 +1,78 @@ +# 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. +"""Example DAG demonstrating the usage of the ``AssetEventSensor``.""" + +from __future__ import annotations + +from typing import Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.standard.version_compat import AIRFLOW_V_3_4_PLUS + +if not AIRFLOW_V_3_4_PLUS: + raise AirflowOptionalProviderFeatureException("AssetEventSensor needs Airflow 3.4+.") + +import pendulum + +from airflow.providers.standard.sensors.asset import AssetEventSensor +from airflow.sdk import DAG, Asset + +my_asset = Asset("s3://bucket/my-file.csv") + + +# [START example_asset_event_process_result] +def latest_event_only(events: list[Any]) -> list[Any]: + """Keep only the most recent event (must be a top-level importable function in deferrable mode).""" + return events[-1:] + + +# [END example_asset_event_process_result] + + +with DAG( + dag_id="example_asset_sensor", + schedule=None, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, + tags=["example"], +) as dag: + # [START example_asset_event_sensor] + wait_for_asset_event = AssetEventSensor( + task_id="wait_for_asset_event", + asset=my_asset, + ) + # [END example_asset_event_sensor] + + # [START example_asset_event_sensor_filtered] + wait_for_partition_events = AssetEventSensor( + task_id="wait_for_partition_events", + asset=my_asset, + partition_key="2021-01-01", + expected_count=3, + ) + # [END example_asset_event_sensor_filtered] + + # [START example_asset_event_sensor_async] + wait_for_asset_event_async = AssetEventSensor( + task_id="wait_for_asset_event_async", + asset=my_asset, + process_result=latest_event_only, + deferrable=True, + ) + # [END example_asset_event_sensor_async] + + [wait_for_asset_event, wait_for_partition_events, wait_for_asset_event_async] diff --git a/providers/standard/src/airflow/providers/standard/exceptions.py b/providers/standard/src/airflow/providers/standard/exceptions.py index df13fd5040b2d..b3550f1a2e624 100644 --- a/providers/standard/src/airflow/providers/standard/exceptions.py +++ b/providers/standard/src/airflow/providers/standard/exceptions.py @@ -57,6 +57,14 @@ class DuplicateStateError(AirflowExternalTaskSensorException): """Raised when duplicate states are provided across allowed, skipped and failed states.""" +class AssetEventSensorException(AirflowException): + """Base exception for all AssetEventSensor related errors.""" + + +class UnexpectedAssetEventTriggerEventError(AssetEventSensorException): + """Raised when the AssetEventTrigger returns an unexpected event to the sensor.""" + + class HITLTriggerEventError(Exception): """Raised when TriggerEvent contains error.""" diff --git a/providers/standard/src/airflow/providers/standard/get_provider_info.py b/providers/standard/src/airflow/providers/standard/get_provider_info.py index 1f7b2049454d1..13155d14f86e6 100644 --- a/providers/standard/src/airflow/providers/standard/get_provider_info.py +++ b/providers/standard/src/airflow/providers/standard/get_provider_info.py @@ -43,6 +43,7 @@ def get_provider_info(): "/docs/apache-airflow-providers-standard/sensors/datetime.rst", "/docs/apache-airflow-providers-standard/sensors/file.rst", "/docs/apache-airflow-providers-standard/sensors/external_task_sensor.rst", + "/docs/apache-airflow-providers-standard/sensors/asset.rst", ], } ], @@ -67,6 +68,7 @@ def get_provider_info(): { "integration-name": "Standard", "python-modules": [ + "airflow.providers.standard.sensors.asset", "airflow.providers.standard.sensors.date_time", "airflow.providers.standard.sensors.time_delta", "airflow.providers.standard.sensors.time", @@ -92,6 +94,7 @@ def get_provider_info(): { "integration-name": "Standard", "python-modules": [ + "airflow.providers.standard.triggers.asset", "airflow.providers.standard.triggers.external_task", "airflow.providers.standard.triggers.file", "airflow.providers.standard.triggers.temporal", diff --git a/providers/standard/src/airflow/providers/standard/sensors/asset.py b/providers/standard/src/airflow/providers/standard/sensors/asset.py new file mode 100644 index 0000000000000..393af347f7c4d --- /dev/null +++ b/providers/standard/src/airflow/providers/standard/sensors/asset.py @@ -0,0 +1,244 @@ +# 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 + +from airflow.providers.common.compat.module_loading import import_string +from airflow.providers.common.compat.sdk import ( + Asset, + AssetAlias, + BaseSensorOperator, + PokeReturnValue, + conf, +) +from airflow.providers.standard.exceptions import UnexpectedAssetEventTriggerEventError +from airflow.providers.standard.triggers.asset import ( + AssetEventTrigger, + _count_satisfied, + _fetch_asset_events, + _serialize_events, +) +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 + + +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 asset: The :class:`~airflow.sdk.Asset` or :class:`~airflow.sdk.AssetAlias` to wait on. + As an alternative, pass ``name``/``uri``/``alias_name`` directly. + :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: Maximum number of events to fetch. + :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 number of events required to succeed. ``-1`` (the default) means + "at least one" (``count >= 1``); ``0`` means "exactly zero"; any other positive value + requires an exact match. Note that if ``limit`` is set below an exact ``expected_count`` + the condition can never be satisfied (the sensor will wait until it times out). + :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. In ``deferrable`` mode it must be a top-level + importable function (not a lambda, nested function, or bound method) that lives in a module + installed on the triggerer -- a function defined in a DAG file will import on the worker but + typically fail to import in the triggerer process. Its return value must be JSON-serializable, + because it runs in the triggerer and its result is returned via the trigger event. It should + be **idempotent / side-effect free**: in deferrable mode it runs once during the initial + synchronous poke and again on every fetch in the triggerer. + :param deferrable: Run the sensor in deferrable mode. + """ + + template_fields: Sequence[str] = ( + "name", + "uri", + "alias_name", + "partition_key", + "partition_key_regexp_pattern", + "extra", + "after", + "before", + ) + + def __init__( + self, + *, + asset: 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, + process_result: Callable[[list[Any]], list[Any]] | str | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **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 asset is not None: + if isinstance(asset, AssetAlias): + alias_name = asset.name + elif isinstance(asset, Asset): + name = asset.name + uri = asset.uri + else: + raise TypeError(f"`asset` must be an Asset or AssetAlias, got {type(asset).__name__}") + if not (name or uri or alias_name): + raise ValueError("One of `asset`, `name`, `uri`, or `alias_name` must be provided.") + if expected_count < -1: + raise ValueError( + f"`expected_count` must be -1 (at least one) or a non-negative integer, got {expected_count}." + ) + + self.name = name + self.uri = uri + self.alias_name = alias_name + 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.process_result = process_result + self.deferrable = deferrable + + 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 _resolve_process_result_path(self) -> str | None: + """Resolve ``process_result`` to a dotted import path for use in the trigger.""" + if self.process_result is None: + return None + if isinstance(self.process_result, str): + path = self.process_result + else: + func = self.process_result + module = getattr(func, "__module__", "") + qualname = getattr(func, "__qualname__", "") + if not module or not qualname or "<" in qualname: + raise ValueError( + "In deferrable mode, `process_result` must be a top-level importable function " + "(or a dotted import path string), not a lambda or a nested/local function." + ) + path = f"{module}.{qualname}" + # Fail fast at defer time rather than in the triggerer: e.g. a bound method resolves to + # ``module.Class.method`` which ``import_string`` cannot import. This validates importability + # on the *worker*; the callable must also live in a module installed on the triggerer (a + # function defined in a DAG file may import here yet fail there). + try: + import_string(path) + except Exception as exc: + raise ValueError( + "In deferrable mode, `process_result` must be importable by a dotted path, but the " + f"resolved path {path!r} is not importable ({exc}). Pass a top-level function " + "(defined in an installed module, not a DAG file) or an explicit import-path string." + ) from exc + return path + + def poke(self, context: Context) -> PokeReturnValue: + events = _fetch_asset_events( + name=self.name, + uri=self.uri, + alias_name=self.alias_name, + after=self.after, + before=self.before, + ascending=self.ascending, + limit=self.limit, + partition_key=self.partition_key, + partition_key_regexp_pattern=self.partition_key_regexp_pattern, + extra=self.extra, + ) + processed = self._apply_process_result(events) + count = len(processed) + done = _count_satisfied(count, self.expected_count) + self.log.info( + "Found %d matching asset events (expected %s): %s", + count, + self.expected_count, + "condition met" if done else "still waiting", + ) + return PokeReturnValue(is_done=done, xcom_value=_serialize_events(processed) if done else None) + + def execute(self, context: Context) -> Any: + if not self.deferrable: + return super().execute(context=context) + + # Deferrable path: do an initial synchronous check to avoid deferring needlessly. + poke_result = self.poke(context) + if poke_result: + return poke_result.xcom_value if isinstance(poke_result, PokeReturnValue) else None + + self.defer( + trigger=AssetEventTrigger( + name=self.name, + uri=self.uri, + alias_name=self.alias_name, + after=self.after, + before=self.before, + ascending=self.ascending, + limit=self.limit, + partition_key=self.partition_key, + partition_key_regexp_pattern=self.partition_key_regexp_pattern, + extra=self.extra, + expected_count=self.expected_count, + process_result_path=self._resolve_process_result_path(), + poke_interval=self.poke_interval, + ), + method_name="execute_complete", + timeout=self.timeout, + ) + + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> Any: + """Return the processed events once the trigger fires successfully.""" + if event and event.get("status") == "success": + return event.get("events") + raise UnexpectedAssetEventTriggerEventError( + f"AssetEventSensor received an unexpected trigger event: {event}" + ) diff --git a/providers/standard/src/airflow/providers/standard/triggers/asset.py b/providers/standard/src/airflow/providers/standard/triggers/asset.py new file mode 100644 index 0000000000000..d8899e56293e4 --- /dev/null +++ b/providers/standard/src/airflow/providers/standard/triggers/asset.py @@ -0,0 +1,206 @@ +# 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 + +import asyncio +from typing import TYPE_CHECKING, Any + +from asgiref.sync import sync_to_async + +from airflow.providers.common.compat.module_loading import import_string +from airflow.triggers.base import BaseTrigger, TriggerEvent + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + from datetime import datetime + + +def _count_satisfied(count: int, expected_count: int) -> bool: + """ + Return whether the number of (processed) asset events satisfies the expectation. + + ``expected_count == -1`` means "at least one" (``count >= 1``); any other value + requires an exact match (``count == expected_count``). + """ + if expected_count == -1: + return count >= 1 + return count == expected_count + + +def _fetch_asset_events( + *, + name: str | None, + uri: str | None, + alias_name: str | None, + after: datetime | str | None, + before: datetime | str | None, + ascending: bool, + limit: int | None, + partition_key: str | None, + partition_key_regexp_pattern: str | None, + extra: dict[str, str] | None, +) -> list[Any]: + """ + Fetch asset events matching the given filters. + + This uses the Task SDK :class:`InletEventsAccessor`, which lazily fetches events + from the supervisor. It is a blocking call and must be wrapped with + ``sync_to_async`` when used from within the triggerer. + """ + from airflow.sdk.execution_time.context import InletEventsAccessor + + accessor = InletEventsAccessor(asset_name=name, asset_uri=uri, alias_name=alias_name) + if after is not None: + accessor.after(after if isinstance(after, str) else after.isoformat()) + if before is not None: + accessor.before(before if isinstance(before, str) else before.isoformat()) + accessor.ascending(ascending) + if limit is not None: + accessor.limit(limit) + if partition_key is not None: + accessor.partition_key(partition_key) + if partition_key_regexp_pattern is not None: + accessor.partition_key_regexp_pattern(partition_key_regexp_pattern) + if extra: + for key, value in extra.items(): + accessor.extra(key, value) + return list(accessor) + + +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 AssetEventTrigger(BaseTrigger): + """ + A trigger that waits until asset events matching the given filters reach an expected count. + + The trigger periodically fetches asset events (by asset name/uri or alias) matching the + supplied filters, optionally applies a ``process_result`` callable to transform, deduplicate + or filter them, and fires once the resulting count satisfies ``expected_count``. + + :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: Maximum number of events to fetch. + :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 number of events required to succeed. ``-1`` means at least one. + :param process_result_path: Dotted import path to a callable applied to the fetched events + before the count check. Its return value must be JSON-serializable. + :param poke_interval: Number of seconds to wait between fetches. + """ + + def __init__( + self, + *, + 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, + process_result_path: str | None = None, + poke_interval: float = 60, + ) -> None: + super().__init__() + if expected_count < -1: + raise ValueError( + f"`expected_count` must be -1 (at least one) or a non-negative integer, got {expected_count}." + ) + self.name = name + self.uri = uri + self.alias_name = alias_name + 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.process_result_path = process_result_path + self.poke_interval = poke_interval + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + "airflow.providers.standard.triggers.asset.AssetEventTrigger", + { + "name": self.name, + "uri": self.uri, + "alias_name": self.alias_name, + "after": self.after, + "before": self.before, + "ascending": self.ascending, + "limit": self.limit, + "partition_key": self.partition_key, + "partition_key_regexp_pattern": self.partition_key_regexp_pattern, + "extra": self.extra, + "expected_count": self.expected_count, + "process_result_path": self.process_result_path, + "poke_interval": self.poke_interval, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + process_result = import_string(self.process_result_path) if self.process_result_path else None + while True: + events = await sync_to_async(_fetch_asset_events)( + name=self.name, + uri=self.uri, + alias_name=self.alias_name, + after=self.after, + before=self.before, + ascending=self.ascending, + limit=self.limit, + partition_key=self.partition_key, + partition_key_regexp_pattern=self.partition_key_regexp_pattern, + extra=self.extra, + ) + if process_result is not None: + # ``process_result`` may be an ordinary (non-async) callable, so run it in a + # separate thread to avoid blocking the triggerer event loop. + events = await sync_to_async(process_result, thread_sensitive=False)(events) + count = len(events) + if _count_satisfied(count, self.expected_count): + self.log.info("Found %d matching asset events; firing trigger.", count) + yield TriggerEvent({"status": "success", "events": _serialize_events(events)}) + return + self.log.info( + "Found %d matching asset events, expected %s; sleeping %s seconds.", + count, + self.expected_count, + self.poke_interval, + ) + await asyncio.sleep(self.poke_interval) diff --git a/providers/standard/src/airflow/providers/standard/version_compat.py b/providers/standard/src/airflow/providers/standard/version_compat.py index f99a558937807..146a5828fa4bc 100644 --- a/providers/standard/src/airflow/providers/standard/version_compat.py +++ b/providers/standard/src/airflow/providers/standard/version_compat.py @@ -37,6 +37,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 3) AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0) # BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator 3.0.x # This is needed for DecoratedOperator compatibility @@ -60,6 +61,7 @@ def is_arg_set(value): # type: ignore[misc,no-redef] "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "AIRFLOW_V_3_3_PLUS", + "AIRFLOW_V_3_4_PLUS", "ArgNotSet", "BaseOperator", "is_arg_set", diff --git a/providers/standard/tests/unit/standard/conftest.py b/providers/standard/tests/unit/standard/conftest.py new file mode 100644 index 0000000000000..dbc1d2d453eb7 --- /dev/null +++ b/providers/standard/tests/unit/standard/conftest.py @@ -0,0 +1,215 @@ +# 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. +""" +Shared fixtures for the database-backed asset sensor/trigger tests. + +The :class:`AssetEventSensor` / :class:`AssetEventTrigger` always fetch through the Task SDK +:class:`~airflow.sdk.execution_time.context.InletEventsAccessor`, which talks to +``SUPERVISOR_COMMS`` (normally a socket to the supervisor, which calls the execution API, which +queries the DB). There is no supervisor in a unit test, so :class:`_DBBackedComms` replaces only +that transport seam: it runs the *real* execution-API SQL against the test metadata DB, so the +partition-key / extra / time-range / limit filters execute for real. ``process_result`` and the +count check then run unmocked on the returned events. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +import pytest + +ASSET_NAME = "my_asset" +ASSET_URI = "s3://bucket/key" +ALIAS_NAME = "my_alias" + + +def event_timestamp(day: int) -> datetime: + return datetime(2024, 1, day, tzinfo=timezone.utc) + + +class _DBBackedComms: + """A ``SUPERVISOR_COMMS`` stand-in that answers asset-event queries from the metadata DB. + + It mirrors what :class:`~airflow.sdk.execution_time.supervisor.ActivitySubprocess` does for + ``GetAssetEventByAsset`` / ``GetAssetEventByAssetAlias``, but instead of an HTTP round-trip it + calls the execution-API route function directly against a fresh session (so it is safe to use + from the triggerer's worker thread as well). + """ + + def send(self, msg: Any) -> Any: + from airflow.api_fastapi.common.parameters import FilterParam, _RegexParam + from airflow.api_fastapi.execution_api.routes.asset_events import ( + get_asset_event_by_asset_alias, + get_asset_event_by_asset_name_uri, + ) + from airflow.models.asset import AssetEvent + from airflow.sdk.execution_time.comms import ( + AssetEventsResult, + GetAssetEventByAsset, + GetAssetEventByAssetAlias, + ) + from airflow.utils.session import create_session + + def _extra_list(extra: dict[str, str] | None) -> list[str] | None: + return [f"{k}={v}" for k, v in extra.items()] if extra else None + + with create_session() as session: + partition_key = FilterParam(AssetEvent.partition_key, msg.partition_key) + partition_key_regexp = _RegexParam(AssetEvent.partition_key, msg.partition_key_regexp_pattern) + if isinstance(msg, GetAssetEventByAsset): + response = get_asset_event_by_asset_name_uri( + name=msg.name, + uri=msg.uri, + session=session, + partition_key=partition_key, + partition_key_regexp_pattern=partition_key_regexp, + after=msg.after, + before=msg.before, + ascending=msg.ascending, + limit=msg.limit, + extra=_extra_list(msg.extra), + ) + elif isinstance(msg, GetAssetEventByAssetAlias): + response = get_asset_event_by_asset_alias( + name=msg.alias_name, + session=session, + partition_key=partition_key, + partition_key_regexp_pattern=partition_key_regexp, + after=msg.after, + before=msg.before, + ascending=msg.ascending, + limit=msg.limit, + extra=_extra_list(msg.extra), + ) + else: # pragma: no cover - defensive + raise AssertionError(f"unexpected supervisor message: {msg!r}") + return AssetEventsResult.from_asset_events_response(response) + + +@pytest.fixture +def db_supervisor_comms(monkeypatch): + """Route the Task SDK ``SUPERVISOR_COMMS`` to the metadata DB for the duration of a test.""" + from airflow.sdk.execution_time import task_runner + + comms = _DBBackedComms() + monkeypatch.setattr(task_runner, "SUPERVISOR_COMMS", comms, raising=False) + return comms + + +@pytest.fixture +def asset_event_rows(session): + """Create one active asset with five events (varied partition keys / extras / timestamps). + + Layout (ascending by timestamp): + id=1 us|2024-01-01 {"region": "us"} + id=2 us|2024-01-02 {"region": "us"} + id=3 eu|2024-01-01 {"region": "eu"} + id=4 us|2024-01-01 {"region": "us"} # duplicate partition key of id=1 (exercises dedup) + id=5 None {} + Also links every event to an alias named ``my_alias``. + """ + from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel + + from tests_common.test_utils.db import clear_db_assets + + clear_db_assets() + + asset = AssetModel(id=1, name=ASSET_NAME, uri=ASSET_URI, group="asset", extra={}) + session.add_all([asset, AssetActive.for_asset(asset)]) + session.commit() + + common = { + "asset_id": 1, + "source_dag_id": "producer", + "source_task_id": "emit", + "source_run_id": "run_1", + "source_map_index": -1, + } + events = [ + AssetEvent( + id=1, + timestamp=event_timestamp(1), + partition_key="us|2024-01-01", + extra={"region": "us"}, + **common, + ), + AssetEvent( + id=2, + timestamp=event_timestamp(2), + partition_key="us|2024-01-02", + extra={"region": "us"}, + **common, + ), + AssetEvent( + id=3, + timestamp=event_timestamp(3), + partition_key="eu|2024-01-01", + extra={"region": "eu"}, + **common, + ), + AssetEvent( + id=4, + timestamp=event_timestamp(4), + partition_key="us|2024-01-01", + extra={"region": "us"}, + **common, + ), + AssetEvent(id=5, timestamp=event_timestamp(5), partition_key=None, extra={}, **common), + ] + session.add_all(events) + session.commit() + + alias = AssetAliasModel(id=1, name=ALIAS_NAME) + alias.asset_events = events + alias.assets.append(asset) + session.add(alias) + session.commit() + + yield events + + clear_db_assets() + + +@pytest.fixture +def add_asset_event(session): + """Return a helper that inserts (and commits) an extra asset event for ``asset_id=1``. + + Useful for exercising the trigger's polling loop: insert a new matching event "between" two + fetches (e.g. from a patched ``asyncio.sleep``) and watch the trigger fire. + """ + from airflow.models.asset import AssetEvent + from airflow.utils.session import create_session + + def _add(*, id: int, day: int, partition_key: str | None, extra: dict[str, str] | None = None) -> None: + with create_session() as new_session: + new_session.add( + AssetEvent( + id=id, + timestamp=event_timestamp(day), + partition_key=partition_key, + extra=extra or {}, + asset_id=1, + source_dag_id="producer", + source_task_id="emit", + source_run_id="run_1", + source_map_index=-1, + ) + ) + new_session.commit() + + return _add diff --git a/providers/standard/tests/unit/standard/sensors/test_asset.py b/providers/standard/tests/unit/standard/sensors/test_asset.py new file mode 100644 index 0000000000000..fe342e567214f --- /dev/null +++ b/providers/standard/tests/unit/standard/sensors/test_asset.py @@ -0,0 +1,435 @@ +# 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. +""" +Tests for :class:`AssetEventSensor`. + +Most tests are database-backed (``db_test``): they create real ``AssetEvent`` rows and exercise the +whole fetch path (real ``InletEventsAccessor`` -> DB-backed ``SUPERVISOR_COMMS`` -> real +execution-API SQL), rather than patching ``_fetch_asset_events``. The DB harness lives in the +``conftest.py`` one directory up (``db_supervisor_comms`` / ``asset_event_rows`` fixtures). +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +import pytest + +from airflow.providers.common.compat.sdk import Asset, AssetAlias, PokeReturnValue, TaskDeferred +from airflow.providers.standard.exceptions import UnexpectedAssetEventTriggerEventError +from airflow.providers.standard.sensors.asset import AssetEventSensor +from airflow.providers.standard.triggers.asset import AssetEventTrigger + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS + +pytestmark = pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="AssetEventSensor requires Airflow 3.4+") + +FETCH_PATH = "airflow.providers.standard.sensors.asset._fetch_asset_events" + +ASSET_NAME = "my_asset" +ASSET_URI = "s3://bucket/key" +ALIAS_NAME = "my_alias" + + +def _ts(day: int) -> datetime: + return datetime(2024, 1, day, tzinfo=timezone.utc) + + +# top-level (importable) process_result callables, so they also work in deferrable mode +def only_us_partitions(events: list[Any]) -> list[Any]: + """Keep only events whose partition key targets the ``us`` region.""" + return [event for event in events if (event.partition_key or "").startswith("us|")] + + +def dedup_by_partition_key(events: list[Any]) -> list[Any]: + """Keep the first event seen for each distinct partition key (order preserved).""" + seen: set[str | None] = set() + out: list[Any] = [] + for event in events: + if event.partition_key in seen: + continue + seen.add(event.partition_key) + out.append(event) + return out + + +class _Processor: + """A callable holder whose bound method cannot be serialized to an import path.""" + + def run(self, events: list[Any]) -> list[Any]: + return events + + +def _partition_keys(result: Any) -> list[str | None]: + """Extract partition keys from a poke result's (serialized) xcom value.""" + assert isinstance(result, PokeReturnValue) + return [event["partition_key"] for event in (result.xcom_value or [])] + + +class TestInit: + """Pure construction/validation logic (no DB access).""" + + def test_requires_a_target(self): + with pytest.raises(ValueError, match="One of `asset`, `name`, `uri`, or `alias_name`"): + AssetEventSensor(task_id="s") + + def test_target_from_asset(self): + sensor = AssetEventSensor(task_id="s", asset=Asset(name="my_asset", uri="s3://bucket/key")) + assert sensor.name == "my_asset" + assert sensor.uri == "s3://bucket/key" + assert sensor.alias_name is None + + def test_target_from_asset_alias(self): + sensor = AssetEventSensor(task_id="s", asset=AssetAlias(name="my_alias")) + assert sensor.alias_name == "my_alias" + assert sensor.name is None + assert sensor.uri is None + + def test_target_rejects_bad_type(self): + with pytest.raises(TypeError, match="must be an Asset or AssetAlias"): + AssetEventSensor(task_id="s", asset=object()) # type: ignore[arg-type] + + def test_rejects_invalid_expected_count(self): + with pytest.raises(ValueError, match="expected_count"): + AssetEventSensor(task_id="s", name=ASSET_NAME, expected_count=-2) + + def test_requires_airflow_3_4(self, mocker): + # The version guard is dead code in CI (the module is skipped below 3.4), so patch the flag. + mocker.patch("airflow.providers.standard.sensors.asset.AIRFLOW_V_3_4_PLUS", False) + with pytest.raises(RuntimeError, match="requires Apache Airflow 3.4"): + AssetEventSensor(task_id="s", name=ASSET_NAME) + + def test_template_fields(self): + assert set(AssetEventSensor.template_fields) >= { + "name", + "uri", + "alias_name", + "partition_key", + "partition_key_regexp_pattern", + "extra", + "after", + "before", + } + + def test_deferrable_default_from_config(self): + # ``deferrable``'s default is bound to ``[operators] default_deferrable`` at import time, so + # the module has to be reloaded under the patched config to observe the new default. + import importlib + + from airflow.providers.standard.sensors import asset as asset_module + + try: + with conf_vars({("operators", "default_deferrable"): "True"}): + importlib.reload(asset_module) + sensor = asset_module.AssetEventSensor(task_id="s", name=ASSET_NAME) + assert sensor.deferrable is True + finally: + importlib.reload(asset_module) # restore the default-config version for other tests + + def test_non_deferrable_execute_delegates_to_base(self, mocker): + base_execute = mocker.patch( + "airflow.providers.standard.sensors.asset.BaseSensorOperator.execute", + return_value="delegated", + ) + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, deferrable=False) + assert sensor.execute({"k": "v"}) == "delegated" + base_execute.assert_called_once() + + def test_execute_complete_success(self): + sensor = AssetEventSensor(task_id="s", name="my_asset") + events = [{"id": 1, "partition_key": "a"}] + assert sensor.execute_complete({}, {"status": "success", "events": events}) == events + + def test_execute_complete_failure(self): + sensor = AssetEventSensor(task_id="s", name="my_asset") + with pytest.raises(UnexpectedAssetEventTriggerEventError, match="unexpected trigger event"): + sensor.execute_complete({}, {"status": "boom"}) + + def test_execute_complete_none_event(self): + sensor = AssetEventSensor(task_id="s", name="my_asset") + with pytest.raises(UnexpectedAssetEventTriggerEventError, match="unexpected trigger event"): + sensor.execute_complete({}) + + def test_execute_complete_missing_status(self): + sensor = AssetEventSensor(task_id="s", name="my_asset") + with pytest.raises(UnexpectedAssetEventTriggerEventError, match="unexpected trigger event"): + sensor.execute_complete({}, {"events": []}) + + +@pytest.mark.db_test +class TestPoke: + """poke() against real ``AssetEvent`` rows (no patching of the fetch).""" + + def test_returns_all_events(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME) # expected_count=-1 (at least one) + result = sensor.poke({}) + assert bool(result) is True + assert _partition_keys(result) == [ + "us|2024-01-01", + "us|2024-01-02", + "eu|2024-01-01", + "us|2024-01-01", + None, + ] + + def test_no_match_is_not_done(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", name="does_not_exist") + result = sensor.poke({}) + assert bool(result) is False + assert result.xcom_value is None + + def test_exact_partition_key_filter(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, partition_key="us|2024-01-01") + assert _partition_keys(sensor.poke({})) == ["us|2024-01-01", "us|2024-01-01"] + + def test_extra_filter(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, extra={"region": "eu"}) + assert _partition_keys(sensor.poke({})) == ["eu|2024-01-01"] + + def test_time_range_limit_and_descending(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, after=_ts(2), before=_ts(4), ascending=False, limit=2 + ) + # events in [day2, day4] are ids 2,3,4; descending -> 4,3; limited to 2 + assert _partition_keys(sensor.poke({})) == ["us|2024-01-01", "eu|2024-01-01"] + + @pytest.mark.parametrize(("expected_count", "done"), [(5, True), (4, False), (1, False)]) + def test_exact_expected_count(self, asset_event_rows, db_supervisor_comms, expected_count, done): + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, expected_count=expected_count) + assert bool(sensor.poke({})) is done + + def test_by_uri(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", uri=ASSET_URI) + assert len(sensor.poke({}).xcom_value) == 5 + + def test_by_alias(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", asset=AssetAlias(name=ALIAS_NAME)) + assert len(sensor.poke({}).xcom_value) == 5 + + def test_expected_count_zero_matches_none(self, asset_event_rows, db_supervisor_comms): + # expected_count=0 means "exactly zero events"; a filter matching nothing satisfies it. + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, partition_key="does-not-exist", expected_count=0 + ) + result = sensor.poke({}) + assert bool(result) is True + assert result.xcom_value == [] + + def test_limit_below_exact_expected_count_never_satisfied(self, asset_event_rows, db_supervisor_comms): + # limit caps the DB result below the exact expected_count, so the condition can never be met. + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, limit=1, expected_count=3) + assert bool(sensor.poke({})) is False + + def test_combined_filters(self, asset_event_rows, db_supervisor_comms): + # partition_key AND extra AND time-range together, resolved by real SQL -> only id=1. + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + partition_key="us|2024-01-01", + extra={"region": "us"}, + after=_ts(1), + before=_ts(1), + ) + assert _partition_keys(sensor.poke({})) == ["us|2024-01-01"] + + def test_string_datetime_bounds(self, asset_event_rows, db_supervisor_comms): + # after/before accept ISO strings (coerced to datetime by the comms message model). + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + after="2024-01-02T00:00:00+00:00", + before="2024-01-03T00:00:00+00:00", + ) + assert _partition_keys(sensor.poke({})) == ["us|2024-01-02", "eu|2024-01-01"] + + +@pytest.mark.db_test +class TestProcessResult: + """process_result (filter / dedup) applied to real DB events before the count check.""" + + def test_filters_out_events(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, process_result=only_us_partitions, expected_count=3 + ) + result = sensor.poke({}) + assert bool(result) is True + assert _partition_keys(result) == ["us|2024-01-01", "us|2024-01-02", "us|2024-01-01"] + + def test_dedup(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, process_result=dedup_by_partition_key, expected_count=4 + ) + result = sensor.poke({}) + assert bool(result) is True + assert _partition_keys(result) == ["us|2024-01-01", "us|2024-01-02", "eu|2024-01-01", None] + + def test_by_import_path(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + process_result=f"{__name__}.only_us_partitions", + expected_count=3, + ) + assert bool(sensor.poke({})) is True + + def test_reduces_below_expected(self, asset_event_rows, db_supervisor_comms): + # only 3 "us" events remain after filtering, so an exact expectation of 5 is not met + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, process_result=only_us_partitions, expected_count=5 + ) + assert bool(sensor.poke({})) is False + + def test_returns_empty_list_not_done(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, process_result=lambda events: []) + result = sensor.poke({}) # expected_count=-1 -> "at least one" -> not satisfied by [] + assert bool(result) is False + assert result.xcom_value is None + + def test_returns_empty_list_with_expected_zero(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, process_result=lambda events: [], expected_count=0 + ) + result = sensor.poke({}) + assert bool(result) is True + assert result.xcom_value == [] + + def test_returns_more_events_than_fetched(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, process_result=lambda events: events + events, expected_count=10 + ) + result = sensor.poke({}) + assert bool(result) is True + assert len(result.xcom_value) == 10 + + def test_exception_propagates(self, asset_event_rows, db_supervisor_comms): + def boom(events): + raise RuntimeError("process_result exploded") + + sensor = AssetEventSensor(task_id="s", name=ASSET_NAME, process_result=boom) + with pytest.raises(RuntimeError, match="process_result exploded"): + sensor.poke({}) + + +@pytest.mark.db_test +class TestDeferrable: + def test_short_circuits_when_already_satisfied(self, asset_event_rows, db_supervisor_comms): + # Initial poke already satisfies expected_count, so execute() returns without deferring. + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + process_result=only_us_partitions, + expected_count=3, + deferrable=True, + ) + events = sensor.execute({}) + assert [event["partition_key"] for event in events] == [ + "us|2024-01-01", + "us|2024-01-02", + "us|2024-01-01", + ] + + def test_defers_and_forwards_process_result_path(self, asset_event_rows, db_supervisor_comms): + # Ask for more events than exist so the initial poke fails and the sensor defers. + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + partition_key="us|2024-01-01", + process_result=only_us_partitions, + expected_count=99, + deferrable=True, + ) + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + trigger = exc.value.trigger + assert isinstance(trigger, AssetEventTrigger) + assert trigger.name == ASSET_NAME + assert trigger.partition_key == "us|2024-01-01" + assert trigger.expected_count == 99 + assert trigger.process_result_path.endswith("test_asset.only_us_partitions") + + def test_rejects_non_importable_callable(self, asset_event_rows, db_supervisor_comms): + # A lambda cannot be serialized to an import path for the trigger, so deferring must fail. + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + expected_count=99, # force a defer past the initial poke + process_result=lambda events: events, + deferrable=True, + ) + with pytest.raises(ValueError, match="top-level importable function"): + sensor.execute({}) + + def test_rejects_bound_method_callable(self, asset_event_rows, db_supervisor_comms): + # A bound method resolves to ``module.Class.method`` which import_string cannot import; + # this must fail at defer time (not later in the triggerer). + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + expected_count=99, + process_result=_Processor().run, + deferrable=True, + ) + with pytest.raises(ValueError, match="not importable"): + sensor.execute({}) + + def test_forwards_poke_interval_to_trigger(self, asset_event_rows, db_supervisor_comms): + sensor = AssetEventSensor( + task_id="s", name=ASSET_NAME, expected_count=99, deferrable=True, poke_interval=42 + ) + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + assert exc.value.trigger.poke_interval == 42 + + +class TestDeferrableFilterForwarding: + """The regexp filter cannot run on SQLite, so verify filter -> trigger forwarding with a mock. + + This is the one place we still patch the fetch: it lets us assert that *all* filters (including + ``partition_key_regexp_pattern``) are forwarded to the trigger without executing a regexp query + that the local SQLite backend does not support. + """ + + def test_forwards_all_filters_to_trigger(self, mocker): + mocker.patch(FETCH_PATH, return_value=[]) + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + uri=ASSET_URI, + after=_ts(1), + before=_ts(9), + ascending=False, + limit=7, + partition_key="us|2024-01-01", + partition_key_regexp_pattern="us.*", + extra={"region": "us"}, + expected_count=3, + deferrable=True, + ) + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + trigger = exc.value.trigger + assert isinstance(trigger, AssetEventTrigger) + assert (trigger.name, trigger.uri) == (ASSET_NAME, ASSET_URI) + assert (trigger.after, trigger.before) == (_ts(1), _ts(9)) + assert trigger.ascending is False + assert trigger.limit == 7 + assert trigger.partition_key == "us|2024-01-01" + assert trigger.partition_key_regexp_pattern == "us.*" + assert trigger.extra == {"region": "us"} + assert trigger.expected_count == 3 diff --git a/providers/standard/tests/unit/standard/triggers/test_asset.py b/providers/standard/tests/unit/standard/triggers/test_asset.py new file mode 100644 index 0000000000000..fe8b19ca59bbb --- /dev/null +++ b/providers/standard/tests/unit/standard/triggers/test_asset.py @@ -0,0 +1,316 @@ +# 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. +""" +Tests for :class:`AssetEventTrigger`. + +The ``TestRun`` cases are database-backed (``db_test``): the trigger fetches from real +``AssetEvent`` rows via a DB-backed ``SUPERVISOR_COMMS`` (see the ``conftest.py`` one directory up), +imports the ``process_result`` callable by path, applies it, and fires once the count is satisfied. +Only the pure helpers and the serialization contract are tested without a DB. +""" + +from __future__ import annotations + +from typing import Any +from unittest import mock + +import pytest + +from airflow.providers.common.compat.sdk import TaskDeferred +from airflow.providers.standard.sensors.asset import AssetEventSensor +from airflow.providers.standard.triggers.asset import ( + AssetEventTrigger, + _count_satisfied, + _serialize_events, +) +from airflow.triggers.base import TriggerEvent + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS + +pytestmark = pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="AssetEventTrigger requires Airflow 3.4+") + +ASSET_NAME = "my_asset" +ASSET_URI = "s3://bucket/key" +ALIAS_NAME = "my_alias" + + +# top-level (importable) process_result callables +def only_us_partitions(events: list[Any]) -> list[Any]: + """Keep only events whose partition key targets the ``us`` region.""" + return [event for event in events if (event.partition_key or "").startswith("us|")] + + +def dedup_by_partition_key(events: list[Any]) -> list[Any]: + """Keep the first event seen for each distinct partition key (order preserved).""" + seen: set[str | None] = set() + out: list[Any] = [] + for event in events: + if event.partition_key in seen: + continue + seen.add(event.partition_key) + out.append(event) + return out + + +def drop_all(events: list[Any]) -> list[Any]: + """A process_result that always yields nothing (keeps the trigger waiting).""" + return [] + + +def duplicate_all(events: list[Any]) -> list[Any]: + """A process_result that returns more events than it received.""" + return events + events + + +def boom(events: list[Any]) -> list[Any]: + raise RuntimeError("process_result exploded") + + +class _Stop(Exception): + """Sentinel raised from a patched ``asyncio.sleep`` to break the polling loop.""" + + +def _keys(event: TriggerEvent) -> list[str | None]: + return [e["partition_key"] for e in event.payload["events"]] + + +@pytest.mark.parametrize( + ("count", "expected_count", "result"), + [(0, -1, False), (1, -1, True), (5, -1, True), (2, 2, True), (1, 2, False), (3, 2, False)], +) +def test_count_satisfied(count, expected_count, result): + assert _count_satisfied(count, expected_count) is result + + +def test_serialize_events_uses_model_dump(): + class FakeEvent: + def model_dump(self, mode=None): + return {"dumped": mode} + + assert _serialize_events([FakeEvent(), {"plain": 1}]) == [{"dumped": "json"}, {"plain": 1}] + + +def test_serialize_round_trip(): + trigger = AssetEventTrigger( + name="my_asset", + uri="s3://bucket/key", + partition_key_regexp_pattern="da.*", + extra={"k": "v"}, + expected_count=2, + process_result_path="my.module.func", + poke_interval=30, + ) + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.standard.triggers.asset.AssetEventTrigger" + assert kwargs == { + "name": "my_asset", + "uri": "s3://bucket/key", + "alias_name": None, + "after": None, + "before": None, + "ascending": True, + "limit": None, + "partition_key": None, + "partition_key_regexp_pattern": "da.*", + "extra": {"k": "v"}, + "expected_count": 2, + "process_result_path": "my.module.func", + "poke_interval": 30, + } + # Re-instantiation from the serialized kwargs must work. + AssetEventTrigger(**kwargs) + + +def test_rejects_invalid_expected_count(): + with pytest.raises(ValueError, match="expected_count"): + AssetEventTrigger(name="my_asset", expected_count=-2) + + +def test_serialize_round_trip_with_datetimes(): + from datetime import datetime, timezone + + after = datetime(2024, 1, 1, tzinfo=timezone.utc) + before = datetime(2024, 1, 9, tzinfo=timezone.utc) + trigger = AssetEventTrigger(name="my_asset", after=after, before=before) + _, kwargs = trigger.serialize() + assert kwargs["after"] == after + assert kwargs["before"] == before + reborn = AssetEventTrigger(**kwargs) + assert reborn.after == after + assert reborn.before == before + + +@pytest.mark.db_test +class TestRun: + """run() against real ``AssetEvent`` rows (no patching of the fetch).""" + + @pytest.mark.asyncio + async def test_fires_at_least_one(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger(name=ASSET_NAME, poke_interval=0.01) # expected_count=-1 + event = await trigger.run().__anext__() + assert isinstance(event, TriggerEvent) + assert event.payload["status"] == "success" + assert len(event.payload["events"]) == 5 + + @pytest.mark.asyncio + async def test_fires_on_exact_count(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger(name=ASSET_NAME, expected_count=5, poke_interval=0.01) + event = await trigger.run().__anext__() + assert len(event.payload["events"]) == 5 + + @pytest.mark.asyncio + async def test_exact_partition_key_filter(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger(name=ASSET_NAME, partition_key="us|2024-01-01", poke_interval=0.01) + event = await trigger.run().__anext__() + assert _keys(event) == ["us|2024-01-01", "us|2024-01-01"] + + @pytest.mark.asyncio + async def test_extra_filter(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger(name=ASSET_NAME, extra={"region": "eu"}, poke_interval=0.01) + event = await trigger.run().__anext__() + assert _keys(event) == ["eu|2024-01-01"] + + @pytest.mark.asyncio + async def test_by_alias(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger(alias_name=ALIAS_NAME, poke_interval=0.01) + event = await trigger.run().__anext__() + assert len(event.payload["events"]) == 5 + + @pytest.mark.asyncio + async def test_applies_process_result_filter(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger( + name=ASSET_NAME, + expected_count=3, + process_result_path=f"{__name__}.only_us_partitions", + poke_interval=0.01, + ) + event = await trigger.run().__anext__() + assert _keys(event) == ["us|2024-01-01", "us|2024-01-02", "us|2024-01-01"] + + @pytest.mark.asyncio + async def test_applies_process_result_dedup(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger( + name=ASSET_NAME, + expected_count=4, + process_result_path=f"{__name__}.dedup_by_partition_key", + poke_interval=0.01, + ) + event = await trigger.run().__anext__() + assert _keys(event) == ["us|2024-01-01", "us|2024-01-02", "eu|2024-01-01", None] + + @pytest.mark.asyncio + async def test_sleeps_until_matching_event_appears( + self, asset_event_rows, db_supervisor_comms, add_asset_event, mocker + ): + # No event matches the filter yet; a new one is inserted "during" the first sleep, after + # which the next fetch finds it and the trigger fires. This exercises the real polling loop. + async def _insert(_seconds): + add_asset_event(id=100, day=9, partition_key="zz|only", extra={"region": "zz"}) + + sleep_mock = mocker.patch("asyncio.sleep", new=mock.AsyncMock(side_effect=_insert)) + trigger = AssetEventTrigger(name=ASSET_NAME, partition_key="zz|only", poke_interval=5) + event = await trigger.run().__anext__() + assert event.payload["status"] == "success" + assert _keys(event) == ["zz|only"] + sleep_mock.assert_awaited_once_with(5) + + @pytest.mark.asyncio + async def test_process_result_reduces_below_expected_keeps_waiting( + self, asset_event_rows, db_supervisor_comms, add_asset_event, mocker + ): + # only_us -> 3 events initially (< expected 4), so the trigger must sleep; a 4th "us" event + # is inserted during that sleep, after which the trigger fires. + async def _insert(_seconds): + add_asset_event(id=101, day=9, partition_key="us|2024-01-09", extra={"region": "us"}) + + sleep_mock = mocker.patch("asyncio.sleep", new=mock.AsyncMock(side_effect=_insert)) + trigger = AssetEventTrigger( + name=ASSET_NAME, + expected_count=4, + process_result_path=f"{__name__}.only_us_partitions", + poke_interval=5, + ) + event = await trigger.run().__anext__() + assert len(event.payload["events"]) == 4 + assert all((k or "").startswith("us|") for k in _keys(event)) + sleep_mock.assert_awaited_once_with(5) + + @pytest.mark.asyncio + async def test_callable_serialized_by_sensor_runs_in_trigger(self, asset_event_rows, db_supervisor_comms): + # End-to-end: pass a *real callable* to the sensor. The sensor's serialization util turns it + # into a dotted path and hands it to the trigger; we then run the exact trigger the sensor + # built and confirm the trigger imports and executes that callable against the DB. + sensor = AssetEventSensor( + task_id="s", + name=ASSET_NAME, + process_result=only_us_partitions, # a real function object, not an import path + expected_count=99, # force the initial poke to fail so the sensor defers + deferrable=True, + ) + with pytest.raises(TaskDeferred) as exc: + sensor.execute({}) + trigger = exc.value.trigger + assert isinstance(trigger, AssetEventTrigger) + assert trigger.process_result_path.endswith("test_asset.only_us_partitions") + + # Relax the (deliberately unreachable) expectation so the same trigger can fire, then run it. + trigger.expected_count = 3 + trigger.poke_interval = 0.01 + event = await trigger.run().__anext__() + assert event.payload["status"] == "success" + assert _keys(event) == ["us|2024-01-01", "us|2024-01-02", "us|2024-01-01"] + + @pytest.mark.asyncio + async def test_empty_extra_is_no_filter(self, asset_event_rows, db_supervisor_comms): + # An empty ``extra`` dict must be treated as "no filter", not as an impossible match. + trigger = AssetEventTrigger(name=ASSET_NAME, extra={}, poke_interval=0.01) + event = await trigger.run().__anext__() + assert len(event.payload["events"]) == 5 + + @pytest.mark.asyncio + async def test_process_result_returns_more_events(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger( + name=ASSET_NAME, + expected_count=10, + process_result_path=f"{__name__}.duplicate_all", + poke_interval=0.01, + ) + event = await trigger.run().__anext__() + assert len(event.payload["events"]) == 10 + + @pytest.mark.asyncio + async def test_process_result_exception_propagates(self, asset_event_rows, db_supervisor_comms): + trigger = AssetEventTrigger( + name=ASSET_NAME, process_result_path=f"{__name__}.boom", poke_interval=0.01 + ) + with pytest.raises(RuntimeError, match="process_result exploded"): + await trigger.run().__anext__() + + @pytest.mark.asyncio + async def test_empty_process_result_keeps_waiting(self, asset_event_rows, db_supervisor_comms, mocker): + # process_result drops everything -> count never satisfied -> the trigger must sleep. + # A patched sleep raises to break the otherwise-infinite polling loop after one iteration. + async def _stop(_seconds): + raise _Stop() + + sleep_mock = mocker.patch("asyncio.sleep", new=mock.AsyncMock(side_effect=_stop)) + trigger = AssetEventTrigger( + name=ASSET_NAME, process_result_path=f"{__name__}.drop_all", poke_interval=3 + ) + with pytest.raises(_Stop): + await trigger.run().__anext__() + sleep_mock.assert_awaited_once_with(3)