Add AssetEventSensor to standard provider (#70225)

* Work around mapped inlet population bug

* Tighten limit and expected_count checks

* Make sensor 3.2-compatible

Specific arguments that only make sense in 3.4+ get additional checks
at parse time to be rejected. Importing and using the sensor works on
3.2 and 3.3 as long as you do not use the 3.4+ parameters.

* Remove dotted process_events format

* Better BaseModel dump

* Improve AssetEventSensor docs and example

* Add AssetEventSensor to standard provider

* Scope asset event sensor/trigger tests to Airflow 3.4+ and regenerate provider build files

* Address review: dedicated exception, validation, docstring, and tests for AssetEventSensor

* Coerce datetime after/before to ISO strings, drop type: ignore (Copilot review)

* Add AssetEventSensor howto docs and example DAG

* Fix docs build: drop unresolved asset ref and reword to fix spellcheck

* Drop deferrable support from AssetEventSensor

The execution API calls backing asset event fetching (GetAssetEventByAsset
and GetAssetEventByAssetAlias) are not available in the triggerer, so the
deferrable path could not work in practice. Ship the synchronous sensor
only; deferrable support can follow once the triggerer can fetch asset
events and InletEventsAccessor exposes a native async interface.

Removes AssetEventTrigger and its tests, folds the shared fetch/count/
serialize helpers into the sensor module, and drops the now-unused
asset sensor exceptions.

* Declare asset sensor dependencies and support minimum event counts

Task-level asset waits need visible lineage without changing the Dag schedule. Explicit count policies let consumers wait for a minimum number of events without missing the condition when several events arrive between pokes.

* Reject ambiguous asset sensor targets and invalid event counts

Conflicting selectors can silently wait for a different asset, while fractional counts and invalid limits can produce misleading waits. Result processing also needs clear limits so consumers know which events their callback can inspect.

* Use UTC in asset sensor tests

* Unify sensor asset argument

* Support run-relative asset event bounds

Asset waits need time bounds tied to a run without requiring Jinja expressions
or limiting Dag authors to a fixed set of anchors. A runtime resolver contract
lets authors derive custom bounds from task context.

* Fix AssetEventSensor documentation heading

* Add explicit asset event filters and stable run bounds

* Allow empty extra filters on older Airflow versions

An empty extra dictionary adds no query predicate, so it should not require Airflow 3.4 when another filter is present.

---------

Co-authored-by: Tzu-ping Chung <tp@astronomer.io>
This commit is contained in:
Hussein AwalaandTzu-ping Chung authored and GitHub committed 2026-11-03 11:56:59 +01:00
1 parent 51140809d2
commit 4df59b9a03
6 files changed
+1529

No files matched your search

+304
View File
@@ -0,0 +1,304 @@
.. 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.2+**. The ``partition_key``,
``partition_key_regexp_pattern`` and ``extra`` filters require
**Apache Airflow 3.4+**.
When to use a schedule instead
------------------------------
If the dependency does not need to sit in the middle of a run, a schedule is
usually a better fit. See
:doc:`apache-airflow:authoring-and-scheduling/asset-scheduling` for details.
* Schedule the Dag on the asset when the whole run can wait for it.
* Use ``AssetAndTimeSchedule`` (Airflow 3.4+) when a time-scheduled run needs
the asset before it starts.
* Use a partitioned asset timetable with ``RollupMapper`` to wait for every
partition in a window.
Use this sensor when some tasks should run immediately and only later tasks
depend on the asset. Here, ``prepare`` starts as soon as the run does, and only
``process`` waits for the partition (``wait_for_partition_events`` is defined
:ref:`below <howto/sensor:AssetEventSensor:filtering>`):
.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py
:language: python
:dedent: 4
:start-after: [START example_asset_event_sensor_mid_run]
:end-before: [END example_asset_event_sensor_mid_run]
Basic usage
-----------
Pass an :class:`~airflow.sdk.Asset`, :meth:`~airflow.sdk.Asset.ref`, or
:class:`~airflow.sdk.AssetAlias` as ``asset``. A reference must resolve to an
existing asset when the task starts. The target is 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.
Supply at least one query filter: ``after``, ``before``, ``partition_key``,
``partition_key_regexp_pattern``, or a non-empty ``extra`` dictionary. The sensor
raises ``ValueError`` at construction if no filter is supplied. It checks again
after template rendering, before querying events. An empty ``extra={}`` does
not count as a filter.
Time bounds are optional when another query filter is supplied. Ordering,
``limit``, and ``process_result`` do not replace query filters.
This daily example waits for events at or after the data interval ends:
.. 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]
A mapped sensor works, but its asset is only recorded as an inlet in lineage if
it is also passed through ``partial(inlets=[...])``.
Run-relative time bounds
------------------------
The ``after`` and ``before`` arguments accept datetime values, ISO strings,
Jinja templates, and these bound objects from
``airflow.providers.standard.sensors.asset``:
.. list-table:: Choosing a time bound
:header-rows: 1
:widths: 20 25 25 30
* - Bound
- Available on
- When tasks are cleared
- Watch out for
* - ``RunStart()``
- All runs, including asset-triggered runs and backfills.
- A finished run gets a new start time when it starts again. Clearing tasks
in a queued or running run preserves the existing start time.
- Uses ``dag_run.start_date``, the actual start even for a backfill.
With ``after``, a negative offset can include events from an earlier execution.
* - ``RunAfter()``
- All runs, including asset-triggered runs and backfills.
- Stays fixed for the same run.
- Uses ``dag_run.run_after``. For some asset timetables, this equals a
triggering event's timestamp. Since ``after`` is inclusive, that event
can satisfy the sensor immediately.
* - ``LogicalDate()``
- Runs with a logical date, including manual runs that supply one.
Asset-triggered runs have none.
- Stays fixed for the same run.
- Its meaning depends on the timetable or the value supplied for a manual
run. It does not necessarily represent when execution started.
* - ``DataIntervalStart()`` / ``DataIntervalEnd()``
- Runs with a data interval. Availability depends on the timetable and
run type. Asset-triggered runs have none.
- Stay fixed for the same run.
- Use the interval represented by the run for historical work.
``CronTriggerTimetable`` defaults to a zero-length interval, so both
bounds are equal.
Use ``after=RunStart()`` when clearing a finished run should wait for new data.
Use ``after=RunAfter()`` to keep the same time boundary when an asset-triggered
run is cleared. Runs with logical dates or data intervals can also use those
fixed bounds. An open-ended query can still include new events that arrive
later, so fixed bounds do not freeze the result set.
Each built-in object accepts a signed ``offset`` duration, with zero as the
default. A negative offset moves the bound earlier. A positive offset moves it
later. For example, ``after=DataIntervalEnd(offset=-timedelta(hours=1))`` includes
events from one hour before the interval ends onward.
``before=DataIntervalEnd(offset=timedelta(hours=1))`` includes events up to one
hour after the interval ends. Import ``timedelta`` from ``datetime`` to use these
examples.
To exclude events exactly at ``dag_run.run_after``, use
``after=RunAfter(offset=timedelta(microseconds=1))``. Later events can still match,
including events seen during an earlier execution of the run.
For example, use both interval bounds to select a historical window:
.. code-block:: python
from airflow.providers.standard.sensors.asset import (
AssetEventSensor,
DataIntervalEnd,
DataIntervalStart,
)
wait_for_interval_events = AssetEventSensor(
task_id="wait_for_interval_events",
asset=my_asset,
after=DataIntervalStart(),
before=DataIntervalEnd(),
)
Both bounds are inclusive and filter the asset event's timestamp. An event at
an interval boundary can therefore match two adjacent intervals. These bounds
do not filter the event's partition key.
To create a nonzero window with ``CronTriggerTimetable``, use an offset or
configure the timetable's ``interval``.
A bound raises ``ValueError`` if its selected timestamp is missing; it does not
substitute another timestamp. Manual runs can have a data interval when a logical
date is supplied. For asset-triggered runs, use ``RunStart()`` or ``RunAfter()``
instead. See the table for how they differ when a run is cleared.
Airflow resolves the objects during template rendering, before polling. Their
run-relative values stay fixed while the sensor polls and after rescheduling.
Custom resolvers
~~~~~~~~~~~~~~~~
Pass any object with a ``resolve(context)`` method that returns a datetime or
ISO string. :class:`~airflow.providers.standard.sensors.asset.AssetEventTimeBound`
defines this interface. Inheritance is optional.
For stable Dag serialization, give custom objects a deterministic ``__repr__``
that includes their configuration. Airflow currently uses this representation
for ``partial`` and ``expand`` arguments. Older versions, including Airflow 3.2.0,
also use it for unmapped template fields. The default object representation
contains a memory address and can create a new Dag version on each parse.
This example includes the parameter name in its representation:
.. code-block:: python
from airflow.providers.standard.sensors.asset import AssetEventSensor, AssetEventTimeBound
class ParameterBound(AssetEventTimeBound):
def __init__(self, name):
self.name = name
def __repr__(self):
return f"ParameterBound(name={self.name!r})"
def resolve(self, context):
return context["params"][self.name]
wait_for_custom_events = AssetEventSensor.partial(
task_id="wait_for_custom_events",
asset=my_asset,
inlets=[my_asset],
after=ParameterBound("events_after"),
params={"events_after": "2024-01-01T00:00:00+00:00"},
).expand(limit=[1, 2])
Custom resolvers do not need ``serialize`` or ``deserialize`` methods. The
``offset`` argument is a convenience provided by the built-in helpers.
Alternatively, Airflow 3.2+ supports a template callable that accepts the ``context`` and
``jinja_env`` keyword arguments and returns a datetime or ISO string. Airflow
evaluates the callable during template rendering. Custom resolvers and callables
must return the same bound after rescheduling. A bound based on the current clock
can move the query window and exclude events that arrived between pokes.
.. _howto/sensor:AssetEventSensor:filtering:
Filtering events and expected count
-----------------------------------
Filter events with ``after``, ``before`` (time range), ``partition_key``,
``partition_key_regexp_pattern``, or ``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.
Use ``ascending`` to set the order and ``limit`` to cap the number of matching
events returned.
``partition_key_regexp_pattern`` requires a positive
:ref:`[api] regexp_query_timeout <apache-airflow:config:api__regexp_query_timeout>`.
With the default ``0``, regexp filtering is disabled, the query is rejected,
and the task fails.
Use ``expected_count`` and ``count_policy`` to control how many processed
events are required:
* ``count_policy="minimum"`` (the default) succeeds at or above
``expected_count``, which defaults to one.
* ``count_policy="exact"`` succeeds only at ``expected_count``. Use this policy
with zero to check that no matching events exist. An exact count can be
missed if several events arrive between pokes.
Counts must be non-negative integers, and ``limit`` must be a positive integer
when supplied. With the *minimum* policy, zero always satisfies the count
condition.
If a ``process_result`` callback is not provided, the sensor raises *ValueError*
at construction if ``limit`` is less than ``expected_count`` with the *minimum*
policy, or not greater than ``expected_count`` with the *exact* policy.
.. 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 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. It runs on every
poke, so it should be idempotent and free of side effects.
The query applies ``limit`` **before** processing. A callback may expand the fetched list, so
``limit < expected_count`` is allowed when a callback is supplied; the count is checked on its
output. Choose the limit with any filtering or deduplication in mind.
With ascending order and a limit, each poke returns the same oldest matching events while those
events remain in the query range. If the callback rejects them, newer acceptable events beyond
the limit are never inspected. Prefer query filters such as ``partition_key``, ``extra`` and time
bounds where possible. For a check of recent events, use ``ascending=False`` with a suitable limit.
The example below keeps the latest event at or after the data interval ends.
If no post-processing is needed, use ``ascending=False, limit=1`` with the query
filter to fetch only its latest matching event.
.. 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]
.. exampleinclude:: /../src/airflow/providers/standard/example_dags/example_asset_sensor.py
:language: python
:dedent: 4
:start-after: [START example_asset_event_sensor_process_result]
:end-before: [END example_asset_event_sensor_process_result]
+2
View File
@@ -87,6 +87,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
@@ -104,6 +105,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
@@ -0,0 +1,85 @@
# 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.operators.empty import EmptyOperator
from airflow.providers.standard.sensors.asset import AssetEventSensor, DataIntervalEnd
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."""
return events[-1:]
# [END example_asset_event_process_result]
with DAG(
dag_id="example_asset_sensor",
schedule="@daily",
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,
after=DataIntervalEnd(),
)
# [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="{{ data_interval_start | ds }}",
expected_count=3,
count_policy="minimum",
)
# [END example_asset_event_sensor_filtered]
# [START example_asset_event_sensor_process_result]
wait_for_latest_asset_event = AssetEventSensor(
task_id="wait_for_latest_asset_event",
asset=my_asset,
after=DataIntervalEnd(),
process_result=latest_event_only,
)
# [END example_asset_event_sensor_process_result]
# [START example_asset_event_sensor_mid_run]
prepare = EmptyOperator(task_id="prepare")
process = EmptyOperator(task_id="process")
prepare >> wait_for_partition_events >> process
# [END example_asset_event_sensor_mid_run]
@@ -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",
@@ -0,0 +1,330 @@
# 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 contextlib
from abc import abstractmethod
from datetime import timedelta
from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, cast
from pydantic import BaseModel
from airflow.providers.common.compat.sdk import (
AirflowOptionalProviderFeatureException,
BaseSensorOperator,
PokeReturnValue,
)
from airflow.providers.standard.version_compat import AIRFLOW_V_3_2_PLUS, AIRFLOW_V_3_4_PLUS
if not AIRFLOW_V_3_2_PLUS:
raise AirflowOptionalProviderFeatureException("AssetEventSensor needs Airflow 3.2+.")
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 Asset, AssetAlias, AssetRef
from airflow.sdk.execution_time.comms import AssetEventResult
from airflow.sdk.execution_time.context import InletEventsAccessor
def _serialize_events(events: list[Any]) -> list[Any]:
"""
Convert asset events Pydantic models to JSON-compatible dicts.
This keeps the XCom value plain JSON, rather than relying on how XCom
serializes pydantic models.
"""
return [e.model_dump(mode="json") if isinstance(e, BaseModel) else e for e in events]
class AssetEventTimeBound(Protocol):
"""
A time bound that Airflow resolves during template rendering.
Implement ``resolve(context)`` to return a datetime or ISO string. Custom
objects do not need to inherit this protocol.
The result must stay fixed across reschedules of the same Dag run.
"""
def resolve(self, context: Context) -> datetime | str:
"""Return the event timestamp bound for this task's runtime context."""
...
class _ContextTimeBound(AssetEventTimeBound):
# This is a serialized identifier; changing its value changes the serialized Dag's hash.
_anchor: ClassVar[str]
def __init__(self, *, offset: timedelta = timedelta(0)) -> None:
self.offset = offset
def __repr__(self) -> str:
return f"{type(self).__name__}(offset={self.offset!r})"
def serialize(self) -> dict[str, str | float]:
"""Return the bound configuration for serialized task metadata."""
return {"anchor": self._anchor, "offset_seconds": self.offset.total_seconds()}
@abstractmethod
def _get_timestamp(self, context: Context) -> datetime | None: ...
def resolve(self, context: Context) -> datetime:
"""Resolve the bound during template rendering, before the sensor polls."""
timestamp = self._get_timestamp(context)
if timestamp is None:
raise ValueError(f"{type(self).__name__} requires {self._anchor!r} in the task context.")
return timestamp + self.offset
class RunStart(_ContextTimeBound):
"""
Bound events by the Dag run's actual start time, including for backfills.
:param offset: Signed duration to add to ``dag_run.start_date``.
"""
_anchor = "dag_run.start_date"
def _get_timestamp(self, context: Context) -> datetime | None:
dag_run = context.get("dag_run")
return dag_run.start_date if dag_run is not None else None
class RunAfter(_ContextTimeBound):
"""
Bound events by the Dag run's earliest allowed start time.
This timestamp stays fixed when tasks are cleared. With ``after``, an event
at this timestamp is included, even if it triggered the run.
:param offset: Signed duration to add to ``dag_run.run_after``.
"""
_anchor = "dag_run.run_after"
def _get_timestamp(self, context: Context) -> datetime | None:
dag_run = context.get("dag_run")
return dag_run.run_after if dag_run is not None else None
class LogicalDate(_ContextTimeBound):
"""
Bound events by the Dag run's logical date.
:param offset: Signed duration to add to ``logical_date``.
"""
_anchor = "logical_date"
def _get_timestamp(self, context: Context) -> datetime | None:
return context.get("logical_date")
class DataIntervalStart(_ContextTimeBound):
"""
Bound events by the start of the Dag run's data interval.
:param offset: Signed duration to add to ``data_interval_start``.
"""
_anchor = "data_interval_start"
def _get_timestamp(self, context: Context) -> datetime | None:
return context.get("data_interval_start")
class DataIntervalEnd(_ContextTimeBound):
"""
Bound events by the end of the Dag run's data interval.
:param offset: Signed duration to add to ``data_interval_end``.
"""
_anchor = "data_interval_end"
def _get_timestamp(self, context: Context) -> datetime | None:
return context.get("data_interval_end")
class AssetEventSensor(BaseSensorOperator):
"""
Wait for asset events matching the given filters to reach an expected count.
The sensor fetches events of an asset, asset reference, 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``.
At least one of ``after``, ``before``, ``partition_key``,
``partition_key_regexp_pattern``, or non-empty ``extra`` is required.
This sensor requires Apache Airflow 3.2+. The ``partition_key``,
``partition_key_regexp_pattern`` and ``extra`` filters require Airflow 3.4+.
:param asset: The :class:`~airflow.sdk.Asset`, :meth:`~airflow.sdk.Asset.ref`,
or :class:`~airflow.sdk.AssetAlias` to wait on. It is declared as an
inlet. A reference must resolve to an existing asset when the task
starts. The target is static; event filters can be templated.
:param after: Only include events at or after this timestamp.
:param before: Only include events at or before this timestamp.
Both bounds accept a datetime, ISO string, any object implementing
:class:`AssetEventTimeBound`, or a template callable accepting
``context`` and ``jinja_env``. Airflow
resolves them before polling. Custom callables must return stable
bounds across reschedules.
: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``.
:param expected_count: The non-negative number of processed events required
to succeed. Defaults to 1. If ``process_result`` is omitted, ``limit``
must not be smaller, or must be greater with the ``"exact"`` policy.
: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 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 and free of side-effects. Its return value must
be JSON-serializable to be pushed into XCom.
"""
template_fields: Sequence[str] = (
"partition_key",
"partition_key_regexp_pattern",
"extra",
"after",
"before",
)
def __init__(
self,
*,
asset: Asset | AssetAlias | AssetRef,
after: datetime | str | AssetEventTimeBound | Callable[..., datetime | str] | None = None,
before: datetime | str | AssetEventTimeBound | Callable[..., 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[AssetEventResult]], list[Any]] | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
if expected_count < 0:
raise ValueError(f"`expected_count` must be a non-negative integer, got {expected_count}.")
if limit is not None and 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:
if count_policy == "exact" and limit <= expected_count:
raise ValueError(
"`limit` must be greater than `expected_count` with the exact count policy, "
"otherwise extra events cannot be detected."
)
if limit < expected_count:
raise ValueError(
"`limit` must be at least `expected_count` when `process_result` is not provided."
)
self.asset = asset
if asset not in self.inlets:
self.inlets.append(asset)
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
self._validate_filters()
def _validate_filters(self) -> None:
if not AIRFLOW_V_3_4_PLUS and (
self.partition_key is not None or self.partition_key_regexp_pattern is not None or self.extra
):
raise AirflowOptionalProviderFeatureException(
"AssetEventSensor filters partition_key, partition_key_regexp_pattern and extra "
"need Airflow 3.4+."
)
if (
self.after is None
and self.before is None
and self.partition_key is None
and self.partition_key_regexp_pattern is None
and not self.extra
):
raise ValueError(
"AssetEventSensor requires at least one filter: after, before, partition_key, "
"partition_key_regexp_pattern, or non-empty extra."
)
def _get_accessor(self, context: Context) -> InletEventsAccessor:
# The mapped operator's inlets don't include the asset added in init.
# Remove when the minimum Airflow includes https://github.com/apache/airflow/pull/75751.
with contextlib.suppress(KeyError):
return context["inlet_events"][self.asset]
from airflow.sdk.execution_time.context import InletEventsAccessors
return InletEventsAccessors([self.asset])[self.asset]
def poke(self, context: Context) -> PokeReturnValue:
self._validate_filters()
# Template rendering replaces bound objects and callables before polling.
after = cast("datetime | str | None", self.after)
before = cast("datetime | str | None", self.before)
accessor = self._get_accessor(context)
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(self.ascending)
if self.limit is not None:
accessor.limit(self.limit)
if self.partition_key is not None:
accessor.partition_key(self.partition_key)
if self.partition_key_regexp_pattern is not None:
accessor.partition_key_regexp_pattern(self.partition_key_regexp_pattern)
if self.extra is not None:
for key, value in self.extra.items():
accessor.extra(key, value)
processed = list(accessor)
if self.process_result is not None:
processed = self.process_result(processed)
count = len(processed)
done = (
count >= self.expected_count if self.count_policy == "minimum" else count == self.expected_count
)
self.log.info(
"Found %d matching asset events (expected %s, policy %s): %s",
count,
self.expected_count,
self.count_policy,
"condition met" if done else "still waiting",
)
return PokeReturnValue(is_done=done, xcom_value=_serialize_events(processed) if done else None)
@@ -0,0 +1,806 @@
# 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 json
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from typing import Any
from unittest.mock import call
import pytest
from airflow.providers.common.compat.sdk import (
DAG,
AirflowOptionalProviderFeatureException,
Asset,
AssetAlias,
PokeReturnValue,
)
from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS, AIRFLOW_V_3_4_PLUS
if AIRFLOW_V_3_2_PLUS:
from airflow.exceptions import AirflowRescheduleException
from airflow.providers.standard.sensors import asset as asset_sensor
from airflow.providers.standard.sensors.asset import AssetEventSensor, AssetEventTimeBound
from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse
from airflow.sdk.execution_time.comms import (
AssetEventsResult,
AssetResult,
GetAssetByName,
GetAssetByUri,
GetAssetEventByAsset,
GetAssetEventByAssetAlias,
TaskRescheduleStartDate,
)
from airflow.sdk.execution_time.context import InletEventsAccessors
from tests_common.test_utils.compat import DagSerialization, OperatorSerialization
class _ParameterBound:
def resolve(self, context):
return context["params"]["events_after"]
class _InheritedParameterBound(AssetEventTimeBound):
def resolve(self, context):
return context["params"]["events_after"]
class _MappedParameterBound:
def __init__(self, name):
self.name = name
def __repr__(self):
return f"_MappedParameterBound(name={self.name!r})"
def resolve(self, context):
return context["params"][self.name]
pytestmark = pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="AssetEventSensor requires Airflow 3.2+")
requires_3_4 = pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Filter requires Airflow 3.4+")
ASSET_NAME = "my_asset"
ASSET_URI = "s3://bucket/key"
ALIAS_NAME = "my_alias"
def _get_timestamp(day: int) -> datetime:
return datetime(2024, 1, day, tzinfo=UTC)
def _get_context(sensor: AssetEventSensor) -> dict[str, Any]:
return {"inlet_events": InletEventsAccessors(sensor.inlets)}
def only_us_partitions(events: list[Any]) -> list[Any]:
return [event for event in events if (event.partition_key or "").startswith("us|")]
def dedup_by_partition_key(events: list[Any]) -> list[Any]:
seen: set[str | None] = set()
result = []
for event in events:
if event.partition_key not in seen:
seen.add(event.partition_key)
result.append(event)
return result
@pytest.fixture
def asset():
return Asset(name=ASSET_NAME, uri=ASSET_URI)
@pytest.fixture
def events_response():
partition_keys = ["us|2024-01-01", "us|2024-01-02", "eu|2024-01-01", "us|2024-01-01", None]
return AssetEventsResult(
asset_events=[
AssetEventResponse(
id=index,
timestamp=_get_timestamp(index),
partition_key=key,
extra={"region": key.split("|")[0]} if key else {},
asset=AssetResponse(name=ASSET_NAME, uri=ASSET_URI, group="asset", extra={}),
created_dagruns=[],
source_dag_id="producer",
source_task_id="emit",
source_run_id="run_1",
source_map_index=-1,
)
for index, key in enumerate(partition_keys, start=1)
]
)
@pytest.fixture(
params=[
("RunStart", "dag_run.start_date", 8),
("RunAfter", "dag_run.run_after", 7),
("LogicalDate", "logical_date", 1),
("DataIntervalStart", "data_interval_start", 2),
("DataIntervalEnd", "data_interval_end", 3),
],
ids=lambda case: case[0],
)
def time_bound_case(request):
name, anchor, day = request.param
return getattr(asset_sensor, name), anchor, _get_timestamp(day)
@pytest.fixture
def time_bound_context():
return {
"dag_run": SimpleNamespace(start_date=_get_timestamp(8), run_after=_get_timestamp(7)),
"logical_date": _get_timestamp(1),
"data_interval_start": _get_timestamp(2),
"data_interval_end": _get_timestamp(3),
}
class TestInit:
@pytest.mark.parametrize(
"kwargs",
[
{},
{"after": None, "before": None},
{"limit": 1},
{"ascending": False},
{"process_result": only_us_partitions},
{"expected_count": 0, "count_policy": "exact"},
{"extra": {}},
],
ids=["none", "null-bounds", "limit", "order", "processing", "count", "empty-extra"],
)
def test_requires_query_filter(self, asset, kwargs):
with pytest.raises(ValueError, match="requires at least one filter"):
AssetEventSensor(task_id="s", asset=asset, **kwargs)
@pytest.mark.parametrize(
"filters",
[
{"after": _get_timestamp(1)},
{"before": _get_timestamp(2)},
pytest.param({"partition_key": "a"}, marks=requires_3_4),
pytest.param({"partition_key_regexp_pattern": "a.*"}, marks=requires_3_4),
pytest.param({"extra": {"region": "us"}}, marks=requires_3_4),
],
ids=["after", "before", "partition", "partition-regexp", "extra"],
)
def test_accepts_each_query_filter_alone(self, asset, mock_supervisor_comms, filters):
sensor = AssetEventSensor(task_id="s", asset=asset, **filters)
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
sensor.poke(_get_context(sensor))
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, **filters)
)
@pytest.mark.parametrize("use_alias", [False, True])
def test_declares_target_as_inlet(self, asset, use_alias):
target = AssetAlias(name=ALIAS_NAME) if use_alias else asset
sensor = AssetEventSensor(task_id="s", asset=target, after=_get_timestamp(1))
assert sensor.asset == target
assert sensor.inlets == [target]
@pytest.mark.parametrize("include_target", [False, True])
def test_preserves_existing_inlets_without_duplicates(self, asset, include_target):
other_asset = Asset("s3://bucket/other")
inlets = [other_asset, asset] if include_target else [other_asset]
original_inlets = inlets.copy()
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), inlets=inlets)
assert sensor.inlets == [other_asset, asset]
assert inlets == original_inlets
def test_rejects_negative_expected_count(self, asset):
with pytest.raises(ValueError, match="expected_count"):
AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), expected_count=-1)
@pytest.mark.parametrize("limit", [0, -1])
def test_rejects_nonpositive_limit(self, asset, limit):
with pytest.raises(ValueError, match="limit"):
AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), limit=limit, expected_count=0)
def test_rejects_invalid_count_policy(self, asset):
with pytest.raises(ValueError, match="count_policy"):
AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), count_policy="unknown")
@pytest.mark.parametrize(
("count_policy", "limit", "expected_count"),
[("minimum", 1, 3), ("exact", 1, 3), ("exact", 3, 3)],
)
def test_rejects_insufficient_limit_without_processing(self, asset, count_policy, limit, expected_count):
with pytest.raises(ValueError, match="limit"):
AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
limit=limit,
expected_count=expected_count,
count_policy=count_policy,
)
@pytest.mark.parametrize(("count_policy", "limit"), [("minimum", 3), ("exact", 4)])
def test_accepts_sufficient_limit_without_processing(self, asset, count_policy, limit):
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
limit=limit,
expected_count=3,
count_policy=count_policy,
)
assert sensor.limit == limit
@pytest.mark.parametrize(
"filters",
[{"partition_key": "a"}, {"partition_key_regexp_pattern": "a.*"}, {"extra": {"k": "v"}}],
)
def test_partition_and_extra_filters_require_airflow_3_4(self, mocker, asset, filters):
mocker.patch("airflow.providers.standard.sensors.asset.AIRFLOW_V_3_4_PLUS", False)
with pytest.raises(AirflowOptionalProviderFeatureException, match="need Airflow 3.4"):
AssetEventSensor(task_id="s", asset=asset, **filters)
@pytest.mark.parametrize("extra", [None, {}], ids=["no-extra", "empty-extra"])
def test_other_filters_do_not_require_airflow_3_4(self, mocker, asset, mock_supervisor_comms, extra):
mocker.patch("airflow.providers.standard.sensors.asset.AIRFLOW_V_3_4_PLUS", False)
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), limit=1, extra=extra)
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
assert not sensor.poke(_get_context(sensor))
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(1), limit=1)
)
def test_empty_extra_requires_another_filter_before_airflow_3_4(self, mocker, asset):
mocker.patch("airflow.providers.standard.sensors.asset.AIRFLOW_V_3_4_PLUS", False)
with pytest.raises(ValueError, match="requires at least one filter"):
AssetEventSensor(task_id="s", asset=asset, extra={})
def test_template_fields(self):
assert set(AssetEventSensor.template_fields) == {
"partition_key",
"partition_key_regexp_pattern",
"extra",
"after",
"before",
}
@requires_3_4
def test_renders_filters_without_changing_lineage(self, asset):
sensor = AssetEventSensor(
task_id="s",
asset=asset,
partition_key="{{ ds }}",
partition_key_regexp_pattern="{{ params.region }}.*",
extra={"region": "{{ params.region }}"},
after="{{ ds }}T00:00:00+00:00",
before="{{ ds }}T23:59:59+00:00",
)
sensor.render_template_fields({"ds": "2024-01-01", "params": {"region": "us"}})
assert sensor.partition_key == "2024-01-01"
assert sensor.partition_key_regexp_pattern == "us.*"
assert sensor.extra == {"region": "us"}
assert sensor.after == "2024-01-01T00:00:00+00:00"
assert sensor.before == "2024-01-01T23:59:59+00:00"
assert sensor.asset == asset
assert sensor.inlets == [asset]
class TestTimeBounds:
@pytest.mark.parametrize("field", ["after", "before"])
@pytest.mark.parametrize("inherit_protocol", [False, True], ids=["plain", "subclass"])
@pytest.mark.parametrize("as_string", [False, True], ids=["datetime", "iso"])
def test_custom_resolver_needs_only_resolve(
self, asset, events_response, mock_supervisor_comms, field, inherit_protocol, as_string
):
bound = _InheritedParameterBound() if inherit_protocol else _ParameterBound()
timestamp = _get_timestamp(8)
sensor = AssetEventSensor(task_id="s", asset=asset, **{field: bound})
context = {
**_get_context(sensor),
"params": {"events_after": timestamp.isoformat() if as_string else timestamp},
}
sensor.render_template_fields(context)
mock_supervisor_comms.send.side_effect = [AssetEventsResult(asset_events=[]), events_response]
assert not sensor.poke(context)
assert sensor.poke(context)
query = GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, **{field: timestamp})
assert mock_supervisor_comms.send.call_args_list == [call(query), call(query)]
def test_custom_resolver_serialization_is_optional(self, asset):
with DAG("test_custom_bound_serialization"):
sensor = AssetEventSensor(task_id="s", asset=asset, after=_ParameterBound())
encoded = json.loads(json.dumps(OperatorSerialization.serialize_operator(sensor)))
assert isinstance(encoded["after"], str)
assert "_ParameterBound" in encoded["after"]
@pytest.mark.parametrize(
"offset", [timedelta(0), -timedelta(hours=1, microseconds=7), timedelta(hours=1, microseconds=7)]
)
def test_resolves_selected_timestamp(self, time_bound_case, time_bound_context, offset):
bound_type, _, timestamp = time_bound_case
assert bound_type(offset=offset).resolve(time_bound_context) == timestamp + offset
@pytest.mark.parametrize("missing", [False, True], ids=["none", "absent"])
def test_rejects_missing_timestamp(self, time_bound_case, time_bound_context, missing):
bound_type, anchor, _ = time_bound_case
if anchor.startswith("dag_run."):
setattr(time_bound_context["dag_run"], anchor.removeprefix("dag_run."), None)
if missing:
del time_bound_context["dag_run"]
elif missing:
del time_bound_context[anchor]
else:
time_bound_context[anchor] = None
with pytest.raises(ValueError, match=rf"{bound_type.__name__} requires"):
bound_type().resolve(time_bound_context)
@pytest.mark.parametrize("offset", [-timedelta(hours=1, microseconds=7), timedelta(hours=2)])
def test_displays_bound_configuration_in_serialized_task(self, asset, time_bound_case, offset):
bound_type, anchor, _ = time_bound_case
bound = bound_type(offset=offset)
metadata = {"anchor": anchor, "offset_seconds": offset.total_seconds()}
assert json.loads(json.dumps(bound.serialize())) == metadata
with DAG("test_bounds_serialization"):
sensor = AssetEventSensor(task_id="s", asset=asset, after=bound)
encoded = json.loads(json.dumps(OperatorSerialization.serialize_operator(sensor)))
assert encoded["after"] in (metadata, repr(bound))
@pytest.mark.parametrize("field", ["after", "before"])
def test_renders_and_queries_bound(
self, asset, time_bound_case, time_bound_context, field, mock_supervisor_comms
):
bound_type, _, timestamp = time_bound_case
offset = -timedelta(hours=1)
bound = bound_type(offset=offset)
sensor = AssetEventSensor(task_id="s", asset=asset, **{field: bound})
context = {**time_bound_context, **_get_context(sensor)}
sensor.render_template_fields(context)
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
sensor.poke(context)
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, **{field: timestamp + offset})
)
assert bound.offset == offset
@pytest.mark.parametrize("offset", [timedelta(0), timedelta(microseconds=1)])
def test_run_after_preserves_explicit_offset_in_query(
self, asset, time_bound_context, mock_supervisor_comms, offset
):
sensor = AssetEventSensor(task_id="s", asset=asset, after=asset_sensor.RunAfter(offset=offset))
context = {**time_bound_context, **_get_context(sensor)}
sensor.render_template_fields(context)
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
sensor.poke(context)
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(7) + offset)
)
@pytest.mark.parametrize("field", ["after", "before"])
@pytest.mark.parametrize("as_string", [False, True])
def test_custom_callable_resolves_before_polling(
self, asset, events_response, mock_supervisor_comms, mocker, field, as_string
):
timestamp = _get_timestamp(8)
def get_bound(*, context, jinja_env):
return timestamp.isoformat() if as_string else timestamp
get_bound_spy = mocker.create_autospec(get_bound, side_effect=get_bound)
sensor = AssetEventSensor(task_id="s", asset=asset, **{field: get_bound_spy})
context = _get_context(sensor)
sensor.render_template_fields(context)
mock_supervisor_comms.send.side_effect = [AssetEventsResult(asset_events=[]), events_response]
assert not sensor.poke(context)
assert sensor.poke(context)
get_bound_spy.assert_called_once()
query = GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, **{field: timestamp})
assert mock_supervisor_comms.send.call_args_list == [call(query), call(query)]
@pytest.mark.parametrize("mapped_bound", [False, True], ids=["partial", "expand"])
@pytest.mark.parametrize(("bound_name", "anchor_day"), [("RunStart", 8), ("RunAfter", 7)])
def test_mapped_sensor(
self,
asset,
events_response,
create_runtime_ti,
mock_supervisor_comms,
mapped_bound,
bound_name,
anchor_day,
):
bound = getattr(asset_sensor, bound_name)(offset=-timedelta(hours=1))
with DAG("test_mapped_bounds"):
if mapped_bound:
mapped = AssetEventSensor.partial(task_id="s", asset=asset).expand(after=[bound, bound])
else:
mapped = AssetEventSensor.partial(task_id="s", asset=asset, after=bound).expand(limit=[1, 2])
ti = create_runtime_ti(task=mapped, map_index=1, start_date=_get_timestamp(8), logical_date=None)
context = ti.get_template_context()
context["dag_run"].run_after = _get_timestamp(7)
ti.render_templates(context=context)
mock_supervisor_comms.send.return_value = events_response
assert ti.task.poke(context)
mock_supervisor_comms.send.assert_called_with(
GetAssetEventByAsset(
name=ASSET_NAME,
uri=ASSET_URI,
after=_get_timestamp(anchor_day) - timedelta(hours=1),
limit=None if mapped_bound else 2,
)
)
@pytest.mark.parametrize("mapped_bound", [False, True], ids=["partial", "expand"])
def test_mapped_custom_resolver(
self, asset, events_response, create_runtime_ti, mock_supervisor_comms, mapped_bound
):
bound = _MappedParameterBound("events_after")
timestamp = _get_timestamp(8)
with DAG("test_mapped_custom_bound"):
if mapped_bound:
mapped = AssetEventSensor.partial(
task_id="s", asset=asset, params={"events_after": timestamp.isoformat()}
).expand(after=[bound, bound])
else:
mapped = AssetEventSensor.partial(
task_id="s",
asset=asset,
after=bound,
params={"events_after": timestamp.isoformat()},
).expand(limit=[1, 2])
ti = create_runtime_ti(task=mapped, map_index=1)
context = ti.get_template_context()
ti.render_templates(context=context)
mock_supervisor_comms.send.return_value = events_response
assert ti.task.poke(context)
mock_supervisor_comms.send.assert_called_with(
GetAssetEventByAsset(
name=ASSET_NAME,
uri=ASSET_URI,
after=timestamp,
limit=None if mapped_bound else 2,
)
)
@pytest.mark.parametrize("mapped_bound", [False, True], ids=["partial", "expand"])
@pytest.mark.parametrize(
"bound_name",
["RunStart", "RunAfter", "LogicalDate", "DataIntervalStart", "DataIntervalEnd", "custom"],
)
def test_mapped_bound_serialization_is_stable(self, asset, mapped_bound, bound_name):
def build_dag(configuration):
if bound_name == "custom":
bound = _MappedParameterBound(f"events_after_{configuration}")
else:
bound = getattr(asset_sensor, bound_name)(offset=timedelta(seconds=configuration))
with DAG("test_mapped_bound_serialization") as dag:
if mapped_bound:
AssetEventSensor.partial(task_id="s", asset=asset).expand(after=[bound])
else:
AssetEventSensor.partial(task_id="s", asset=asset, after=bound).expand(limit=[1])
return dag
first, second, changed = build_dag(1), build_dag(1), build_dag(2)
first_encoded = json.loads(json.dumps(DagSerialization.to_dict(first)))
assert first_encoded == json.loads(json.dumps(DagSerialization.to_dict(second)))
assert first_encoded != json.loads(json.dumps(DagSerialization.to_dict(changed)))
def test_cleared_run_changes_only_start_time_bound(self, asset, time_bound_case, time_bound_context):
bound_type, anchor, timestamp = time_bound_case
first = AssetEventSensor(task_id="s", asset=asset, after=bound_type())
first.render_template_fields(time_bound_context)
time_bound_context["dag_run"].start_date = _get_timestamp(9)
cleared = AssetEventSensor(task_id="s", asset=asset, after=bound_type())
cleared.render_template_fields(time_bound_context)
assert first.after == timestamp
assert cleared.after == (_get_timestamp(9) if anchor == "dag_run.start_date" else timestamp)
@pytest.mark.parametrize(("bound_name", "anchor_day"), [("RunStart", 8), ("RunAfter", 7)])
def test_reconstructed_sensor_keeps_bound_after_rescheduling(
self,
asset,
events_response,
create_runtime_ti,
mock_supervisor_comms,
time_machine,
bound_name,
anchor_day,
):
time_machine.move_to(_get_timestamp(10), tick=False)
bound_type = getattr(asset_sensor, bound_name)
first = AssetEventSensor(task_id="s", asset=asset, after=bound_type(), mode="reschedule")
ti = create_runtime_ti(task=first, start_date=_get_timestamp(8), logical_date=None)
context = ti.get_template_context()
context["dag_run"].run_after = _get_timestamp(7)
ti.render_templates(context=context)
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
with pytest.raises(AirflowRescheduleException):
ti.task.execute(context)
time_machine.shift(timedelta(hours=1))
resumed = AssetEventSensor(task_id="s", asset=asset, after=bound_type(), mode="reschedule")
resumed_ti = create_runtime_ti(
task=resumed, start_date=_get_timestamp(8), logical_date=None, task_reschedule_count=1
)
resumed_context = resumed_ti.get_template_context()
resumed_context["dag_run"].run_after = _get_timestamp(7)
resumed_ti.render_templates(context=resumed_context)
mock_supervisor_comms.send.side_effect = [
TaskRescheduleStartDate(start_date=_get_timestamp(10)),
events_response,
]
assert resumed_ti.task.execute(resumed_context)
query = GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(anchor_day))
assert [c for c in mock_supervisor_comms.send.call_args_list if c == call(query)] == [
call(query),
call(query),
]
class TestPoke:
@pytest.mark.parametrize(
("field", "rendered_value"),
[
("after", None),
("before", None),
pytest.param("partition_key", None, marks=requires_3_4),
pytest.param("partition_key_regexp_pattern", None, marks=requires_3_4),
pytest.param("extra", None, marks=requires_3_4),
pytest.param("extra", {}, marks=requires_3_4),
],
ids=["after", "before", "partition", "partition-regexp", "null-extra", "empty-extra"],
)
def test_rejects_missing_filters_after_rendering(
self, asset, mock_supervisor_comms, field, rendered_value
):
def get_filter(*, context, jinja_env):
return rendered_value
sensor = AssetEventSensor(task_id="s", asset=asset, **{field: get_filter})
context = _get_context(sensor)
sensor.render_template_fields(context)
with pytest.raises(ValueError, match="requires at least one filter"):
sensor.poke(context)
mock_supervisor_comms.send.assert_not_called()
def test_returns_serialized_events(self, asset, events_response, mock_supervisor_comms):
mock_supervisor_comms.send.return_value = events_response
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1))
result = sensor.poke(_get_context(sensor))
assert isinstance(result, PokeReturnValue)
assert bool(result) is True
assert result.xcom_value == [event.model_dump(mode="json") for event in events_response.asset_events]
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(1))
)
@pytest.mark.parametrize(
("count_policy", "expected_count", "actual_count", "done"),
[
("minimum", 1, 0, False),
("minimum", 1, 1, True),
("minimum", 3, 2, False),
("minimum", 3, 3, True),
("minimum", 3, 5, True),
("minimum", 0, 0, True),
("exact", 0, 0, True),
("exact", 0, 1, False),
("exact", 3, 2, False),
("exact", 3, 3, True),
("exact", 3, 5, False),
],
)
def test_count_policy(
self, asset, events_response, mock_supervisor_comms, count_policy, expected_count, actual_count, done
):
events = events_response.asset_events[:actual_count]
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=events)
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
count_policy=count_policy,
expected_count=expected_count,
)
result = sensor.poke(_get_context(sensor))
assert bool(result) is done
assert result.xcom_value == ([event.model_dump(mode="json") for event in events] if done else None)
@requires_3_4
@pytest.mark.parametrize("use_alias", [False, True])
@pytest.mark.parametrize("string_bounds", [False, True])
def test_forwards_all_filters(self, asset, mock_supervisor_comms, use_alias, string_bounds):
mock_supervisor_comms.send.return_value = AssetEventsResult(asset_events=[])
target = AssetAlias(name=ALIAS_NAME) if use_alias else asset
sensor = AssetEventSensor(
task_id="s",
asset=target,
after=_get_timestamp(1).isoformat() if string_bounds else _get_timestamp(1),
before=_get_timestamp(9).isoformat() if string_bounds else _get_timestamp(9),
ascending=False,
limit=7,
partition_key="us|2024-01-01",
partition_key_regexp_pattern="us.*",
extra={"region": "us", "status": "validated"},
)
sensor.poke(_get_context(sensor))
filters = {
"after": _get_timestamp(1),
"before": _get_timestamp(9),
"ascending": False,
"limit": 7,
"partition_key": "us|2024-01-01",
"partition_key_regexp_pattern": "us.*",
"extra": {"region": "us", "status": "validated"},
}
message = (
GetAssetEventByAssetAlias(alias_name=ALIAS_NAME, **filters)
if use_alias
else GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, **filters)
)
mock_supervisor_comms.send.assert_called_once_with(message)
@pytest.mark.parametrize("selector", ["name", "uri"])
def test_resolves_asset_reference(self, selector, events_response, mock_supervisor_comms):
target = {"name": ASSET_NAME} if selector == "name" else {"uri": ASSET_URI}
sensor = AssetEventSensor(task_id="s", asset=Asset.ref(**target), after=_get_timestamp(1))
assert sensor.inlets == [Asset.ref(**target)]
mock_supervisor_comms.send.side_effect = [
AssetResult(name=ASSET_NAME, uri=ASSET_URI, group="asset", extra={}),
events_response,
]
assert bool(sensor.poke(_get_context(sensor))) is True
resolution = GetAssetByName(name=ASSET_NAME) if selector == "name" else GetAssetByUri(uri=ASSET_URI)
query = GetAssetEventByAsset(
name=ASSET_NAME if selector == "name" else None,
uri=ASSET_URI if selector == "uri" else None,
after=_get_timestamp(1),
)
assert mock_supervisor_comms.send.call_args_list == [call(resolution), call(query)]
def test_fetches_fresh_events_on_every_poke(self, asset, events_response, mock_supervisor_comms):
mock_supervisor_comms.send.side_effect = [AssetEventsResult(asset_events=[]), events_response]
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1))
context = _get_context(sensor)
assert bool(sensor.poke(context)) is False
assert bool(sensor.poke(context)) is True
query = GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(1))
assert mock_supervisor_comms.send.call_args_list == [call(query), call(query)]
def test_execute_returns_xcom_payload(self, asset, events_response, mock_supervisor_comms):
mock_supervisor_comms.send.return_value = events_response
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1))
result = sensor.execute(_get_context(sensor))
assert result == [event.model_dump(mode="json") for event in events_response.asset_events]
def test_mapped_task(self, asset, events_response, create_runtime_ti, mock_supervisor_comms):
with DAG("test_dag"):
mapped = AssetEventSensor.partial(task_id="s", asset=asset).expand(
after=[_get_timestamp(1), _get_timestamp(2)]
)
ti = create_runtime_ti(task=mapped, map_index=1)
context = ti.get_template_context()
ti.render_templates(context=context)
mock_supervisor_comms.send.return_value = events_response
assert bool(ti.task.poke(context)) is True
mock_supervisor_comms.send.assert_called_with(
GetAssetEventByAsset(name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(2))
)
class TestProcessResult:
@pytest.mark.parametrize(
("process_result", "expected_count", "expected_ids", "done"),
[
(only_us_partitions, 3, [1, 2, 4], True),
(dedup_by_partition_key, 4, [1, 2, 3, 5], True),
(only_us_partitions, 5, None, False),
(lambda events: [], 1, None, False),
(lambda events: [], 0, [], True),
],
)
def test_processes_events_before_counting(
self,
asset,
events_response,
mock_supervisor_comms,
process_result,
expected_count,
expected_ids,
done,
):
mock_supervisor_comms.send.return_value = events_response
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
process_result=process_result,
expected_count=expected_count,
count_policy="exact",
)
result = sensor.poke(_get_context(sensor))
assert bool(result) is done
if done:
assert [event["id"] for event in result.xcom_value] == expected_ids
else:
assert result.xcom_value is None
def test_processing_can_expand_limited_results(self, asset, events_response, mock_supervisor_comms):
mock_supervisor_comms.send.return_value = AssetEventsResult(
asset_events=events_response.asset_events[:1]
)
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
limit=1,
expected_count=2,
count_policy="exact",
process_result=lambda events: events + events,
)
result = sensor.poke(_get_context(sensor))
assert bool(result) is True
assert [event["id"] for event in result.xcom_value] == [1, 1]
def test_processing_receives_only_limited_results(
self, asset, events_response, mock_supervisor_comms, mocker
):
mock_supervisor_comms.send.return_value = AssetEventsResult(
asset_events=[events_response.asset_events[2]]
)
process_result = mocker.create_autospec(only_us_partitions, side_effect=only_us_partitions)
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(3),
ascending=True,
limit=1,
process_result=process_result,
)
result = sensor.poke(_get_context(sensor))
assert bool(result) is False
assert result.xcom_value is None
process_result.assert_called_once()
assert [event.id for event in process_result.call_args.args[0]] == [3]
mock_supervisor_comms.send.assert_called_once_with(
GetAssetEventByAsset(
name=ASSET_NAME, uri=ASSET_URI, after=_get_timestamp(3), ascending=True, limit=1
)
)
def test_processing_can_return_json_values(self, asset, events_response, mock_supervisor_comms):
mock_supervisor_comms.send.return_value = events_response
sensor = AssetEventSensor(
task_id="s",
asset=asset,
after=_get_timestamp(1),
process_result=lambda events: [{"partition": event.partition_key} for event in events],
)
result = sensor.poke(_get_context(sensor))
assert bool(result) is True
assert result.xcom_value == [
{"partition": event.partition_key} for event in events_response.asset_events
]
def test_exception_propagates(self, asset, events_response, mock_supervisor_comms):
def boom(events):
raise RuntimeError("process_result exploded")
mock_supervisor_comms.send.return_value = events_response
sensor = AssetEventSensor(task_id="s", asset=asset, after=_get_timestamp(1), process_result=boom)
with pytest.raises(RuntimeError, match="process_result exploded"):
sensor.poke(_get_context(sensor))