This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3ac1e4f7f2e Make asset event batching opt-in (#68517)
3ac1e4f7f2e is described below
commit 3ac1e4f7f2e72deaf152b524908ad16a9d320813
Author: renat-sagut <[email protected]>
AuthorDate: Mon Oct 12 05:49:29 2026 +0200
Make asset event batching opt-in (#68517)
* batch asset events
* allow apdr reuse
* tests
* tests
* cleanup
* tests
* formatting
* tests
* tests
* tests
* mypy
* docu
* ruff
* added version guard in openlineage tests
* Update airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst
Co-authored-by: Jens Scheffler <[email protected]>
* cleanup
* version guard
* tests
* tests
* tests
* docu
* trigger ci
* version guard
* cleanup
* tests
* Preserve asset event batching choices across timetables
Dag authors need batching choices to survive serialization without changing
existing constructor calls. Align the feature with current scheduler queue
handling and the reviewed release target.
* Make asset event batching opt-in and bound queued run creation
* Preserve queued asset events when batching changes
Existing compound asset schedules must remain valid after reparsing.
Changing batching modes must also preserve queued partition events and let
users inspect the progress of each pending run.
* Remove asset event batching newsfragment
* Fix docs spelling check failure on asset batching docs
The docs spelling check rejects "rollups", which is not in the wordlist, so
the PR documentation build failed. Refer to RollupMapper instead of using the
plural.
* Move merging of pending partitioned runs out of the batching change
Merging pending partitioned runs when batching is enabled is a separate
concern from making batching opt-in, and its scheduler locking is the
riskiest part of the change. Keep this PR to the opt-in default and track
the merge logic separately.
* Improve asset batching and pending run selection
Asset schedules need to express their event requirements consistently
during parsing and scheduling. The pending-run view should show the latest
attempt when no specific run is selected.
* Clarify asset catchup and simplify partition mapper defaults
---------
Co-authored-by: Jens Scheffler <[email protected]>
Co-authored-by: Ephraim Anierobi <[email protected]>
---
.../authoring-and-scheduling/asset-scheduling.rst | 64 ++++++
airflow-core/newsfragments/68517.significant.rst | 14 ++
.../api_fastapi/core_api/openapi/_private_ui.yaml | 8 +
.../core_api/routes/ui/partitioned_dag_runs.py | 11 +-
airflow-core/src/airflow/assets/manager.py | 26 ++-
.../src/airflow/config_templates/config.yml | 15 ++
.../src/airflow/jobs/scheduler_job_runner.py | 84 +++++---
airflow-core/src/airflow/models/asset.py | 10 +-
.../airflow/serialization/definitions/assets.py | 13 ++
airflow-core/src/airflow/serialization/encoders.py | 7 +-
airflow-core/src/airflow/timetables/assets.py | 35 +--
airflow-core/src/airflow/timetables/base.py | 16 ++
airflow-core/src/airflow/timetables/simple.py | 91 +++++---
.../src/airflow/ui/openapi-gen/queries/common.ts | 5 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 6 +-
.../ui/openapi-gen/requests/services.gen.ts | 4 +-
.../airflow/ui/openapi-gen/requests/types.gen.ts | 1 +
.../ui/src/components/AssetProgressCell.test.tsx | 70 ++++++
.../ui/src/components/AssetProgressCell.tsx | 15 +-
.../src/pages/DagsList/PartitionScheduleModal.tsx | 1 +
.../routes/ui/test_partitioned_dag_runs.py | 48 +++--
airflow-core/tests/unit/assets/test_manager.py | 37 ++++
.../tests/unit/dag_processing/test_dagbag.py | 28 +++
airflow-core/tests/unit/jobs/test_scheduler_job.py | 239 ++++++++++++++++++++-
.../unit/serialization/definitions/test_assets.py | 19 ++
.../tests/unit/timetables/test_assets_timetable.py | 144 ++++++++++++-
.../unit/timetables/test_partitioned_timetable.py | 54 ++++-
.../tests/unit/openlineage/plugins/test_utils.py | 39 +++-
.../tests/unit/openlineage/utils/test_utils.py | 33 ++-
task-sdk/docs/api.rst | 2 +
task-sdk/src/airflow/sdk/__init__.py | 3 +
task-sdk/src/airflow/sdk/__init__.pyi | 2 +
.../src/airflow/sdk/definitions/asset/__init__.py | 13 ++
.../airflow/sdk/definitions/timetables/assets.py | 51 ++++-
task-sdk/tests/task_sdk/definitions/test_asset.py | 16 ++
.../task_sdk/definitions/timetables/test_assets.py | 166 ++++++++++++++
39 files changed, 1263 insertions(+), 145 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst
b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst
index 1bf2507237b..03d9122cc9d 100644
--- a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst
+++ b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst
@@ -446,3 +446,67 @@ Asset-aware timetables combine asset expressions with a
time-based schedule:
* Use ``AssetAndTimeSchedule`` to keep a Dag on a timetable but only create
scheduled runs once the referenced assets have been updated.
For more detailed information on asset-aware timetables, refer to
:ref:`AssetOrTimeSchedule <asset-timetable-section>`.
+
+
+Controlling DagRun creation per asset event
+---------------------------------------------
+
+.. versionadded:: 3.4.0
+
+Asset-triggered Dags create one Dag run per event by default. Set
+``batch_asset_events=True`` on the timetable to consume queued events together
+in one run. The :ref:`config:scheduler__batch_asset_events` setting changes the
+default for the deployment; an explicit timetable argument takes precedence.
+Reparse Dags after changing the setting. Previously serialized timetables
without
+this option retain batching until they are reparsed.
+
+The option is also available on ``AssetOrTimeSchedule`` and
+``PartitionedAssetTimetable``. Asset conditions using ``&`` and partition
mappers based on
+``RollupMapper`` need several events to satisfy their condition and cannot
consume each event
+independently, so they default to batching when ``batch_asset_events`` is not
set.
+Passing ``batch_asset_events=False`` explicitly for them makes timetable
validation
+raise an error. Single-asset schedules and ``|`` expressions support either
mode.
+
+Pending partitioned runs created while batching was disabled are not merged
when you
+enable batching, so let them finish before changing the setting.
+
+For non-partitioned Dags, per-event run creation respects ``max_active_runs``
and the scheduler's
+:ref:`config:scheduler__max_dagruns_to_create_per_loop` budget. Events left
over
+remain queued for subsequent scheduler passes.
+
+For non-partitioned asset-triggered Dags, ``catchup=False`` consumes only
events
+queued for that Dag. It does not discard a queued backlog: with batching
disabled,
+each queued event still produces its own run.
+
+With ``catchup=True``, once queued events satisfy the asset condition, the
scheduler
+also includes older, unconsumed events for the Dag's scheduled assets and
aliases,
+including events emitted before the Dag started consuming them. Events already
+consumed by a run of that Dag are excluded in both modes.
``batch_asset_events``
+controls whether the eligible events produce separate runs or are consumed
together.
+
+.. code-block:: python
+
+ from airflow.sdk import DAG, Asset, AssetTriggeredTimetable
+
+ # Each update to "data-file" produces its own DagRun
+ with DAG(
+ dag_id="per-event-consumer",
+ schedule=AssetTriggeredTimetable(
+ assets=Asset("s3://bucket/data-file"),
+ batch_asset_events=False,
+ ),
+ ):
+ ...
+
+To consume several queued updates to an asset in one run, enable batching:
+
+.. code-block:: python
+
+ with DAG(
+ dag_id="combined-consumer",
+ schedule=AssetTriggeredTimetable(
+ assets=Asset("s3://bucket/orders"),
+ batch_asset_events=True,
+ ),
+ ):
+ ...
diff --git a/airflow-core/newsfragments/68517.significant.rst
b/airflow-core/newsfragments/68517.significant.rst
new file mode 100644
index 00000000000..f07598c3625
--- /dev/null
+++ b/airflow-core/newsfragments/68517.significant.rst
@@ -0,0 +1,14 @@
+Asset-triggered Dags create one run per event by default
+
+Single-asset schedules and asset OR expressions now create one Dag run per
event
+instead of consuming all queued events in one run. This also applies to
+``AssetOrTimeSchedule`` and ``PartitionedAssetTimetable``.
+
+To retain batching, set ``batch_asset_events=True`` on the timetable or enable
+``[scheduler] batch_asset_events`` and reparse the affected Dags. Already
serialized
+timetables without this option retain their previous batching behavior until
reparsed.
+
+Schedules requiring several events, including ``schedule=[a, b]``, asset AND
+expressions, and partitioned timetables using ``RollupMapper``, continue to
batch
+by default. Explicitly setting ``batch_asset_events=False`` for these schedules
+raises a timetable validation error.
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index 2feaf0a7ea2..3b988ffb102 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -560,6 +560,14 @@ paths:
schema:
type: string
title: Partition Key
+ - name: partitioned_dag_run_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Partitioned Dag Run Id
responses:
'200':
description: Successful Response
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
index 77ac9454b6a..f87e4cd12b2 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py
@@ -370,12 +370,13 @@ def get_pending_partitioned_dag_run(
dag_id: str,
partition_key: str,
session: SessionDep,
+ partitioned_dag_run_id: int | None = None,
) -> PartitionedDagRunDetailResponse:
"""Return full details for pending PartitionedDagRun."""
# partition_key is a query param, not a path segment: it is a free-form key
# (up to 250 chars) that may itself contain "/", which would otherwise be
# ambiguous (or mis-routed) as a path segment.
- partitioned_dag_run = session.execute(
+ query = (
select(
AssetPartitionDagRun.id,
AssetPartitionDagRun.target_dag_id,
@@ -390,11 +391,11 @@ def get_pending_partitioned_dag_run(
AssetPartitionDagRun.partition_key == partition_key,
AssetPartitionDagRun.created_dag_run_id.is_(None),
)
- # Duplicate pending rows for the same (dag_id, partition_key) can exist
- # after a crash; mirror _get_or_create_apdr and work on the latest one.
.order_by(AssetPartitionDagRun.id.desc())
- .limit(1)
- ).first()
+ )
+ if partitioned_dag_run_id is not None:
+ query = query.where(AssetPartitionDagRun.id == partitioned_dag_run_id)
+ partitioned_dag_run = session.execute(query.limit(1)).first()
if partitioned_dag_run is None:
raise HTTPException(
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index 8aa834d4a4a..ad133881db7 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -694,6 +694,7 @@ class AssetManager(LoggingMixin):
target_dag=target_dag,
rollup_fingerprint=fingerprint,
asset_id=asset_id,
+ allow_reuse=timetable.batch_asset_events,
session=session,
)
log_record = PartitionedAssetKeyLog(
@@ -715,6 +716,7 @@ class AssetManager(LoggingMixin):
target_dag: DagModel,
rollup_fingerprint: dict,
asset_id: int,
+ allow_reuse: bool,
session: Session,
) -> AssetPartitionDagRun:
"""
@@ -743,17 +745,25 @@ class AssetManager(LoggingMixin):
one, the producing assets disagree; picking one would be
order-dependent, so the
carried date is suppressed to ``None`` (and re-adoptable by a later
event).
- Otherwise (the dates agree, or this event carries none) the existing
value is kept.
+
+ When ``allow_reuse=True``, an existing pending APDR for the same
+ ``(target_dag, partition_key)`` is reused — multiple events accumulate
on one
+ APDR. When ``allow_reuse=False`` (set when the timetable's
``batch_asset_events``
+ is ``False``), a new APDR is always created so each event gets its own
APDR
+ and the scheduler produces one DagRun per event.
"""
with _lock_asset_model(session=session, asset_id=asset_id):
- latest_apdr: AssetPartitionDagRun | None = session.scalar(
- select(AssetPartitionDagRun)
- .where(
- AssetPartitionDagRun.partition_key == target_key,
- AssetPartitionDagRun.target_dag_id == target_dag.dag_id,
+ latest_apdr: AssetPartitionDagRun | None = None
+ if allow_reuse:
+ latest_apdr = session.scalar(
+ select(AssetPartitionDagRun)
+ .where(
+ AssetPartitionDagRun.partition_key == target_key,
+ AssetPartitionDagRun.target_dag_id ==
target_dag.dag_id,
+ )
+ .order_by(AssetPartitionDagRun.id.desc())
+ .limit(1)
)
- .order_by(AssetPartitionDagRun.id.desc())
- .limit(1)
- )
if latest_apdr and latest_apdr.created_dag_run_id is None:
existing_partition_date = latest_apdr.partition_date
if existing_partition_date is None:
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index 0de787a4c59..951b85194aa 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -2688,6 +2688,19 @@ sentry:
scheduler:
description: ~
options:
+ batch_asset_events:
+ description: |
+ Default for the ``batch_asset_events`` argument of asset-triggered
timetables.
+ When false, each asset event creates a separate Dag run. Set this to
true
+ to restore batching of queued events into one run. An explicit
timetable
+ argument overrides this setting. Asset AND conditions and
``RollupMapper`` partition
+ mappers always default to batching, and passing
``batch_asset_events=False`` for them
+ is rejected. The setting is read when a Dag's timetable is constructed;
+ reparse Dags after changing it.
+ version_added: 3.4.0
+ type: boolean
+ example: ~
+ default: "False"
job_heartbeat_sec:
description: |
Task instances listen for external kill signal (when you clear tasks
@@ -2886,6 +2899,8 @@ scheduler:
max_dagruns_to_create_per_loop:
description: |
Max number of DAGs to create DagRuns for per scheduler loop.
+ Also limits non-partitioned asset-triggered Dag run creation across
all selected
+ Dags, including those with event batching enabled.
example: ~
version_added: 2.0.0
type: integer
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 052d7d12b1a..c78bbf1d4bb 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -2901,6 +2901,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
*,
dag: SerializedDAG,
records: Sequence[AssetDagRunQueue],
+ limit: int | None = None,
session: Session,
) -> list[AssetEvent]:
"""Select unconsumed events referenced by a Dag's locked ADRQ rows."""
@@ -2934,6 +2935,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
),
)
.order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
+ .limit(limit)
)
)
@@ -2966,7 +2968,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
session: Session,
) -> None:
"""For Dags that are triggered by assets, create Dag runs."""
+ remaining_runs = DagModel.NUM_DAGS_PER_DAGRUN_QUERY
+ per_dag_limit = max(1, remaining_runs // max(1, len(dag_models)))
+ active_runs_by_dag: dict[str, int] | None = None
for dag_model in dag_models:
+ if remaining_runs <= 0:
+ break
dag = self._get_current_dag(dag_id=dag_model.dag_id,
session=session)
if not dag:
self.log.error("Dag '%s' not found in serialized_dag table",
dag_model.dag_id)
@@ -2990,11 +2997,31 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
continue
+ run_limit = per_dag_limit
+ if not dag.timetable.batch_asset_events:
+ if active_runs_by_dag is None:
+ active_runs_by_dag = DagRun.active_runs_of_dags(
+ dag_ids=(model.dag_id for model in dag_models),
+ exclude_backfill=True,
+ session=session,
+ )
+ run_limit = min(
+ run_limit,
+ remaining_runs,
+ (dag_model.max_active_runs or 0) -
active_runs_by_dag.get(dag.dag_id, 0),
+ )
+ if run_limit <= 0:
+ continue
+
asset_events = self._select_consumed_asset_events(
dag=dag,
records=queued_adrqs,
+ limit=None if dag.timetable.batch_asset_events else run_limit
+ 1,
session=session,
)
+ has_more_events = not dag.timetable.batch_asset_events and
len(asset_events) > run_limit
+ if has_more_events:
+ asset_events = asset_events[:run_limit]
if asset_events:
triggered_date = timezone.coerce_datetime(max(event.timestamp
for event in asset_events))
self.log.debug(
@@ -3003,41 +3030,50 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
len(queued_adrqs),
triggered_date,
)
- dag_run = dag.create_dagrun(
- run_id=DagRun.generate_run_id(
- run_type=DagRunType.ASSET_TRIGGERED,
logical_date=None, run_after=triggered_date
- ),
- logical_date=None,
- data_interval=None,
- run_after=triggered_date,
- run_type=DagRunType.ASSET_TRIGGERED,
- triggered_by=DagRunTriggeredByType.ASSET,
- state=DagRunState.QUEUED,
- creating_job_id=self.job.id,
- session=session,
- )
team_name = (
self._get_team_names_for_dag_ids([dag.dag_id],
session).get(dag.dag_id)
if self._multi_team
else None
)
- stats.incr("asset.triggered_dagruns",
tags=prune_dict({"team_name": team_name}))
- dag_run.consumed_asset_events.extend(asset_events)
- self.log.info(
- "Created asset-triggered DagRun for '%s': run_id=%s,
consumed %d asset events",
- dag.dag_id,
- dag_run.run_id,
- len(asset_events),
- )
+ for run_after, events in
dag.timetable.group_asset_events(asset_events, triggered_date):
+ dag_run = dag.create_dagrun(
+ run_id=DagRun.generate_run_id(
+ run_type=DagRunType.ASSET_TRIGGERED,
logical_date=None, run_after=run_after
+ ),
+ logical_date=None,
+ data_interval=None,
+ run_after=run_after,
+ run_type=DagRunType.ASSET_TRIGGERED,
+ triggered_by=DagRunTriggeredByType.ASSET,
+ state=DagRunState.QUEUED,
+ creating_job_id=self.job.id,
+ session=session,
+ )
+ remaining_runs -= 1
+ stats.incr("asset.triggered_dagruns",
tags=prune_dict({"team_name": team_name}))
+ dag_run.consumed_asset_events.extend(events)
+ self.log.info(
+ "Created asset-triggered DagRun for '%s': run_id=%s,
consumed %d asset events",
+ dag.dag_id,
+ dag_run.run_id,
+ len(events),
+ )
else:
self.log.info(
"No DagRun created for '%s' - asset events already
consumed or none found",
dag.dag_id,
)
- # Always delete ADRQ rows for this batch to prevent stale entries
accumulating,
- # including when all events were already consumed by a concurrent
DagRun.
+ records_to_delete = queued_adrqs
+ if has_more_events:
+ consumed_ids = {event.id for event in asset_events}
+ if dag.catchup:
+ # Historical events may have no ADRQ. Keep a wakeup row
until that backlog drains.
+ consumed_ids.discard(max(record.asset_event_id for record
in queued_adrqs))
+ records_to_delete = [
+ record for record in queued_adrqs if record.asset_event_id
in consumed_ids
+ ]
self._delete_consumed_asset_records(
- records=queued_adrqs,
+ records=records_to_delete,
dag_id=dag.dag_id,
session=session,
)
diff --git a/airflow-core/src/airflow/models/asset.py
b/airflow-core/src/airflow/models/asset.py
index afc71a59ad6..ab780be2e77 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -949,11 +949,11 @@ class AssetPartitionDagRun(Base):
target_dag_id / partition_key are expected in general; each dag run that
gets created leaves its APDR record behind.
- Where created_dag_run_id is null, the dag run has not yet been created.
- We should not allow more than one row with the same target_dag_id /
- partition_key where created_dag_run_id is null, and this is what the
- `_lock_asset_model` mutex control is for. In case a duplicate somehow
- gets created, we always work on the latest matching APDR record.
+ Where created_dag_run_id is null, the Dag run has not yet been created.
+ With event batching enabled, `_lock_asset_model` protects reuse of the
+ latest pending row for a target_dag_id / partition_key pair. With batching
+ disabled, each event creates its own row, so multiple pending rows for
+ the same pair are expected.
"""
id: Mapped[int] = mapped_column(Integer, primary_key=True,
autoincrement=True)
diff --git a/airflow-core/src/airflow/serialization/definitions/assets.py
b/airflow-core/src/airflow/serialization/definitions/assets.py
index b1f28909373..0ba7137e998 100644
--- a/airflow-core/src/airflow/serialization/definitions/assets.py
+++ b/airflow-core/src/airflow/serialization/definitions/assets.py
@@ -70,6 +70,11 @@ class SerializedAssetBase:
:meta private:
"""
+ @property
+ def requires_batching(self) -> bool:
+ """Whether this condition needs multiple asset events to be
satisfied."""
+ return False
+
def __bool__(self) -> bool:
return True
@@ -260,6 +265,10 @@ class SerializedAssetBooleanCondition(SerializedAssetBase):
agg_func: ClassVar[Callable[[Iterable], bool]]
+ @property
+ def requires_batching(self) -> bool:
+ return any(obj.requires_batching for obj in self.objects)
+
def iter_assets(self) -> Iterator[tuple[SerializedAssetUniqueKey,
SerializedAsset]]:
for o in self.objects:
yield from o.iter_assets()
@@ -301,6 +310,10 @@ class SerializedAssetAll(SerializedAssetBooleanCondition):
agg_func = all
+ @property
+ def requires_batching(self) -> bool:
+ return len(self.objects) > 1 or super().requires_batching
+
def __repr__(self) -> str:
return f"AssetAny({', '.join(map(str, self.objects))})"
diff --git a/airflow-core/src/airflow/serialization/encoders.py
b/airflow-core/src/airflow/serialization/encoders.py
index e4b417de213..d1f801d0b53 100644
--- a/airflow-core/src/airflow/serialization/encoders.py
+++ b/airflow-core/src/airflow/serialization/encoders.py
@@ -371,7 +371,10 @@ class _Serializer:
@serialize_timetable.register
def _(self, timetable: AssetTriggeredTimetable) -> dict[str, Any]:
- return {"asset_condition":
encode_asset_like(timetable.asset_condition)}
+ return {
+ "asset_condition": encode_asset_like(timetable.asset_condition),
+ "batch_asset_events": timetable.batch_asset_events,
+ }
@serialize_timetable.register
def _(self, timetable: EventsTimetable) -> dict[str, Any]:
@@ -439,6 +442,7 @@ class _Serializer:
return {
"asset_condition": encode_asset_like(timetable.asset_condition),
"timetable": encode_timetable(timetable.timetable),
+ "batch_asset_events": timetable.batch_asset_events,
}
@serialize_timetable.register
@@ -449,6 +453,7 @@ class _Serializer:
def _(self, timetable: PartitionedAssetTimetable) -> dict[str, Any]:
return {
"asset_condition": encode_asset_like(timetable.asset_condition),
+ "batch_asset_events": timetable.batch_asset_events,
"default_partition_mapper":
encode_partition_mapper(timetable.default_partition_mapper),
"partition_mapper_config": [
(encode_asset_like(asset),
encode_partition_mapper(partition_mapper))
diff --git a/airflow-core/src/airflow/timetables/assets.py
b/airflow-core/src/airflow/timetables/assets.py
index 7359ba23d0e..8ec26434f7c 100644
--- a/airflow-core/src/airflow/timetables/assets.py
+++ b/airflow-core/src/airflow/timetables/assets.py
@@ -20,10 +20,16 @@ from __future__ import annotations
import typing
from collections.abc import Collection
+import attrs
+
from airflow.exceptions import AirflowTimetableInvalid
from airflow.serialization.definitions.assets import SerializedAsset,
SerializedAssetAll, SerializedAssetBase
from airflow.timetables.base import Timetable
-from airflow.timetables.simple import AssetTriggeredTimetable
+from airflow.timetables.simple import (
+ AssetTriggeredTimetable,
+ _coerce_asset_condition,
+ _get_default_batch_asset_events,
+)
from airflow.utils.types import DagRunType
if typing.TYPE_CHECKING:
@@ -39,21 +45,21 @@ def _validate_asset_time_schedule(*, timetable: Timetable,
asset_condition: Seri
raise AirflowTimetableInvalid("All elements in 'assets' must be
assets")
[email protected](eq=False, repr=False, kw_only=True, slots=False)
class AssetOrTimeSchedule(AssetTriggeredTimetable):
"""Combine time-based scheduling with event-based scheduling."""
- def __init__(
- self,
- *,
- timetable: Timetable,
- assets: Collection[SerializedAsset] | SerializedAssetBase,
- ) -> None:
- super().__init__(assets)
- self.timetable = timetable
- self.description = f"Triggered by assets or {timetable.description}"
- self.periodic = timetable.periodic
- self.can_be_scheduled = timetable.can_be_scheduled
- self.active_runs_limit = timetable.active_runs_limit
+ asset_condition: SerializedAssetBase = attrs.field(alias="assets",
converter=_coerce_asset_condition)
+ timetable: Timetable
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events, takes_self=True)
+ )
+
+ def __attrs_post_init__(self) -> None:
+ self.description = f"Triggered by assets or
{self.timetable.description}"
+ self.periodic = self.timetable.periodic
+ self.can_be_scheduled = self.timetable.can_be_scheduled
+ self.active_runs_limit = self.timetable.active_runs_limit
@classmethod
def deserialize(cls, data: dict[str, typing.Any]) -> Timetable:
@@ -62,9 +68,11 @@ class AssetOrTimeSchedule(AssetTriggeredTimetable):
return cls(
assets=decode_asset_like(data["asset_condition"]),
timetable=decode_timetable(data["timetable"]),
+ batch_asset_events=data.get("batch_asset_events", True),
)
def validate(self) -> None:
+ super().validate()
_validate_asset_time_schedule(
timetable=self.timetable,
asset_condition=self.asset_condition,
@@ -76,6 +84,7 @@ class AssetOrTimeSchedule(AssetTriggeredTimetable):
return {
"asset_condition": encode_asset_like(self.asset_condition),
"timetable": encode_timetable(self.timetable),
+ "batch_asset_events": self.batch_asset_events,
}
@property
diff --git a/airflow-core/src/airflow/timetables/base.py
b/airflow-core/src/airflow/timetables/base.py
index b2ad57cae76..a2325395a08 100644
--- a/airflow-core/src/airflow/timetables/base.py
+++ b/airflow-core/src/airflow/timetables/base.py
@@ -28,6 +28,7 @@ if TYPE_CHECKING:
from pendulum import DateTime
+ from airflow.models.asset import AssetEvent
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun
from airflow.partition_mappers import PartitionMapper
@@ -228,6 +229,21 @@ class Timetable(Protocol):
asset_condition: SerializedAssetBase = _NullAsset()
"""The asset condition that triggers a DAG using this timetable."""
+ batch_asset_events: bool = False
+ """Whether queued asset events are consumed together in one Dag run."""
+
+ def group_asset_events(
+ self, events: list[AssetEvent], triggered_date: DateTime
+ ) -> Iterable[tuple[DateTime, list[AssetEvent]]]:
+ """Yield a run time and its consumed events for each asset-triggered
run."""
+ if not events:
+ return
+ if self.batch_asset_events:
+ yield triggered_date, events
+ else:
+ for event in events:
+ yield timezone.coerce_datetime(event.timestamp), [event]
+
partitioned: bool = False
"""Whether this timetable considers asset partitions.
diff --git a/airflow-core/src/airflow/timetables/simple.py
b/airflow-core/src/airflow/timetables/simple.py
index a1c64a213b0..d16627a56c8 100644
--- a/airflow-core/src/airflow/timetables/simple.py
+++ b/airflow-core/src/airflow/timetables/simple.py
@@ -20,10 +20,12 @@ from contextlib import suppress
from datetime import datetime
from typing import TYPE_CHECKING, Any, TypeAlias
+import attrs
import structlog
from airflow._shared.timezones import timezone
-from airflow.exceptions import InvalidPartitionKeyError
+from airflow.configuration import conf
+from airflow.exceptions import AirflowTimetableInvalid,
InvalidPartitionKeyError
from airflow.partition_mappers.identity import IdentityMapper
from airflow.serialization.definitions.assets import (
SerializedAsset,
@@ -207,6 +209,22 @@ class PartitionedAtRuntime(NullTimetable):
return "PartitionedAtRuntime"
+def _coerce_asset_condition(
+ assets: Collection[SerializedAsset | BaseAsset] | SerializedAssetBase |
BaseAsset,
+) -> SerializedAssetBase:
+ if isinstance(assets, SerializedAssetBase | BaseAsset):
+ return ensure_serialized_asset(assets)
+ return SerializedAssetAll([ensure_serialized_asset(asset) for asset in
assets])
+
+
+def _get_default_batch_asset_events(timetable: AssetTriggeredTimetable) ->
bool:
+ return (
+ conf.getboolean("scheduler", "batch_asset_events", fallback=False)
+ or timetable.get_batching_requirement() is not None
+ )
+
+
[email protected](eq=False, repr=False, slots=False)
class AssetTriggeredTimetable(_TrivialTimetable):
"""
Timetable that never schedules anything.
@@ -216,29 +234,42 @@ class AssetTriggeredTimetable(_TrivialTimetable):
:meta private:
"""
- description: str = "Triggered by assets"
+ description = "Triggered by assets"
asset_triggered = True
- def __init__(self, assets: Collection[SerializedAsset] |
SerializedAssetBase) -> None:
- super().__init__()
- # Compatibility: Handle SDK assets if needed so this class works in
dag files.
- if isinstance(assets, SerializedAssetBase | BaseAsset):
- self.asset_condition = ensure_serialized_asset(assets)
- else:
- self.asset_condition =
SerializedAssetAll([ensure_serialized_asset(a) for a in assets])
+ asset_condition: SerializedAssetBase = attrs.field(alias="assets",
converter=_coerce_asset_condition)
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events,
takes_self=True), kw_only=True
+ )
@classmethod
def deserialize(cls, data: dict[str, Any]) -> Timetable:
from airflow.serialization.decoders import decode_asset_like
- return cls(decode_asset_like(data["asset_condition"]))
+ return cls(
+ decode_asset_like(data["asset_condition"]),
+ batch_asset_events=data.get("batch_asset_events", True),
+ )
@property
def summary(self) -> str:
return "Asset"
def serialize(self) -> dict[str, Any]:
- return {"asset_condition": encode_asset_like(self.asset_condition)}
+ return {
+ "asset_condition": encode_asset_like(self.asset_condition),
+ "batch_asset_events": self.batch_asset_events,
+ }
+
+ def get_batching_requirement(self) -> str | None:
+ """Return why this timetable cannot run without event batching, or
``None``."""
+ if self.asset_condition.requires_batching:
+ return "Asset AND conditions require batch_asset_events=True"
+ return None
+
+ def validate(self) -> None:
+ if not self.batch_asset_events and (reason :=
self.get_batching_requirement()):
+ raise AirflowTimetableInvalid(reason)
def generate_run_id(
self,
@@ -273,6 +304,7 @@ class AssetTriggeredTimetable(_TrivialTimetable):
DEFAULT_PARTITION_MAPPER = IdentityMapper()
[email protected](eq=False, repr=False, kw_only=True, slots=False)
class PartitionedAssetTimetable(AssetTriggeredTimetable):
"""Asset-driven timetable that listens for partitioned assets."""
@@ -282,21 +314,30 @@ class PartitionedAssetTimetable(AssetTriggeredTimetable):
def summary(self) -> str:
return "Partitioned Asset"
- def __init__(
- self,
- *,
- assets: SerializedAssetBase,
- partition_mapper_config: dict[SerializedAssetBase, PartitionMapper] |
None = None,
- default_partition_mapper: PartitionMapper = DEFAULT_PARTITION_MAPPER,
- ) -> None:
- super().__init__(assets=assets)
- self.partition_mapper_config = partition_mapper_config or {}
- self.default_partition_mapper = default_partition_mapper
-
- self._name_to_partition_mapper: dict[str, PartitionMapper] = {}
- self._uri_to_partition_mapper: dict[str, PartitionMapper] = {}
+ asset_condition: SerializedAssetBase = attrs.field(alias="assets",
converter=_coerce_asset_condition)
+ partition_mapper_config: dict[SerializedAssetBase, PartitionMapper] =
attrs.field(
+ factory=dict,
+ # Mypy's attrs plugin does not recognize converter factories.
+ converter=attrs.converters.default_if_none(factory=dict), # type:
ignore[misc]
+ )
+ default_partition_mapper: PartitionMapper = DEFAULT_PARTITION_MAPPER
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events, takes_self=True)
+ )
+ _name_to_partition_mapper: dict[str, PartitionMapper] =
attrs.field(factory=dict, init=False)
+ _uri_to_partition_mapper: dict[str, PartitionMapper] =
attrs.field(factory=dict, init=False)
+
+ def __attrs_post_init__(self) -> None:
self._build_name_uri_mapping()
+ def get_batching_requirement(self) -> str | None:
+ if any(
+ mapper.is_rollup
+ for mapper in (self.default_partition_mapper,
*self.partition_mapper_config.values())
+ ):
+ return "Partition rollups require batch_asset_events=True"
+ return super().get_batching_requirement()
+
def _build_name_uri_mapping(self) -> None:
for base_asset, partition_mapper in
self.partition_mapper_config.items():
for unique_key, _ in base_asset.iter_assets():
@@ -403,6 +444,7 @@ class PartitionedAssetTimetable(AssetTriggeredTimetable):
return {
"asset_condition": encode_asset_like(self.asset_condition),
+ "batch_asset_events": self.batch_asset_events,
"partition_mapper_config": [
(encode_asset_like(asset),
encode_partition_mapper(partition_mapper))
for asset, partition_mapper in
self.partition_mapper_config.items()
@@ -420,6 +462,7 @@ class PartitionedAssetTimetable(AssetTriggeredTimetable):
timetable = cls(
assets=decode_asset_like(data["asset_condition"]),
+ batch_asset_events=data.get("batch_asset_events", True),
default_partition_mapper=decode_partition_mapper(default_partition_mapper_data),
partition_mapper_config={
decode_asset_like(ser_asest):
decode_partition_mapper(ser_partition_mapper)
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
index bcebe9e0d7b..b4aed6744dd 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -1106,10 +1106,11 @@ export const
UsePartitionedDagRunServiceGetPartitionedDagRunsKeyFn = ({ dagId, h
export type PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse
= Awaited<ReturnType<typeof
PartitionedDagRunService.getPendingPartitionedDagRun>>;
export type
PartitionedDagRunServiceGetPendingPartitionedDagRunQueryResult<TData =
PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse, TError =
unknown> = UseQueryResult<TData, TError>;
export const usePartitionedDagRunServiceGetPendingPartitionedDagRunKey =
"PartitionedDagRunServiceGetPendingPartitionedDagRun";
-export const UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn = ({
dagId, partitionKey }: {
+export const UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn = ({
dagId, partitionedDagRunId, partitionKey }: {
dagId: string;
+ partitionedDagRunId?: number;
partitionKey: string;
-}, queryKey?: Array<unknown>) =>
[usePartitionedDagRunServiceGetPendingPartitionedDagRunKey, ...(queryKey ?? [{
dagId, partitionKey }])];
+}, queryKey?: Array<unknown>) =>
[usePartitionedDagRunServiceGetPendingPartitionedDagRunKey, ...(queryKey ?? [{
dagId, partitionedDagRunId, partitionKey }])];
export type DependenciesServiceGetDependenciesDefaultResponse =
Awaited<ReturnType<typeof DependenciesService.getDependencies>>;
export type DependenciesServiceGetDependenciesQueryResult<TData =
DependenciesServiceGetDependenciesDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
export const useDependenciesServiceGetDependenciesKey =
"DependenciesServiceGetDependencies";
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index b3d106c6f04..718f1947310 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -2192,13 +2192,15 @@ export const
ensureUsePartitionedDagRunServiceGetPartitionedDagRunsData = (query
* @param data The data for the request.
* @param data.dagId
* @param data.partitionKey
+* @param data.partitionedDagRunId
* @returns PartitionedDagRunDetailResponse Successful Response
* @throws ApiError
*/
-export const ensureUsePartitionedDagRunServiceGetPendingPartitionedDagRunData
= (queryClient: QueryClient, { dagId, partitionKey }: {
+export const ensureUsePartitionedDagRunServiceGetPendingPartitionedDagRunData
= (queryClient: QueryClient, { dagId, partitionedDagRunId, partitionKey }: {
dagId: string;
+ partitionedDagRunId?: number;
partitionKey: string;
-}) => queryClient.ensureQueryData({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionKey }), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey })
});
+}) => queryClient.ensureQueryData({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionedDagRunId, partitionKey }), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId,
partitionedDagRunId, partitionKey }) });
/**
* Get Dependencies
* Dependencies graph.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index 485654a7a40..bcd34c714af 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -2192,13 +2192,15 @@ export const
prefetchUsePartitionedDagRunServiceGetPartitionedDagRuns = (queryCl
* @param data The data for the request.
* @param data.dagId
* @param data.partitionKey
+* @param data.partitionedDagRunId
* @returns PartitionedDagRunDetailResponse Successful Response
* @throws ApiError
*/
-export const prefetchUsePartitionedDagRunServiceGetPendingPartitionedDagRun =
(queryClient: QueryClient, { dagId, partitionKey }: {
+export const prefetchUsePartitionedDagRunServiceGetPendingPartitionedDagRun =
(queryClient: QueryClient, { dagId, partitionedDagRunId, partitionKey }: {
dagId: string;
+ partitionedDagRunId?: number;
partitionKey: string;
-}) => queryClient.prefetchQuery({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionKey }), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey })
});
+}) => queryClient.prefetchQuery({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionedDagRunId, partitionKey }), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId,
partitionedDagRunId, partitionKey }) });
/**
* Get Dependencies
* Dependencies graph.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index a225a608ac7..11b7e9f5851 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -2192,13 +2192,15 @@ export const
usePartitionedDagRunServiceGetPartitionedDagRuns = <TData = Common.
* @param data The data for the request.
* @param data.dagId
* @param data.partitionKey
+* @param data.partitionedDagRunId
* @returns PartitionedDagRunDetailResponse Successful Response
* @throws ApiError
*/
-export const usePartitionedDagRunServiceGetPendingPartitionedDagRun = <TData =
Common.PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse,
TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
partitionKey }: {
+export const usePartitionedDagRunServiceGetPendingPartitionedDagRun = <TData =
Common.PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse,
TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
partitionedDagRunId, partitionKey }: {
dagId: string;
+ partitionedDagRunId?: number;
partitionKey: string;
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionKey }, queryKey), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey })
as TData, ...options });
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionedDagRunId, partitionKey }, queryKey), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId,
partitionedDagRunId, partitionKey }) as TData, ...options });
/**
* Get Dependencies
* Dependencies graph.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index d77094ceaed..84c37145799 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -2192,13 +2192,15 @@ export const
usePartitionedDagRunServiceGetPartitionedDagRunsSuspense = <TData =
* @param data The data for the request.
* @param data.dagId
* @param data.partitionKey
+* @param data.partitionedDagRunId
* @returns PartitionedDagRunDetailResponse Successful Response
* @throws ApiError
*/
-export const usePartitionedDagRunServiceGetPendingPartitionedDagRunSuspense =
<TData =
Common.PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse,
TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
partitionKey }: {
+export const usePartitionedDagRunServiceGetPendingPartitionedDagRunSuspense =
<TData =
Common.PartitionedDagRunServiceGetPendingPartitionedDagRunDefaultResponse,
TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagId,
partitionedDagRunId, partitionKey }: {
dagId: string;
+ partitionedDagRunId?: number;
partitionKey: string;
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionKey }, queryKey), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId, partitionKey })
as TData, ...options });
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UsePartitionedDagRunServiceGetPendingPartitionedDagRunKeyFn({ dagId,
partitionedDagRunId, partitionKey }, queryKey), queryFn: () =>
PartitionedDagRunService.getPendingPartitionedDagRun({ dagId,
partitionedDagRunId, partitionKey }) as TData, ...options });
/**
* Get Dependencies
* Dependencies graph.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index c22d64e684d..4b799d58b9d 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -5380,6 +5380,7 @@ export class PartitionedDagRunService {
* @param data The data for the request.
* @param data.dagId
* @param data.partitionKey
+ * @param data.partitionedDagRunId
* @returns PartitionedDagRunDetailResponse Successful Response
* @throws ApiError
*/
@@ -5391,7 +5392,8 @@ export class PartitionedDagRunService {
dag_id: data.dagId
},
query: {
- partition_key: data.partitionKey
+ partition_key: data.partitionKey,
+ partitioned_dag_run_id: data.partitionedDagRunId
},
errors: {
404: 'Not Found',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index bc19ad00938..69a3ad485c6 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -5404,6 +5404,7 @@ export type GetPartitionedDagRunsResponse =
PartitionedDagRunCollectionResponse;
export type GetPendingPartitionedDagRunData = {
dagId: string;
+ partitionedDagRunId?: number | null;
partitionKey: string;
};
diff --git
a/airflow-core/src/airflow/ui/src/components/AssetProgressCell.test.tsx
b/airflow-core/src/airflow/ui/src/components/AssetProgressCell.test.tsx
new file mode 100644
index 00000000000..291f3198a9b
--- /dev/null
+++ b/airflow-core/src/airflow/ui/src/components/AssetProgressCell.test.tsx
@@ -0,0 +1,70 @@
+/*!
+ * 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.
+ */
+import { render, waitFor } from "@testing-library/react";
+import { expect, it, vi } from "vitest";
+
+import { PartitionedDagRunService } from "openapi/requests/services.gen";
+
+import { Wrapper } from "src/utils/Wrapper";
+
+import { AssetProgressCell } from "./AssetProgressCell";
+
+it("fetches separate details for pending runs sharing a partition key", async
() => {
+ const request = vi.spyOn(PartitionedDagRunService,
"getPendingPartitionedDagRun").mockResolvedValue({
+ assets: [],
+ dag_id: "consumer",
+ id: 1,
+ partition_key: "same-key",
+ total_received: 1,
+ total_required: 1,
+ });
+
+ render(
+ <>
+ <AssetProgressCell
+ dagId="consumer"
+ partitionedDagRunId={1}
+ partitionKey="same-key"
+ totalReceived={1}
+ totalRequired={1}
+ />
+ <AssetProgressCell
+ dagId="consumer"
+ partitionedDagRunId={2}
+ partitionKey="same-key"
+ totalReceived={1}
+ totalRequired={1}
+ />
+ </>,
+ { wrapper: Wrapper },
+ );
+
+ await waitFor(() => {
+ expect(request).toHaveBeenCalledWith({
+ dagId: "consumer",
+ partitionedDagRunId: 1,
+ partitionKey: "same-key",
+ });
+ expect(request).toHaveBeenCalledWith({
+ dagId: "consumer",
+ partitionedDagRunId: 2,
+ partitionKey: "same-key",
+ });
+ });
+});
diff --git a/airflow-core/src/airflow/ui/src/components/AssetProgressCell.tsx
b/airflow-core/src/airflow/ui/src/components/AssetProgressCell.tsx
index bf62e40b569..799499b6bcf 100644
--- a/airflow-core/src/airflow/ui/src/components/AssetProgressCell.tsx
+++ b/airflow-core/src/airflow/ui/src/components/AssetProgressCell.tsx
@@ -32,14 +32,25 @@ import { formatNumber } from "src/utils";
type Props = {
readonly dagId: string;
+ readonly partitionedDagRunId: number;
readonly partitionKey: string;
readonly totalReceived: number;
readonly totalRequired: number;
};
-export const AssetProgressCell = ({ dagId, partitionKey, totalReceived,
totalRequired }: Props) => {
+export const AssetProgressCell = ({
+ dagId,
+ partitionedDagRunId,
+ partitionKey,
+ totalReceived,
+ totalRequired,
+}: Props) => {
const { i18n } = useTranslation();
- const { data, isLoading } =
usePartitionedDagRunServiceGetPendingPartitionedDagRun({ dagId, partitionKey });
+ const { data, isLoading } =
usePartitionedDagRunServiceGetPendingPartitionedDagRun({
+ dagId,
+ partitionedDagRunId,
+ partitionKey,
+ });
const assetExpression = data?.asset_expression ?? undefined;
const assets: Array<PartitionedDagRunAssetResponse> = data?.assets ?? [];
diff --git
a/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
index 6835033ad2c..b53769e10f3 100644
--- a/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/DagsList/PartitionScheduleModal.tsx
@@ -67,6 +67,7 @@ const getColumns = (
cell: ({ row }) => (
<AssetProgressCell
dagId={dagId}
+ partitionedDagRunId={row.original.id}
partitionKey={row.original.partition_key}
totalReceived={row.original.total_received}
totalRequired={row.original.total_required}
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
index 7016db222dd..1fe5df0b9a3 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_partitioned_dag_runs.py
@@ -1015,14 +1015,10 @@ class TestGetPendingPartitionedDagRun:
assert body["dag_id"] == "slash_key_dag"
assert body["partition_key"] == "region/us"
- def test_duplicate_pending_apdr_rows_return_latest(self, test_client,
dag_maker, session):
- """
- Duplicate pending APDR rows for the same (dag_id, partition_key) must
not 500.
-
- The model docstring for ``AssetPartitionDagRun`` says callers should
always
- work on the latest row when duplicates exist; the route must do the
same
- instead of raising ``MultipleResultsFound`` from ``.one_or_none()``.
- """
+ @pytest.mark.parametrize(
+ "selection", [None, "oldest", "newest", "missing", "other-dag",
"other-partition", "created"]
+ )
+ def test_pending_partition_details_select_run(self, test_client,
dag_maker, session, selection):
asset = Asset(uri="s3://bucket/dup_apdr", name="dup_apdr")
with dag_maker(
dag_id="dup_apdr_dag",
@@ -1030,40 +1026,52 @@ class TestGetPendingPartitionedDagRun:
serialized=True,
):
EmptyOperator(task_id="t")
- dag_maker.create_dagrun()
+ dag_run = dag_maker.create_dagrun()
dag_maker.sync_dagbag_to_db()
asset = session.scalar(select(AssetModel).where(AssetModel.uri ==
"s3://bucket/dup_apdr"))
- # Older duplicate row: no received events.
- stale_pdr = AssetPartitionDagRun(target_dag_id="dup_apdr_dag",
partition_key="2024-08-01")
- session.add(stale_pdr)
+ next_pdr = AssetPartitionDagRun(target_dag_id="dup_apdr_dag",
partition_key="2024-08-01")
+ session.add(next_pdr)
session.flush()
- # Newer duplicate row (higher id): has a received event.
- latest_pdr = AssetPartitionDagRun(target_dag_id="dup_apdr_dag",
partition_key="2024-08-01")
- session.add(latest_pdr)
+ later_pdr = AssetPartitionDagRun(target_dag_id="dup_apdr_dag",
partition_key="2024-08-01")
+ session.add(later_pdr)
session.flush()
- event = AssetEvent(asset_id=asset.id, timestamp=pendulum.now())
+ event = AssetEvent(asset_id=asset.id,
timestamp=pendulum.datetime(2024, 8, 1))
session.add(event)
session.flush()
session.add(
PartitionedAssetKeyLog(
asset_id=asset.id,
asset_event_id=event.id,
- asset_partition_dag_run_id=latest_pdr.id,
+ asset_partition_dag_run_id=next_pdr.id,
source_partition_key="2024-08-01",
target_dag_id="dup_apdr_dag",
target_partition_key="2024-08-01",
)
)
+ params = {"partition_key": "2024-08-01"}
+ if selection is not None:
+ params["partitioned_dag_run_id"] = next_pdr.id if selection ==
"oldest" else later_pdr.id
+ if selection == "missing":
+ params["partitioned_dag_run_id"] = -1
+ elif selection == "other-dag":
+ later_pdr.target_dag_id = "other-dag"
+ elif selection == "other-partition":
+ later_pdr.partition_key = "other-partition"
+ elif selection == "created":
+ later_pdr.created_dag_run_id = dag_run.id
session.commit()
- resp =
test_client.get("/pending_partitioned_dag_run/dup_apdr_dag?partition_key=2024-08-01")
+ resp = test_client.get("/pending_partitioned_dag_run/dup_apdr_dag",
params=params)
+ if selection in {"missing", "other-dag", "other-partition", "created"}:
+ assert resp.status_code == 404
+ return
assert resp.status_code == 200
body = resp.json()
- assert body["id"] == latest_pdr.id
- assert body["total_received"] == 1
+ assert body["id"] == (next_pdr.id if selection == "oldest" else
later_pdr.id)
+ assert body["total_received"] == (1 if selection == "oldest" else 0)
def test_non_rollup_many_to_one_received_capped_at_one(self, test_client,
dag_maker, session):
"""
diff --git a/airflow-core/tests/unit/assets/test_manager.py
b/airflow-core/tests/unit/assets/test_manager.py
index 8b72d1db57d..049df518569 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -55,6 +55,7 @@ from airflow.providers.standard.operators.empty import
EmptyOperator
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
+from tests_common.test_utils.asserts import capture_orm_selects
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
clear_db_apdr,
@@ -439,6 +440,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=rollup_fingerprint,
asset_id=asm.id,
+ allow_reuse=True,
session=_session,
).id
finally:
@@ -477,6 +479,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=fp,
asset_id=asm.id,
+ allow_reuse=True,
session=session,
)
assert first.partition_date == timezone.parse("2026-05-20T00:00:00")
@@ -488,6 +491,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=fp,
asset_id=asm.id,
+ allow_reuse=True,
session=session,
)
assert second.id == first.id # same pending APDR
@@ -508,6 +512,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=fp,
asset_id=asm.id,
+ allow_reuse=True,
session=session,
)
first =
AssetManager._get_or_create_apdr(target_partition_date=source_date, **kwargs)
@@ -536,6 +541,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=fp,
asset_id=asm.id,
+ allow_reuse=True,
session=session,
)
# First event carries no date (e.g. producer had no partition_date).
@@ -562,6 +568,7 @@ class TestAssetManager:
target_dag=testing_dag,
rollup_fingerprint=fp,
asset_id=asm.id,
+ allow_reuse=True,
session=session,
)
first = AssetManager._get_or_create_apdr(target_partition_date=date_1,
**kwargs)
@@ -616,6 +623,36 @@ class TestAssetManager:
assert apdr.partition_date is None
mock_log.exception.assert_called_once()
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ @pytest.mark.parametrize("allow_reuse", [True, False])
+ def test_get_or_create_apdr_allow_reuse(self, session, allow_reuse):
+ clear_db_apdr()
+ clear_db_pakl()
+ asset = AssetModel(uri="test://reuse/", name="reuse_asset",
group="asset")
+ dag = DagModel(dag_id="reuse_test_dag", is_stale=False,
bundle_name="testing")
+ session.add_all([asset, dag])
+ session.flush()
+
+ with capture_orm_selects("asset_partition_dag_run") as lookups:
+ apdrs = [
+ AssetManager._get_or_create_apdr(
+ target_key="key-1",
+ target_partition_date=None,
+ target_dag=dag,
+ rollup_fingerprint={},
+ asset_id=asset.id,
+ allow_reuse=allow_reuse,
+ session=session,
+ )
+ for _ in range(2)
+ ]
+
+ assert len(lookups) == (2 if allow_reuse else 0)
+ expected_count = 1 if allow_reuse else 2
+ assert len({apdr.id for apdr in apdrs}) == expected_count
+ assert
session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) ==
expected_count
+ assert all(apdr.created_dag_run_id is None for apdr in apdrs)
+
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("testing_dag_bundle")
def test_queue_partitioned_dags_stamps_rollup_fingerprint(self, session,
dag_maker):
diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py
b/airflow-core/tests/unit/dag_processing/test_dagbag.py
index 73a718580f0..d24443f709f 100644
--- a/airflow-core/tests/unit/dag_processing/test_dagbag.py
+++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py
@@ -892,6 +892,34 @@ class TestDagBag:
assert len(dagbag.import_errors) == 1
assert len(dagbag.dags) == 0
+ @pytest.mark.parametrize("batch_asset_events", [False, True])
+ @pytest.mark.parametrize("explicit", [False, None])
+ def test_process_file_compound_assets_require_batching(self, tmp_path,
batch_asset_events, explicit):
+ dag_file = tmp_path / "compound_assets.py"
+ batching_argument = "" if explicit is None else f",
batch_asset_events={explicit!r}"
+ dag_file.write_text(
+ textwrap.dedent(
+ f"""
+ from airflow.sdk import DAG, Asset, PartitionedAssetTimetable
+
+ with DAG("compound_assets", schedule=PartitionedAssetTimetable(
+ Asset("a") & Asset("b"){batching_argument}
+ )):
+ pass
+ """
+ )
+ )
+ with conf_vars({("scheduler", "batch_asset_events"):
str(batch_asset_events)}):
+ dagbag = DagBag(dag_folder=os.fspath(tmp_path))
+ if explicit is None:
+ assert not dagbag.import_errors
+ assert "compound_assets" in dagbag.dags
+ else:
+ assert not dagbag.dags
+ assert list(dagbag.import_errors.values()) == [
+ "AirflowTimetableInvalid: Asset AND conditions require
batch_asset_events=True"
+ ]
+
def test_process_file_invalid_param_check(self, tmp_path):
"""
test if an invalid param in the dags can be identified
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 5204fe776c0..530a5f99033 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -151,12 +151,13 @@ from airflow.sdk import (
)
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
-from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
+from airflow.sdk.definitions.timetables.assets import AssetTriggeredTimetable,
PartitionedAssetTimetable
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.serialization.encoders import ensure_serialized_asset
from airflow.serialization.serialized_objects import LazyDeserializedDAG
from airflow.timetables.base import DagRunInfo, DataInterval, Timetable,
compute_rollup_fingerprint
from airflow.timetables.simple import (
+ AssetTriggeredTimetable as CoreAssetTriggeredTimetable,
PartitionedAssetTimetable as CorePartitionedAssetTimetable,
PartitionedAtRuntime,
)
@@ -6518,6 +6519,7 @@ class TestSchedulerJob:
assert dr.creating_job_id == scheduler_job.id
@pytest.mark.need_serialized_dag
+ @conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_create_dag_runs_assets(self, session, dag_maker):
"""
Test various invariants of _create_dag_runs.
@@ -6652,6 +6654,7 @@ class TestSchedulerJob:
pytest.param(True, True, id="catchup-on-consumes-backlog"),
],
)
+ @conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_new_asset_triggered_dag_backlog_gated_by_catchup(
self, catchup, expects_old_event, session, dag_maker
):
@@ -6714,6 +6717,7 @@ class TestSchedulerJob:
assert {e.id for e in created_run.consumed_asset_events} == expected
@pytest.mark.need_serialized_dag
+ @conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_asset_events_out_of_order_are_both_consumed(self, session,
dag_maker):
"""Regression test for GH-54659.
@@ -6878,11 +6882,19 @@ class TestSchedulerJob:
@pytest.mark.need_serialized_dag
@pytest.mark.backend("postgres", "mysql")
- def test_create_dag_runs_when_concurrent_asset_events_created(self,
session: Session, dag_maker, caplog):
+ @pytest.mark.parametrize("batch_asset_events", [True, False])
+ def test_create_dag_runs_when_concurrent_asset_events_created(
+ self, session: Session, dag_maker, caplog, batch_asset_events
+ ):
ASSET_EVENT_COUNT = 30
asset = Asset(name="test_asset")
- with dag_maker(dag_id="consumer", schedule=asset, session=session):
+ with dag_maker(
+ dag_id="consumer",
+ schedule=AssetTriggeredTimetable(assets=asset,
batch_asset_events=batch_asset_events),
+ max_active_runs=ASSET_EVENT_COUNT,
+ session=session,
+ ):
pass
dag_model = dag_maker.dag_model
# Capture the dag_id as a plain string in the main thread. The worker
threads must not
@@ -13230,6 +13242,218 @@ def _produce_and_register_asset_event(
return apdr
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows")
[email protected]("batch_asset_events", [True, False])
+def test_partitioned_batch_asset_events(dag_maker: DagMaker, session: Session,
batch_asset_events: bool):
+ asset = Asset(name="partitioned-batch-asset")
+ dag_id = "partitioned-batch-consumer"
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=PartitionedAssetTimetable(assets=asset,
batch_asset_events=batch_asset_events),
+ session=session,
+ ):
+ EmptyOperator(task_id="task")
+ session.commit()
+ apdrs = [
+ _produce_and_register_asset_event(
+ dag_id=f"partitioned-batch-producer-{index}",
+ asset=asset,
+ partition_key="key-1",
+ session=session,
+ dag_maker=dag_maker,
+ )
+ for index in range(2)
+ ]
+ expected_count = 1 if batch_asset_events else 2
+ assert len({apdr.id for apdr in apdrs}) == expected_count
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ assert runner._create_dagruns_for_partitioned_asset_dags(session=session)
== {dag_id}
+ dag_runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_id)).all()
+ assert len(dag_runs) == expected_count
+ assert len({run.run_id for run in dag_runs}) == expected_count
+ assert {apdr.created_dag_run_id for apdr in apdrs} == {run.id for run in
dag_runs}
+ assert [len(run.consumed_asset_events) for run in dag_runs] == [2 //
expected_count] * expected_count
+ assert len({event.id for run in dag_runs for event in
run.consumed_asset_events}) == 2
+
+
[email protected]_serialized_dag
[email protected]("custom_grouping", [False, True])
[email protected]("batch_asset_events", [True, False])
[email protected]("event_spacing", [timedelta(0), timedelta(seconds=1)])
+def test_non_partitioned_batch_asset_events(
+ dag_maker: DagMaker,
+ session: Session,
+ batch_asset_events: bool,
+ event_spacing: timedelta,
+ custom_grouping,
+ mocker,
+):
+ asset = Asset(name="batch-consumer-asset")
+ with dag_maker(
+ dag_id="batch-consumer",
+ schedule=AssetTriggeredTimetable(assets=asset,
batch_asset_events=batch_asset_events),
+ session=session,
+ ):
+ EmptyOperator(task_id="task")
+ session.commit()
+ dag_model = session.get(DagModel, "batch-consumer")
+ asset_model = session.scalar(select(AssetModel).where(AssetModel.uri ==
asset.uri))
+ assert dag_model is not None
+ assert asset_model is not None
+ base = timezone.parse("2026-10-01T00:00:00+00:00")
+ events = [
+ AssetEvent(asset_id=asset_model.id, timestamp=base + index *
event_spacing) for index in range(2)
+ ]
+ session.add_all(events)
+ session.flush()
+ session.add_all(
+ AssetDagRunQueue(asset_id=asset_model.id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id)
+ for event in events
+ )
+ session.flush()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ if custom_grouping:
+ mocker.patch.object(
+ CoreAssetTriggeredTimetable,
+ "group_asset_events",
+ autospec=True,
+ return_value=[(base.add(hours=1), events)],
+ )
+ runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
+ session.flush()
+
+ dag_runs = session.scalars(
+ select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id).order_by(DagRun.id)
+ ).all()
+ expected_events = [events] if batch_asset_events or custom_grouping else
[[event] for event in events]
+ assert [{event.id for event in run.consumed_asset_events} for run in
dag_runs] == [
+ {event.id for event in group} for group in expected_events
+ ]
+ expected_run_dates = (
+ [base.add(hours=1)]
+ if custom_grouping
+ else [max(event.timestamp for event in group) for group in
expected_events]
+ )
+ assert [run.run_after for run in dag_runs] == expected_run_dates
+ assert len({run.run_id for run in dag_runs}) == len(expected_events)
+ assert all(
+ run.run_type == DagRunType.ASSET_TRIGGERED and run.state ==
DagRunState.QUEUED for run in dag_runs
+ )
+ assert session.scalar(select(func.count()).select_from(AssetDagRunQueue))
== 0
+
+ runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
+ assert session.scalar(
+ select(func.count()).select_from(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)
+ ) == len(expected_events)
+
+
[email protected]_serialized_dag
[email protected]("catchup", [False, True])
[email protected](
+ ("max_active_runs", "run_budget", "first_count"), [(1, 10, 1), (10, 1, 1),
(3, 2, 2)]
+)
+def test_unbatched_asset_runs_drain_backlog_within_limits(
+ dag_maker, session, monkeypatch, catchup, max_active_runs, run_budget,
first_count
+):
+ monkeypatch.setattr(DagModel, "NUM_DAGS_PER_DAGRUN_QUERY", run_budget)
+ asset = Asset(name="bounded-events")
+ dag_id = "bounded-consumer"
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=AssetTriggeredTimetable(assets=asset,
batch_asset_events=False),
+ max_active_runs=max_active_runs,
+ catchup=catchup,
+ session=session,
+ ):
+ EmptyOperator(task_id="consume")
+ session.commit()
+ dag_model = session.get(DagModel, dag_id)
+ asset_model = session.scalar(select(AssetModel).where(AssetModel.uri ==
asset.uri))
+ timestamp = timezone.parse("2026-10-01T00:00:00Z")
+ events = [AssetEvent(asset_id=asset_model.id, timestamp=timestamp) for _
in range(3)]
+ session.add_all(events)
+ session.flush()
+ queued_events = events[:1] if catchup else events
+ session.add_all(
+ AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_id,
asset_event_id=event.id)
+ for event in queued_events
+ )
+ session.flush()
+ runner = SchedulerJobRunner(job=Job(),
executors=[MockExecutor(do_update=False)])
+ runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
+ session.flush()
+ runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_id).order_by(DagRun.id)).all()
+ assert len(runs) == first_count
+ assert session.scalar(select(func.count()).select_from(AssetDagRunQueue))
> 0
+ if first_count == max_active_runs:
+ runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
+ assert (
+
session.scalar(select(func.count()).select_from(DagRun).where(DagRun.dag_id ==
dag_id))
+ == first_count
+ )
+
+ for _ in range(3):
+ for run in runs:
+ run.state = DagRunState.SUCCESS
+ session.flush()
+ query, triggered_dates = DagModel.dags_needing_dagruns(session=session)
+ runner._create_dag_runs_asset_triggered(
+ dag_models=[model for model in query.all() if model.dag_id in
triggered_dates], session=session
+ )
+ session.flush()
+ runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_id).order_by(DagRun.id)).all()
+
+ assert [run.consumed_asset_events for run in runs] == [[event] for event
in events]
+ assert session.scalar(select(func.count()).select_from(AssetDagRunQueue))
== 0
+
+
[email protected]_serialized_dag
[email protected](DagModel, "NUM_DAGS_PER_DAGRUN_QUERY", 3)
[email protected]("batch_first_dag", [False, True])
+def test_asset_runs_share_creation_budget(dag_maker, session, batch_first_dag):
+ asset = Asset(name="shared-budget")
+ dag_ids = ["budget-consumer-1", "budget-consumer-2"]
+ for dag_id in dag_ids:
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=AssetTriggeredTimetable(
+ assets=asset, batch_asset_events=batch_first_dag and dag_id ==
dag_ids[0]
+ ),
+ session=session,
+ ):
+ EmptyOperator(task_id="consume")
+ session.commit()
+ asset_model = session.scalar(select(AssetModel).where(AssetModel.uri ==
asset.uri))
+ events = [
+ AssetEvent(asset_id=asset_model.id,
timestamp=timezone.parse("2026-10-01T00:00:00Z"))
+ for _ in range(3)
+ ]
+ session.add_all(events)
+ session.flush()
+ session.add_all(
+ AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_id,
asset_event_id=event.id)
+ for dag_id in dag_ids
+ for event in events
+ )
+ session.flush()
+ query, triggered_dates = DagModel.dags_needing_dagruns(session=session)
+ models = [model for model in query.all() if model.dag_id in
triggered_dates]
+ assert {model.dag_id for model in models} == set(dag_ids)
+ runner = SchedulerJobRunner(job=Job(),
executors=[MockExecutor(do_update=False)])
+ runner._create_dag_runs_asset_triggered(dag_models=models, session=session)
+ session.flush()
+ counts = dict(session.execute(select(DagRun.dag_id,
func.count()).group_by(DagRun.dag_id)).all())
+ assert counts == dict.fromkeys(dag_ids, 1)
+ assert session.scalar(select(func.count()).select_from(AssetDagRunQueue))
== (2 if batch_first_dag else 4)
+
+
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
@pytest.mark.parametrize(
@@ -13807,6 +14031,7 @@ def
test_consumer_dag_run_partition_date_is_none_when_task_key_diverges(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_consumer_dag_listen_to_two_partitioned_asset(
dag_maker: DagMaker,
session: Session,
@@ -13887,6 +14112,7 @@ def test_consumer_dag_listen_to_two_partitioned_asset(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partitioned_dag_run_rollup_holds_until_window_complete(
dag_maker: DagMaker,
session: Session,
@@ -13945,6 +14171,7 @@ def
test_partitioned_dag_run_rollup_holds_until_window_complete(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partitioned_dag_run_segment_rollup_holds_until_all_segments_arrive(
dag_maker: DagMaker,
session: Session,
@@ -14013,6 +14240,7 @@ def
test_partitioned_dag_run_segment_rollup_holds_until_all_segments_arrive(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def
test_partitioned_dag_run_rollup_minimum_count_negative_fires_with_tolerated_gaps(
dag_maker: DagMaker,
session: Session,
@@ -14074,6 +14302,7 @@ def
test_partitioned_dag_run_rollup_minimum_count_negative_fires_with_tolerated_
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partitioned_dag_run_rollup_minimum_count_fires_when_threshold_met(
dag_maker: DagMaker,
session: Session,
@@ -14186,6 +14415,7 @@ def
test_partitioned_dag_run_rollup_treats_mapper_exception_as_not_satisfied(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partitioned_dag_run_rollup_survives_scheduler_restart_partial_arrival(
dag_maker: DagMaker,
session: Session,
@@ -14553,6 +14783,7 @@ def
test_partitioned_dag_run_fingerprint_compute_failure_does_not_clear_apdr(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partitioned_dag_run_rollup_idempotent_with_duplicate_upstream_events(
dag_maker: DagMaker,
session: Session,
@@ -14631,6 +14862,7 @@ def
test_partitioned_dag_run_rollup_idempotent_with_duplicate_upstream_events(
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("clear_asset_partition_rows")
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_consumer_dag_listen_to_two_partitioned_asset_with_key_1_mapper(
dag_maker: DagMaker,
session: Session,
@@ -15625,6 +15857,7 @@ class TestReapStaleConnectionTests:
),
],
)
+@conf_vars({("scheduler", "batch_asset_events"): "True"})
def test_partition_date_populated_on_dagrun(
dag_maker: DagMaker,
session: Session,
diff --git a/airflow-core/tests/unit/serialization/definitions/test_assets.py
b/airflow-core/tests/unit/serialization/definitions/test_assets.py
index 473f645fb54..98ba48d0005 100644
--- a/airflow-core/tests/unit/serialization/definitions/test_assets.py
+++ b/airflow-core/tests/unit/serialization/definitions/test_assets.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+import pytest
+
from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
from airflow.serialization.definitions.assets import (
SerializedAsset,
@@ -36,6 +38,23 @@ asset1 = SerializedAsset(
)
[email protected](
+ ("condition", "requires_batching"),
+ [
+ (asset1, False),
+ (SerializedAssetAlias(name="alias", group=""), False),
+ (SerializedAssetAll([]), False),
+ (SerializedAssetAll([asset1]), False),
+ (SerializedAssetAny([asset1, asset1]), False),
+ (SerializedAssetAll([asset1, asset1]), True),
+ (SerializedAssetAny([asset1, SerializedAssetAll([asset1, asset1])]),
True),
+ (SerializedAssetAll([SerializedAssetAny([SerializedAssetAll([asset1,
asset1])])]), True),
+ ],
+)
+def test_asset_condition_requires_batching(condition, requires_batching):
+ assert condition.requires_batching is requires_batching
+
+
def test_asset_iter_assets():
assert list(asset1.iter_assets()) == [
(SerializedAssetUniqueKey("asset-1", "s3://bucket1/data1/"), asset1)
diff --git a/airflow-core/tests/unit/timetables/test_assets_timetable.py
b/airflow-core/tests/unit/timetables/test_assets_timetable.py
index 5ab7667cd11..7f4b41e578c 100644
--- a/airflow-core/tests/unit/timetables/test_assets_timetable.py
+++ b/airflow-core/tests/unit/timetables/test_assets_timetable.py
@@ -19,6 +19,7 @@
from __future__ import annotations
from collections import defaultdict
+from contextlib import nullcontext
from typing import Any
import pytest
@@ -36,18 +37,64 @@ from airflow.sdk import (
AssetOrTimeSchedule as SdkAssetOrTimeSchedule,
)
from airflow.sdk.bases.timetable import BaseTimetable
-from airflow.sdk.definitions.timetables.assets import AssetTriggeredTimetable
as SdkAssetTriggeredTimetable
+from airflow.sdk.definitions.timetables.assets import (
+ AssetTriggeredTimetable as SdkAssetTriggeredTimetable,
+ PartitionedAssetTimetable as SdkPartitionedAssetTimetable,
+)
+from airflow.sdk.definitions.timetables.simple import NullTimetable as
SdkNullTimetable
from airflow.sdk.exceptions import AirflowTimetableInvalid
-from airflow.serialization.definitions.assets import SerializedAsset,
SerializedAssetAll, SerializedAssetAny
+from airflow.serialization.decoders import decode_timetable
+from airflow.serialization.definitions.assets import (
+ SerializedAsset,
+ SerializedAssetAll,
+ SerializedAssetAny,
+ SerializedAssetBase,
+)
+from airflow.serialization.encoders import encode_timetable,
ensure_serialized_asset
from airflow.serialization.serialized_objects import DagSerialization
from airflow.timetables.assets import (
AssetAndTimeSchedule as CoreAssetAndTimeSchedule,
AssetOrTimeSchedule as CoreAssetOrTimeSchedule,
)
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction,
Timetable
-from airflow.timetables.simple import AssetTriggeredTimetable
+from airflow.timetables.simple import AssetTriggeredTimetable, NullTimetable,
PartitionedAssetTimetable
from airflow.utils.types import DagRunType
+from tests_common.test_utils.config import conf_vars
+
+
[email protected]("batch_asset_events", [False, True])
[email protected]("event_count", [0, 1, 3])
+def test_group_asset_events(batch_asset_events, event_count):
+ timetable = AssetTriggeredTimetable(Asset("a"),
batch_asset_events=batch_asset_events)
+ triggered_date = DateTime(2026, 10, 7, tzinfo=UTC)
+ events = [
+ AssetEvent(asset_id=1, timestamp=triggered_date.subtract(hours=index +
1))
+ for index in range(event_count)
+ ]
+
+ if batch_asset_events and events:
+ expected = [(triggered_date, events)]
+ else:
+ expected = [(event.timestamp, [event]) for event in events]
+ assert list(timetable.group_asset_events(events, triggered_date)) ==
expected
+
+
+def test_custom_asset_condition_requires_batching():
+ class CustomCondition(SerializedAssetBase):
+ @property
+ def requires_batching(self):
+ return True
+
+ condition = SerializedAssetAny([CustomCondition()])
+ with conf_vars({("scheduler", "batch_asset_events"): "False"}):
+ timetable = AssetTriggeredTimetable(condition)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+ with pytest.raises(AirflowTimetableInvalid,
match="batch_asset_events=True"):
+ AssetTriggeredTimetable(condition, batch_asset_events=False).validate()
+
class MockTimetable(Timetable):
"""
@@ -268,6 +315,7 @@ def test_serialization(sdk_asset_timetable:
SdkAssetOrTimeSchedule, monkeypatch:
)
serialized = _serializer.serialize_timetable(sdk_asset_timetable)
assert serialized == {
+ "batch_asset_events": False,
"timetable": "mock_serialized_timetable",
"asset_condition": {
"__type": "asset_all",
@@ -491,6 +539,95 @@ def test_run_ordering_inheritance(core_asset_timetable) ->
None:
assert core_asset_timetable.run_ordering ==
AssetTriggeredTimetable.run_ordering
[email protected]("batch_asset_events", [True, False])
[email protected](
+ "timetable_type",
+ [
+ AssetTriggeredTimetable,
+ SdkAssetTriggeredTimetable,
+ PartitionedAssetTimetable,
+ SdkPartitionedAssetTimetable,
+ CoreAssetOrTimeSchedule,
+ SdkAssetOrTimeSchedule,
+ ],
+)
+def test_batch_asset_events_roundtrip(timetable_type, batch_asset_events):
+ asset = Asset("test")
+ kwargs = {}
+ if timetable_type is CoreAssetOrTimeSchedule:
+ kwargs["timetable"] = NullTimetable()
+ elif timetable_type is SdkAssetOrTimeSchedule:
+ kwargs["timetable"] = SdkNullTimetable()
+ timetable = timetable_type(assets=asset,
batch_asset_events=batch_asset_events, **kwargs)
+
+ with conf_vars({("scheduler", "batch_asset_events"): str(not
batch_asset_events)}):
+ serialized = encode_timetable(timetable)
+ deserialized = decode_timetable(serialized)
+
+ assert serialized["__var"]["batch_asset_events"] is batch_asset_events
+ assert deserialized.batch_asset_events is batch_asset_events
+ assert deserialized.asset_condition == ensure_serialized_asset(asset)
+
+
[email protected]("configured", [False, True])
[email protected]("explicit", [None, False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
CoreAssetOrTimeSchedule]
+)
+def test_batching_configuration(configured, explicit, timetable_type):
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
CoreAssetOrTimeSchedule else {}
+ if explicit is not None:
+ kwargs["batch_asset_events"] = explicit
+ with conf_vars({("scheduler", "batch_asset_events"): str(configured)}):
+ timetable = timetable_type(assets=Asset("test"), **kwargs)
+ assert timetable.batch_asset_events is (configured if explicit is None
else explicit)
+
+
[email protected]("batch_asset_events", [False, True])
[email protected]("nested", [False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
CoreAssetOrTimeSchedule]
+)
+def test_compound_asset_condition_requires_batching(batch_asset_events,
nested, timetable_type):
+ condition = Asset("a") & Asset("b")
+ if nested:
+ condition = Asset("c") | condition
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
CoreAssetOrTimeSchedule else {}
+ timetable = timetable_type(assets=condition,
batch_asset_events=batch_asset_events, **kwargs)
+ expected = nullcontext() if batch_asset_events else
pytest.raises(AirflowTimetableInvalid, match="AND")
+ with expected:
+ timetable.validate()
+
+
[email protected]("configured", [False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
CoreAssetOrTimeSchedule]
+)
+def test_compound_asset_condition_defaults_to_batching(configured,
timetable_type):
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
CoreAssetOrTimeSchedule else {}
+ with conf_vars({("scheduler", "batch_asset_events"): str(configured)}):
+ timetable = timetable_type(assets=Asset("a") & Asset("b"), **kwargs)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+
+def test_or_condition_can_disable_batching():
+ AssetTriggeredTimetable(assets=Asset("a") | Asset("b"),
batch_asset_events=False).validate()
+
+
[email protected]("configured", [False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
CoreAssetOrTimeSchedule]
+)
+def test_legacy_asset_timetable_defaults_to_batching(timetable_type,
configured):
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
CoreAssetOrTimeSchedule else {}
+ serialized = encode_timetable(timetable_type(assets=Asset("test"),
**kwargs))
+ del serialized["__var"]["batch_asset_events"]
+
+ with conf_vars({("scheduler", "batch_asset_events"): str(configured)}):
+ assert decode_timetable(serialized).batch_asset_events is True
+
+
@pytest.mark.db_test
class TestAssetConditionWithTimetable:
@pytest.fixture(autouse=True)
@@ -567,6 +704,7 @@ class TestAssetConditionWithTimetable:
serialized_timetable_dict =
DagSerialization.to_dict(dag)["dag"]["timetable"]["__var"]
assert serialized_timetable_dict == {
+ "batch_asset_events": True,
"asset_condition": {
"__type": "asset_any",
"objects": [
diff --git a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py
b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py
index 72dd2773bc9..8151426e3b8 100644
--- a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py
+++ b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py
@@ -18,7 +18,7 @@
from __future__ import annotations
from collections.abc import Callable, Iterable
-from contextlib import ExitStack
+from contextlib import ExitStack, nullcontext
from typing import TYPE_CHECKING
from unittest import mock
@@ -26,7 +26,7 @@ import pendulum
import pytest
from airflow._shared.module_loading import qualname
-from airflow.exceptions import InvalidPartitionKeyError
+from airflow.exceptions import AirflowTimetableInvalid,
InvalidPartitionKeyError
from airflow.partition_mappers.base import RollupMapper
from airflow.partition_mappers.identity import IdentityMapper as IdentityMapper
from airflow.partition_mappers.temporal import StartOfDayMapper
@@ -37,6 +37,8 @@ from airflow.serialization.encoders import
ensure_serialized_asset
from airflow.serialization.enums import DagAttributeTypes
from airflow.timetables.simple import PartitionedAssetTimetable
+from tests_common.test_utils.config import conf_vars
+
if TYPE_CHECKING:
from airflow.partition_mappers.base import PartitionMapper
@@ -194,6 +196,7 @@ class TestPartitionedAssetTimetable:
assets=ser_asset, partition_mapper_config={ser_asset:
IdentityMapper()}
)
assert timetable.serialize() == {
+ "batch_asset_events": False,
"asset_condition": {
"__type": DagAttributeTypes.ASSET,
"name": "test",
@@ -225,6 +228,7 @@ class TestPartitionedAssetTimetable:
def test_deserialize(self):
timetable = PartitionedAssetTimetable.deserialize(
{
+ "batch_asset_events": True,
"asset_condition": {
"__type": DagAttributeTypes.ASSET,
"name": "test",
@@ -257,6 +261,52 @@ class TestPartitionedAssetTimetable:
assert timetable.asset_condition == ser_asset
assert isinstance(timetable.default_partition_mapper, IdentityMapper)
assert isinstance(timetable.partition_mapper_config[ser_asset],
IdentityMapper)
+ assert timetable.batch_asset_events is True
+
+ def test_serialize_deserialize_batch_asset_events_false(self):
+ """Serialize/deserialize round-trip preserves
batch_asset_events=False."""
+ ser_asset = ensure_serialized_asset(Asset("test"))
+ timetable = PartitionedAssetTimetable(
+ assets=ser_asset,
+ partition_mapper_config={ser_asset: IdentityMapper()},
+ batch_asset_events=False,
+ )
+ serialized = timetable.serialize()
+ assert serialized["batch_asset_events"] is False
+
+ deserialized = PartitionedAssetTimetable.deserialize(serialized)
+ assert deserialized.batch_asset_events is False
+
+ @pytest.mark.parametrize("batch_asset_events", [False, True])
+ @pytest.mark.parametrize("use_default_mapper", [False, True])
+ def test_rollup_requires_batching(self, batch_asset_events,
use_default_mapper):
+ asset = ensure_serialized_asset(Asset("test"))
+ mapper = RollupMapper(upstream_mapper=StartOfDayMapper(),
window=DayWindow())
+ kwargs = (
+ {"default_partition_mapper": mapper}
+ if use_default_mapper
+ else {"partition_mapper_config": {asset: mapper}}
+ )
+ timetable = PartitionedAssetTimetable(assets=asset,
batch_asset_events=batch_asset_events, **kwargs)
+ expected = (
+ nullcontext() if batch_asset_events else
pytest.raises(AirflowTimetableInvalid, match="rollups")
+ )
+ with expected:
+ timetable.validate()
+
+ @pytest.mark.parametrize("use_default_mapper", [False, True])
+ def test_rollup_defaults_to_batching(self, use_default_mapper):
+ asset = ensure_serialized_asset(Asset("test"))
+ mapper = RollupMapper(upstream_mapper=StartOfDayMapper(),
window=DayWindow())
+ kwargs = (
+ {"default_partition_mapper": mapper}
+ if use_default_mapper
+ else {"partition_mapper_config": {asset: mapper}}
+ )
+ with conf_vars({("scheduler", "batch_asset_events"): "False"}):
+ timetable = PartitionedAssetTimetable(assets=asset, **kwargs)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
@pytest.mark.parametrize(
"asset_like",
diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
index 159ac9d316f..88006cf74dc 100644
--- a/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/plugins/test_utils.py
@@ -50,7 +50,11 @@ from airflow.utils.types import DagRunType
from tests_common.test_utils.compat import (
BashOperator,
)
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS
+from tests_common.test_utils.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
if AIRFLOW_V_3_1_PLUS:
from airflow.models.dag import get_next_data_interval
@@ -437,7 +441,7 @@ def test_serialize_timetable_complex_with_alias():
dag.timetable = AssetTriggeredTimetable(asset)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ expected: dict = {
"asset_condition": {
"__type": DagAttributeTypes.ASSET_ANY,
"objects": [
@@ -478,38 +482,47 @@ def test_serialize_timetable_complex_with_alias():
],
},
],
- }
+ },
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["batch_asset_events"] = True
+ assert dag_info.timetable == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
def test_serialize_timetable_single_asset():
dag = DAG(dag_id="test", start_date=datetime.datetime(2025, 1, 1),
schedule=Asset("a"))
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ expected: dict = {
"asset_condition": {
"__type": DagAttributeTypes.ASSET,
"uri": "a",
"name": "a",
"group": "asset",
"extra": {},
- }
+ },
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["batch_asset_events"] = False
+ assert dag_info.timetable == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
def test_serialize_timetable_list_of_assets():
dag = DAG(dag_id="test", start_date=datetime.datetime(2025, 1, 1),
schedule=[Asset("a"), Asset("b")])
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ expected: dict = {
"asset_condition": {
"__type": DagAttributeTypes.ASSET_ALL,
"objects": [
{"__type": DagAttributeTypes.ASSET, "uri": "a", "name": "a",
"group": "asset", "extra": {}},
{"__type": DagAttributeTypes.ASSET, "uri": "b", "name": "b",
"group": "asset", "extra": {}},
],
- }
+ },
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["batch_asset_events"] = True
+ assert dag_info.timetable == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -521,7 +534,7 @@ def
test_serialize_timetable_with_complex_logical_condition():
& (Asset("ds3") | Asset("ds4", extra={"another_extra": 345})),
)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ expected: dict = {
"asset_condition": {
"__type": DagAttributeTypes.ASSET_ALL,
"objects": [
@@ -564,8 +577,11 @@ def
test_serialize_timetable_with_complex_logical_condition():
],
},
],
- }
+ },
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["batch_asset_events"] = True
+ assert dag_info.timetable == expected
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This test checks
serialization only in 3.0 conditions")
@@ -583,7 +599,7 @@ def
test_serialize_timetable_with_dataset_or_time_schedule():
),
)
dag_info = DagInfo(dag)
- assert dag_info.timetable == {
+ expected: dict = {
"timetable": {
Encoding.TYPE: "airflow.timetables.trigger.CronTriggerTimetable",
Encoding.VAR: {
@@ -637,6 +653,9 @@ def
test_serialize_timetable_with_dataset_or_time_schedule():
],
},
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["batch_asset_events"] = True
+ assert dag_info.timetable == expected
@pytest.mark.skipif(
diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
index 3fa10b1e1ab..4078993937e 100644
--- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
@@ -2645,7 +2645,7 @@ class TestDagInfoAirflow3:
)
result = DagInfo(dag)
- assert dict(result) == {
+ expected: dict = {
"dag_id": "dag_id",
"description": None,
"fileloc": pathlib.Path(__file__).resolve().as_posix(),
@@ -2660,10 +2660,13 @@ class TestDagInfoAirflow3:
"name": "uri1",
"group": "asset",
"extra": {"a": 1},
- }
+ },
},
"timetable_summary": "Asset",
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["timetable"]["batch_asset_events"] = False
+ assert dict(result) == expected
def test_dag_info_schedule_list_single_assets(self):
dag = DAG(
@@ -2673,7 +2676,7 @@ class TestDagInfoAirflow3:
)
result = DagInfo(dag)
- assert dict(result) == {
+ expected: dict = {
"dag_id": "dag_id",
"description": None,
"fileloc": pathlib.Path(__file__).resolve().as_posix(),
@@ -2693,10 +2696,13 @@ class TestDagInfoAirflow3:
"extra": {"a": 1},
}
],
- }
+ },
},
"timetable_summary": "Asset",
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["timetable"]["batch_asset_events"] = False
+ assert dict(result) == expected
def test_dag_info_schedule_list_two_assets(self):
dag = DAG(
@@ -2706,7 +2712,7 @@ class TestDagInfoAirflow3:
)
result = DagInfo(dag)
- assert dict(result) == {
+ expected: dict = {
"dag_id": "dag_id",
"description": None,
"fileloc": pathlib.Path(__file__).resolve().as_posix(),
@@ -2727,10 +2733,13 @@ class TestDagInfoAirflow3:
},
{"__type": "asset", "uri": "uri2", "name": "uri2",
"group": "asset", "extra": {}},
],
- }
+ },
},
"timetable_summary": "Asset",
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["timetable"]["batch_asset_events"] = True
+ assert dict(result) == expected
def test_dag_info_schedule_assets_logical_condition(self):
dag = DAG(
@@ -2740,7 +2749,7 @@ class TestDagInfoAirflow3:
)
result = DagInfo(dag)
- assert dict(result) == {
+ expected: dict = {
"dag_id": "dag_id",
"description": None,
"fileloc": pathlib.Path(__file__).resolve().as_posix(),
@@ -2791,10 +2800,13 @@ class TestDagInfoAirflow3:
],
},
],
- }
+ },
},
"timetable_summary": "Asset",
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["timetable"]["batch_asset_events"] = True
+ assert dict(result) == expected
def test_dag_info_schedule_asset_or_time_schedule(self):
from airflow.timetables.assets import AssetOrTimeSchedule
@@ -2809,7 +2821,7 @@ class TestDagInfoAirflow3:
)
result = DagInfo(dag)
- assert dict(result) == {
+ expected: dict = {
"dag_id": "dag_id",
"description": None,
"fileloc": pathlib.Path(__file__).resolve().as_posix(),
@@ -2873,6 +2885,9 @@ class TestDagInfoAirflow3:
},
"timetable_summary": "Asset or */4 3 * * *",
}
+ if AIRFLOW_V_3_4_PLUS:
+ expected["timetable"]["batch_asset_events"] = True
+ assert dict(result) == expected
class TestDagRunInfoDeadlines:
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index 78809e802ef..1de617447fd 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -208,6 +208,8 @@ Timetables
.. autoapiclass:: airflow.sdk.AssetOrTimeSchedule
+.. autoapiclass:: airflow.sdk.AssetTriggeredTimetable
+
.. autoapiclass:: airflow.sdk.CronDataIntervalTimetable
.. autoapiclass:: airflow.sdk.CronTriggerTimetable
diff --git a/task-sdk/src/airflow/sdk/__init__.py
b/task-sdk/src/airflow/sdk/__init__.py
index 91822202553..c163df52577 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -28,6 +28,7 @@ __all__ = [
"AssetAny",
"AssetAndTimeSchedule",
"AssetOrTimeSchedule",
+ "AssetTriggeredTimetable",
"AssetWatcher",
"AsyncCallback",
"BaseAsyncOperator",
@@ -210,6 +211,7 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.timetables.assets import (
AssetAndTimeSchedule,
AssetOrTimeSchedule,
+ AssetTriggeredTimetable,
PartitionedAssetTimetable,
PartitionedAtRuntime,
)
@@ -242,6 +244,7 @@ __lazy_imports: dict[str, str] = {
"AssetAndTimeSchedule": ".definitions.timetables.assets",
"AssetAny": ".definitions.asset",
"AssetOrTimeSchedule": ".definitions.timetables.assets",
+ "AssetTriggeredTimetable": ".definitions.timetables.assets",
"AssetWatcher": ".definitions.asset",
"AsyncCallback": ".definitions.callback",
"BaseAsyncOperator": ".bases.operator",
diff --git a/task-sdk/src/airflow/sdk/__init__.pyi
b/task-sdk/src/airflow/sdk/__init__.pyi
index c3a94d30a35..1f9e10b5f5a 100644
--- a/task-sdk/src/airflow/sdk/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/__init__.pyi
@@ -114,6 +114,7 @@ from airflow.sdk.definitions.template import literal as
literal
from airflow.sdk.definitions.timetables.assets import (
AssetAndTimeSchedule,
AssetOrTimeSchedule,
+ AssetTriggeredTimetable,
PartitionedAssetTimetable,
PartitionedAtRuntime,
)
@@ -147,6 +148,7 @@ __all__ = [
"AssetAndTimeSchedule",
"AssetAny",
"AssetOrTimeSchedule",
+ "AssetTriggeredTimetable",
"AssetWatcher",
"AsyncCallback",
"BaseAsyncOperator",
diff --git a/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
index 9f91bd4a78a..a5318804b68 100644
--- a/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
+++ b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
@@ -231,6 +231,11 @@ class BaseAsset:
:meta private:
"""
+ @property
+ def requires_batching(self) -> bool:
+ """Whether this condition needs multiple asset events to be
satisfied."""
+ return False
+
def __or__(self, other: BaseAsset) -> BaseAsset:
if not isinstance(other, BaseAsset):
return NotImplemented
@@ -487,6 +492,10 @@ class AssetBooleanCondition(BaseAsset):
def __hash__(self) -> int:
return hash(tuple(self.objects))
+ @property
+ def requires_batching(self) -> bool:
+ return any(obj.requires_batching for obj in self.objects)
+
class AssetAny(AssetBooleanCondition):
"""Use to combine assets schedule references in an "or" relationship."""
@@ -506,6 +515,10 @@ class AssetAll(AssetBooleanCondition):
agg_func = all # type: ignore[assignment]
+ @property
+ def requires_batching(self) -> bool:
+ return len(self.objects) > 1 or super().requires_batching
+
def __and__(self, other: BaseAsset) -> BaseAsset:
if not isinstance(other, BaseAsset):
return NotImplemented
diff --git a/task-sdk/src/airflow/sdk/definitions/timetables/assets.py
b/task-sdk/src/airflow/sdk/definitions/timetables/assets.py
index bc09cd8f2b4..0ba492a7b0d 100644
--- a/task-sdk/src/airflow/sdk/definitions/timetables/assets.py
+++ b/task-sdk/src/airflow/sdk/definitions/timetables/assets.py
@@ -24,6 +24,7 @@ import attrs
from airflow.sdk.bases.timetable import BaseTimetable
from airflow.sdk.definitions.asset import AssetAll, BaseAsset
from airflow.sdk.definitions.partition_mappers.identity import IdentityMapper
+from airflow.sdk.exceptions import AirflowTimetableInvalid
if TYPE_CHECKING:
from collections.abc import Collection
@@ -32,28 +33,64 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.partition_mappers.base import PartitionMapper
+def _get_default_batch_asset_events(timetable: AssetTriggeredTimetable) ->
bool:
+ from airflow.sdk.configuration import conf
+
+ return (
+ conf.getboolean("scheduler", "batch_asset_events", fallback=False)
+ or timetable.get_batching_requirement() is not None
+ )
+
+
@attrs.define
class AssetTriggeredTimetable(BaseTimetable):
"""
- Timetable that never schedules anything.
+ Schedule a Dag when its asset condition is satisfied.
- This should not be directly used anywhere, but only set if a DAG is
triggered by assets.
-
- :meta private:
+ :param assets: The asset expression that triggers the Dag.
+ :param batch_asset_events: Consume queued events together in one Dag run.
Defaults
+ to ``[scheduler] batch_asset_events``, which is false, except that
conditions
+ combining multiple assets with ``&`` and ``RollupMapper`` partition
mappers always
+ default to batching. Passing ``False`` explicitly for those is
rejected.
"""
asset_triggered = True
asset_condition: BaseAsset = attrs.field(alias="assets")
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events,
takes_self=True), kw_only=True
+ )
+
+ def get_batching_requirement(self) -> str | None:
+ """Return why this timetable cannot run without event batching, or
``None``."""
+ if self.asset_condition.requires_batching:
+ return "Asset AND conditions require batch_asset_events=True"
+ return None
+
+ def validate(self) -> None:
+ if not self.batch_asset_events and (reason :=
self.get_batching_requirement()):
+ raise AirflowTimetableInvalid(reason)
@attrs.define
class PartitionedAssetTimetable(AssetTriggeredTimetable):
"""Asset-driven timetable that listens for partitioned assets."""
- asset_condition: BaseAsset = attrs.field(alias="assets")
partition_mapper_config: dict[BaseAsset, PartitionMapper] =
attrs.field(factory=dict)
default_partition_mapper: PartitionMapper = IdentityMapper()
+ # The factory needs the partition mappers to be initialized first.
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events,
takes_self=True), kw_only=True
+ )
+
+ def get_batching_requirement(self) -> str | None:
+ if any(
+ mapper.is_rollup
+ for mapper in (self.default_partition_mapper,
*self.partition_mapper_config.values())
+ ):
+ return "Partition rollups require batch_asset_events=True"
+ return super().get_batching_requirement()
+
class PartitionedAtRuntime(BaseTimetable):
"""Marker timetable indicating that partition key(s) are determined at
runtime."""
@@ -82,6 +119,10 @@ class AssetOrTimeSchedule(AssetTriggeredTimetable):
asset_condition: BaseAsset = attrs.field(alias="assets",
converter=_coerce_assets)
timetable: BaseTimetable
+ batch_asset_events: bool = attrs.field(
+ default=attrs.Factory(_get_default_batch_asset_events,
takes_self=True), kw_only=True
+ )
+
def __attrs_post_init__(self) -> None:
self.active_runs_limit = self.timetable.active_runs_limit
self.can_be_scheduled = self.timetable.can_be_scheduled
diff --git a/task-sdk/tests/task_sdk/definitions/test_asset.py
b/task-sdk/tests/task_sdk/definitions/test_asset.py
index 38c9cd2064f..413d0f4a02a 100644
--- a/task-sdk/tests/task_sdk/definitions/test_asset.py
+++ b/task-sdk/tests/task_sdk/definitions/test_asset.py
@@ -45,6 +45,22 @@ from tests_common.test_utils.dag import create_scheduler_dag
ASSET_MODULE_PATH = "airflow.sdk.definitions.asset"
[email protected](
+ ("condition", "requires_batching"),
+ [
+ (Asset("a"), False),
+ (AssetAll(), False),
+ (AssetAll(Asset("a")), False),
+ (Asset("a") | Asset("b"), False),
+ (Asset("a") & Asset("b"), True),
+ (Asset("a") | (Asset("b") & Asset("c")), True),
+ (AssetAll(Asset("a") | (Asset("b") & Asset("c"))), True),
+ ],
+)
+def test_asset_condition_requires_batching(condition, requires_batching):
+ assert condition.requires_batching is requires_batching
+
+
@pytest.mark.parametrize(
("sql_conn_value", "name", "should_raise"),
[
diff --git a/task-sdk/tests/task_sdk/definitions/timetables/test_assets.py
b/task-sdk/tests/task_sdk/definitions/timetables/test_assets.py
new file mode 100644
index 00000000000..2c5fee02206
--- /dev/null
+++ b/task-sdk/tests/task_sdk/definitions/timetables/test_assets.py
@@ -0,0 +1,166 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from contextlib import nullcontext
+
+import pytest
+
+from airflow.sdk import (
+ DAG,
+ Asset,
+ AssetOrTimeSchedule,
+ AssetTriggeredTimetable,
+ DayWindow,
+ IdentityMapper,
+ PartitionedAssetTimetable,
+ RollupMapper,
+ StartOfDayMapper,
+)
+from airflow.sdk.definitions.timetables.simple import NullTimetable
+from airflow.sdk.exceptions import AirflowTimetableInvalid
+
+from tests_common.test_utils.config import conf_vars
+
+
[email protected]("batch_asset_events", [True, False])
+def
test_partitioned_timetable_preserves_positional_mapper_arguments(batch_asset_events):
+ asset = Asset("test")
+ mapper = IdentityMapper()
+ mapping = {asset: mapper}
+ timetable = PartitionedAssetTimetable(asset, mapping, mapper,
batch_asset_events=batch_asset_events)
+
+ assert timetable.asset_condition == asset
+ assert timetable.partition_mapper_config == mapping
+ assert timetable.default_partition_mapper is mapper
+ assert timetable.batch_asset_events is batch_asset_events
+
+
[email protected]("configured", [False, True])
[email protected]("explicit", [None, False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
AssetOrTimeSchedule]
+)
+def test_batching_configuration(configured, explicit, timetable_type):
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
AssetOrTimeSchedule else {}
+ if explicit is not None:
+ kwargs["batch_asset_events"] = explicit
+ with conf_vars({("scheduler", "batch_asset_events"): str(configured)}):
+ timetable = timetable_type(assets=Asset("test"), **kwargs)
+ assert timetable.batch_asset_events is (configured if explicit is None
else explicit)
+
+
+def test_asset_schedules_do_not_batch_by_default():
+ dag = DAG("default-asset-batching", schedule=Asset("test"))
+ assert dag.timetable.batch_asset_events is False
+
+
[email protected]("batch_asset_events", [False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
AssetOrTimeSchedule]
+)
[email protected]("nested", [False, True])
+def test_compound_asset_condition_requires_batching(batch_asset_events,
timetable_type, nested):
+ condition = Asset("a") & Asset("b")
+ if nested:
+ condition = Asset("c") | condition
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
AssetOrTimeSchedule else {}
+ dag = DAG(
+ "compound-assets",
+ schedule=timetable_type(assets=condition,
batch_asset_events=batch_asset_events, **kwargs),
+ )
+ expected = nullcontext() if batch_asset_events else
pytest.raises(AirflowTimetableInvalid, match="AND")
+ with expected:
+ dag.validate()
+
+
[email protected]("batch_asset_events", [False, True])
[email protected]("use_default_mapper", [False, True])
+def test_rollup_requires_batching(batch_asset_events, use_default_mapper):
+ asset = Asset("test")
+ mapper = RollupMapper(upstream_mapper=StartOfDayMapper(),
window=DayWindow())
+ kwargs = (
+ {"default_partition_mapper": mapper}
+ if use_default_mapper
+ else {"partition_mapper_config": {asset: mapper}}
+ )
+ timetable = PartitionedAssetTimetable(assets=asset,
batch_asset_events=batch_asset_events, **kwargs)
+ expected = (
+ nullcontext() if batch_asset_events else
pytest.raises(AirflowTimetableInvalid, match="rollups")
+ )
+ with expected:
+ timetable.validate()
+
+
[email protected]("configured", [False, True])
[email protected](
+ "timetable_type", [AssetTriggeredTimetable, PartitionedAssetTimetable,
AssetOrTimeSchedule]
+)
+def test_compound_asset_condition_defaults_to_batching(configured,
timetable_type):
+ kwargs = {"timetable": NullTimetable()} if timetable_type is
AssetOrTimeSchedule else {}
+ with conf_vars({("scheduler", "batch_asset_events"): str(configured)}):
+ timetable = timetable_type(assets=Asset("a") & Asset("b"), **kwargs)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+
+def test_asset_list_schedule_defaults_to_batching():
+ dag = DAG("asset-list-schedule", schedule=[Asset("a"), Asset("b")])
+ assert dag.timetable.batch_asset_events is True
+ dag.validate()
+
+
[email protected]("use_default_mapper", [False, True])
+def test_rollup_defaults_to_batching(use_default_mapper):
+ asset = Asset("test")
+ mapper = RollupMapper(upstream_mapper=StartOfDayMapper(),
window=DayWindow())
+ kwargs = (
+ {"default_partition_mapper": mapper}
+ if use_default_mapper
+ else {"partition_mapper_config": {asset: mapper}}
+ )
+ with conf_vars({("scheduler", "batch_asset_events"): "False"}):
+ timetable = PartitionedAssetTimetable(assets=asset, **kwargs)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+
+def test_or_condition_can_disable_batching():
+ timetable = AssetTriggeredTimetable(assets=Asset("a") | Asset("b"),
batch_asset_events=False)
+ timetable.validate()
+
+
+def test_asset_or_time_schedule_list_defaults_to_batching():
+ timetable = AssetOrTimeSchedule(assets=[Asset("a"), Asset("b")],
timetable=NullTimetable())
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+
+def test_custom_asset_condition_requires_batching():
+ class CustomAsset(Asset):
+ @property
+ def requires_batching(self):
+ return True
+
+ condition = Asset("a") | CustomAsset("b")
+ with conf_vars({("scheduler", "batch_asset_events"): "False"}):
+ timetable = AssetTriggeredTimetable(condition)
+ assert timetable.batch_asset_events is True
+ timetable.validate()
+
+ with pytest.raises(AirflowTimetableInvalid,
match="batch_asset_events=True"):
+ AssetTriggeredTimetable(condition, batch_asset_events=False).validate()
