mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
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:
1 parent
51140809d2
commit
4df59b9a03
6 files changed
+1529
No files matched your search
@@ -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]
|
||||
@@ -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
|
||||
|
||||
+85
@@ -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))
|
||||
Reference in new issue
Block a user