This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 9e662c23fc0 Fix stale queued asset rows in scheduler cleanup and
readiness (#76002)
9e662c23fc0 is described below
commit 9e662c23fc02739c13215bf1edddbbe12dc1b97f
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Fri Oct 30 17:21:08 2026 +0800
Fix stale queued asset rows in scheduler cleanup and readiness (#76002)
---
airflow-core/src/airflow/assets/manager.py | 27 +-
.../src/airflow/jobs/scheduler_job_runner.py | 330 ++++++++++++---------
airflow-core/src/airflow/models/asset.py | 44 ++-
airflow-core/src/airflow/models/dag.py | 13 +-
airflow-core/tests/unit/assets/test_manager.py | 4 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 257 ++++++++++++----
airflow-core/tests/unit/models/test_dag.py | 45 +++
7 files changed, 481 insertions(+), 239 deletions(-)
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index eb8ea92d207..b732ad0de00 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -31,7 +31,6 @@ from airflow.listeners.listener import get_listener_manager
from airflow.listeners.types import AssetEvent as ListenerAssetEvent
from airflow.models.asset import (
AssetAliasModel,
- AssetDagRunQueue,
AssetEvent,
AssetModel,
AssetPartitionDagRun,
@@ -42,12 +41,12 @@ from airflow.models.asset import (
PartitionedAssetKeyLog,
TaskOutletAssetReference,
asset_alias_asset_event_association_table,
+ insert_asset_dag_run_queue_rows,
)
from airflow.models.log import Log
from airflow.timetables.base import compute_rollup_fingerprint
from airflow.utils.helpers import is_container, prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.sqlalchemy import get_dialect_name
if TYPE_CHECKING:
from datetime import datetime
@@ -840,29 +839,7 @@ class AssetManager(LoggingMixin):
{"asset_id": asset_id, "target_dag_id": dag.dag_id,
"asset_event_id": event.id}
for dag in dags_to_queue
]
-
- if (dialect_name := get_dialect_name(session)) == "mysql":
- from sqlalchemy.dialects.mysql import insert as my_insert
-
- session.execute(
- # The on-dup update a no-op since we are always working on
ADRQs
- # for the asset. It may not hold if you change e.g. ADRQ
schema.
-
my_insert(AssetDagRunQueue).on_duplicate_key_update(asset_id=AssetDagRunQueue.asset_id),
- values,
- )
- return
-
- if dialect_name == "postgresql":
- from sqlalchemy.dialects.postgresql import insert
- else:
- from sqlalchemy.dialects.sqlite import insert # type:
ignore[assignment]
-
- session.execute(
- insert(AssetDagRunQueue).on_conflict_do_nothing(
- index_elements=["target_dag_id", "asset_event_id"]
- ),
- values,
- )
+ insert_asset_dag_run_queue_rows(values, session=session)
def resolve_asset_manager() -> AssetManager:
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 6a5e74bcc0a..7c507276089 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -91,7 +91,8 @@ from airflow.models.asset import (
TaskInletAssetReference,
TaskOutletAssetReference,
alias_association_table,
- association_table,
+ build_consumed_by_dag_clause,
+ insert_asset_dag_run_queue_rows,
)
from airflow.models.asset_state_store import AssetStateStoreModel
from airflow.models.backfill import Backfill, BackfillDagRun
@@ -125,7 +126,7 @@ from airflow.ti_deps.dependencies_states import
ACTIVE_STATES, EXECUTION_STATES
from airflow.timetables.base import Timetable, compute_rollup_fingerprint
from airflow.triggers.base import TriggerEvent
from airflow.utils.event_scheduler import EventScheduler
-from airflow.utils.helpers import prune_dict
+from airflow.utils.helpers import chunks, prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.log.task_log_address import prepare_task_log_contexts
from airflow.utils.retries import MAX_DB_RETRIES, retry_db_transaction,
run_with_db_retries
@@ -146,7 +147,7 @@ if TYPE_CHECKING:
from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session
from sqlalchemy.orm.interfaces import LoaderOption
- from sqlalchemy.sql.elements import ColumnElement, SQLColumnExpression
+ from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.selectable import Select, Subquery
from airflow._shared.logging.types import Logger
@@ -173,6 +174,9 @@ TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT = "stuck in queued
reschedule"
# safety bound, not a behavioural knob operators need to tune.
MAX_PARTITION_DAG_RUNS_PER_LOOP = 500
+# Max number of ADRQ rows deleted per statement, to keep the bind parameter
count bounded.
+_ADRQ_DELETE_CHUNK_SIZE = 500
+
MAX_TI_REFILL_QUERIES_PER_LOOP = 3
"""
Max extra task instance queries per critical section to fill a partially
selected batch.
@@ -314,6 +318,18 @@ def _create_scheduler_dag_bag() -> CachedDBDagBag:
)
+def _build_scheduled_asset_event_clause(dag_id: str) -> ColumnElement[bool]:
+ """Check whether an asset event is of an asset the Dag schedules on,
directly or via an alias."""
+ return or_(
+ AssetEvent.asset_id.in_(
+
select(DagScheduleAssetReference.asset_id).where(DagScheduleAssetReference.dag_id
== dag_id)
+ ),
+ AssetEvent.source_aliases.any(
+
AssetAliasModel.scheduled_dags.any(DagScheduleAssetAliasReference.dag_id ==
dag_id)
+ ),
+ )
+
+
def _get_current_dr_task_concurrency(states: Iterable[TaskInstanceState]) ->
Subquery:
"""Get the dag_run IDs and how many tasks are in the provided states for
each one."""
return (
@@ -324,16 +340,6 @@ def _get_current_dr_task_concurrency(states:
Iterable[TaskInstanceState]) -> Sub
)
-def _consumed_by_dag(dag_id: str, event_id: SQLColumnExpression[int]) ->
ColumnElement[bool]:
- """Check whether the asset event *event_id* has been consumed by a run of
the Dag."""
- return (
- select(association_table.c.event_id)
- .join(DagRun, DagRun.id == association_table.c.dag_run_id)
- .where(DagRun.dag_id == dag_id, association_table.c.event_id ==
event_id)
- .exists()
- )
-
-
class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
"""
SchedulerJobRunner runs for a specific time interval and schedules jobs
that are ready to run.
@@ -2827,9 +2833,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# a later loop retries it.
if (gate := self._collect_gated_asset_events(dag=serdag,
session=session)) is None:
continue
- gated_max_event_id, gated_asset_events = gate
+ gated_locked_records, gated_asset_events = gate
else:
- gated_max_event_id, gated_asset_events = None, ()
+ gated_locked_records, gated_asset_events = None, ()
try:
next_info =
serdag.timetable.next_run_info_from_dag_model(dag_model=dag_model)
@@ -2864,12 +2870,12 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
session=session,
active_non_backfill_runs=active_runs_of_dags[dag_model.dag_id],
)
- if gated_max_event_id is not None:
+ if gated_locked_records is not None:
created_run.consumed_asset_events.extend(gated_asset_events)
session.flush() # Make consumed events visible to ADRQ
cleanup.
self._delete_consumed_asset_records(
- dag_id=dag_model.dag_id,
- max_event_id=gated_max_event_id,
+ serdag,
+ locked_records=gated_locked_records,
session=session,
)
@@ -2889,45 +2895,29 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
def _collect_gated_asset_events(
self, *, dag: SerializedDAG, session: Session
- ) -> tuple[int, Sequence[AssetEvent]] | None:
+ ) -> tuple[Sequence[AssetDagRunQueue], Sequence[AssetEvent]] | None:
"""
Check an asset-gated Dag's asset condition and collect what a new run
consumes.
- Returns the largest event ID among the locked ADRQ rows, and the events
- the new run consumes. Returns ``None`` when the condition is not
satisfied
- by the queued asset events, in which case no run should be created yet.
+ Returns the locked ADRQ rows, and the event objects consumed by the new
+ run. *None* is returned when the condition is not satisfied by the
queued
+ asset events, in which case no run should be created yet.
"""
- records = self._lock_queued_asset_records(dag_id=dag.dag_id,
load_assets=True, session=session)
+ records = self._lock_queued_asset_records(dag_id=dag.dag_id,
session=session)
if not records:
return None
- statuses = {SerializedAssetUniqueKey.from_asset(record.asset): True
for record in records}
- try:
- ready = AssetEvaluator(session).run(dag.timetable.asset_condition,
statuses=statuses)
- except Exception:
- self.log.exception("Dag '%s' failed to be evaluated; assuming not
ready", dag.dag_id)
- return None
- if not ready:
- return None
- max_event_id = max(record.asset_event_id for record in records)
- asset_events = self._select_consumed_asset_events(dag=dag,
max_event_id=max_event_id, session=session)
- if not asset_events:
- # The select returns every queued event this Dag has not consumed
- # yet; getting nothing back means all locked rows point to events
- # that were already consumed. Delete those stale rows.
- self._delete_consumed_asset_records(dag_id=dag.dag_id,
max_event_id=max_event_id, session=session)
+ if not self._is_asset_condition_met(dag, locked_records=records,
session=session):
+ self._delete_consumed_asset_records(dag, locked_records=records,
session=session)
return None
- return max_event_id, asset_events
+ return records, self._select_consumed_asset_events(dag,
locked_records=records, session=session)
- def _lock_queued_asset_records(
- self, *, dag_id: str, load_assets: bool, session: Session
- ) -> Sequence[AssetDagRunQueue]:
+ def _lock_queued_asset_records(self, *, dag_id: str, session: Session) ->
Sequence[AssetDagRunQueue]:
"""Lock and return the Dag's queued asset (ADRQ) rows, skipping rows
another scheduler holds."""
- query = select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id
== dag_id)
- if load_assets:
- query = query.options(joinedload(AssetDagRunQueue.asset))
return session.scalars(
with_row_locks(
- query,
+ select(AssetDagRunQueue)
+ .where(AssetDagRunQueue.target_dag_id == dag_id)
+ .options(joinedload(AssetDagRunQueue.asset)),
of=AssetDagRunQueue,
skip_locked=True,
key_share=False,
@@ -2935,87 +2925,98 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
).all()
+ def _is_asset_condition_met(
+ self, dag: SerializedDAG, *, locked_records:
Iterable[AssetDagRunQueue], session: Session
+ ) -> bool:
+ """
+ Evaluate a Dag's asset condition against its locked ADRQ rows.
+
+ Rows whose events have already been consumed by the Dag do not count.
+ """
+ consumed_event_ids = set(
+ session.scalars(
+ select(AssetDagRunQueue.asset_event_id).where(
+ AssetDagRunQueue.target_dag_id == dag.dag_id,
+ build_consumed_by_dag_clause(dag.dag_id,
AssetDagRunQueue.asset_event_id),
+ )
+ )
+ )
+ statuses = {
+ SerializedAssetUniqueKey.from_asset(record.asset): True
+ for record in locked_records
+ if record.asset_event_id not in consumed_event_ids
+ }
+ if not statuses:
+ return False
+ try:
+ return AssetEvaluator(session).run(dag.timetable.asset_condition,
statuses=statuses)
+ except Exception:
+ self.log.exception("Dag '%s' failed to be evaluated; assuming not
ready", dag.dag_id)
+ return False
+
def _select_consumed_asset_events(
self,
- *,
dag: SerializedDAG,
- max_event_id: int,
+ *,
+ locked_records: Iterable[AssetDagRunQueue],
limit: int | None = None,
session: Session,
) -> Sequence[AssetEvent]:
"""
Select unconsumed events referenced by a Dag's locked ADRQ rows.
- *max_event_id* is the largest event ID among the rows locked by this
- scheduler, so events queued after the lock was taken are left for the
- next scheduling loop. Concurrent schedulers can't consume this Dag's
rows
- since the Dag's own row is locked in ``DagModel.dags_needing_dagruns``.
+ ADRQ rows are matched up to the largest event ID in *locked_records*.
+ This does not exclude every event queued after the rows were locked.
+ With catchup, events are also matched by asset regardless of ID, and
+ event IDs are not assigned in commit order. Concurrent schedulers can't
+ consume this Dag's events since the Dag's own row is locked.
"""
event_predicate: ColumnElement[bool] = (
select(AssetDagRunQueue.asset_event_id)
.where(
AssetDagRunQueue.target_dag_id == dag.dag_id,
AssetDagRunQueue.asset_event_id == AssetEvent.id,
- AssetDagRunQueue.asset_event_id <= max_event_id,
+ AssetDagRunQueue.asset_event_id <= max(record.asset_event_id
for record in locked_records),
)
.exists()
)
if dag.catchup:
- event_predicate = or_(
- event_predicate,
- AssetEvent.asset_id.in_(
- select(DagScheduleAssetReference.asset_id).where(
- DagScheduleAssetReference.dag_id == dag.dag_id
- )
- ),
- AssetEvent.source_aliases.any(
-
AssetAliasModel.scheduled_dags.any(DagScheduleAssetAliasReference.dag_id ==
dag.dag_id)
- ),
- )
+ event_predicate |= _build_scheduled_asset_event_clause(dag.dag_id)
return session.scalars(
select(AssetEvent)
- .where(
- event_predicate,
- ~_consumed_by_dag(dag.dag_id, AssetEvent.id),
- )
+ .where(event_predicate, ~build_consumed_by_dag_clause(dag.dag_id,
AssetEvent.id))
.order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
.limit(limit)
).all()
def _delete_consumed_asset_records(
self,
+ dag: SerializedDAG,
*,
- dag_id: str,
- max_event_id: int,
- keep_event_id: int | None = None,
+ locked_records: Iterable[AssetDagRunQueue],
session: Session,
) -> None:
"""
- Delete ADRQ rows whose events have been consumed by a run of the Dag.
+ Delete locked ADRQ rows whose events have been consumed by a Dag run.
- Only consumed rows are deleted to avoid dropping newly queued events.
- Rows are matched against the events consumed by the Dag's runs (i.e.
- ``DagRun.consumed_asset_events``) with a subquery, so the caller must
- flush newly consumed events before calling this.
+ Only rows locked by this scheduler are deleted, since a DELETE waits on
+ rows locked elsewhere (e.g. by the API) and can deadlock. Of those,
only
+ consumed rows are deleted to keep events not yet consumed queued.
+ Consumed events are matched against the Dag's runs with a subquery, so
+ the caller must flush newly consumed events before calling this.
- :param max_event_id: Largest event ID among the ADRQ rows locked by
this
- scheduler. Rows queued after the lock was taken are not touched.
- :param keep_event_id: If given, the ADRQ row of this event is kept
even if
- the event is consumed.
+ :param locked_records: The ADRQ rows locked by this scheduler.
"""
- stmt = delete(AssetDagRunQueue).where(
- AssetDagRunQueue.target_dag_id == dag_id,
- AssetDagRunQueue.asset_event_id <= max_event_id,
- _consumed_by_dag(dag_id, AssetDagRunQueue.asset_event_id),
- )
- if keep_event_id is not None:
- stmt = stmt.where(AssetDagRunQueue.asset_event_id != keep_event_id)
- result = cast("CursorResult", session.execute(stmt))
- self.log.info(
- "Deleted %d ADRQ rows for '%s'",
- result.rowcount,
- dag_id,
- )
+ event_ids = [r.asset_event_id for r in locked_records]
+ deleted = 0
+ for chunk in chunks(event_ids, _ADRQ_DELETE_CHUNK_SIZE):
+ stmt = delete(AssetDagRunQueue).where(
+ AssetDagRunQueue.target_dag_id == dag.dag_id,
+ AssetDagRunQueue.asset_event_id.in_(chunk),
+ build_consumed_by_dag_clause(dag.dag_id,
AssetDagRunQueue.asset_event_id),
+ )
+ deleted += cast("CursorResult", session.execute(stmt)).rowcount
+ self.log.info("Deleted %d ADRQ rows for '%s'", deleted, dag.dag_id)
def _create_dag_runs_asset_triggered(
self,
@@ -3042,9 +3043,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
continue
- queued_adrqs = self._lock_queued_asset_records(
- dag_id=dag.dag_id, load_assets=False, session=session
- )
+ queued_adrqs = self._lock_queued_asset_records(dag_id=dag.dag_id,
session=session)
# If another scheduler already locked these ADRQ rows, SKIP LOCKED
makes this scheduler skip them.
if not queued_adrqs:
self.log.debug(
@@ -3052,6 +3051,11 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
dag.dag_id,
)
continue
+ # The condition was evaluated without locks when this Dag was
selected.
+ if not self._is_asset_condition_met(dag,
locked_records=queued_adrqs, session=session):
+ self.log.debug("Asset condition no longer met for Dag '%s'",
dag.dag_id)
+ self._delete_consumed_asset_records(dag,
locked_records=queued_adrqs, session=session)
+ continue
run_limit = per_dag_limit
if not dag.timetable.batch_asset_events:
@@ -3069,68 +3073,98 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if run_limit <= 0:
continue
- max_event_id = max(record.asset_event_id for record in
queued_adrqs)
asset_events = self._select_consumed_asset_events(
- dag=dag,
- max_event_id=max_event_id,
+ dag,
+ locked_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:
+ if has_more_events := not dag.timetable.batch_asset_events and
len(asset_events) > run_limit:
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(
- "Creating asset-triggered DagRun for '%s': %d queued
assets, triggered_date=%s",
- dag.dag_id,
- len(queued_adrqs),
- triggered_date,
- )
- team_name = (
- self._get_team_names_for_dag_ids([dag.dag_id],
session).get(dag.dag_id)
- if self._multi_team
- else None
+ triggered_date = timezone.coerce_datetime(max(event.timestamp for
event in asset_events))
+ self.log.debug(
+ "Creating asset-triggered DagRun for '%s': %d queued assets,
triggered_date=%s",
+ dag.dag_id,
+ len(queued_adrqs),
+ triggered_date,
+ )
+ team_name = (
+ self._get_team_names_for_dag_ids([dag.dag_id],
session).get(dag.dag_id)
+ if self._multi_team
+ else None
+ )
+ 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,
)
- 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),
- )
- session.flush() # Make consumed events visible to ADRQ
cleanup.
- else:
+ 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(
- "No DagRun created for '%s' - asset events already
consumed or none found",
+ "Created asset-triggered DagRun for '%s': run_id=%s,
consumed %d asset events",
dag.dag_id,
+ dag_run.run_id,
+ len(events),
)
- # Dequeue rows for events both just consumed above and any stale
- # ones. This runs even if no Dag run was created; otherwise those
- # rows would keep the Dag selected in every scheduling loop.
- self._delete_consumed_asset_records(
- dag_id=dag.dag_id,
- max_event_id=max_event_id,
- # Historical events may have no ADRQ. Keep a wakeup row until
that backlog drains.
- keep_event_id=max_event_id if has_more_events and dag.catchup
else None,
- session=session,
+ session.flush() # Make consumed events visible to ADRQ cleanup.
+ # Dequeue rows for events just consumed above and any stale ones.
+ self._delete_consumed_asset_records(dag,
locked_records=queued_adrqs, session=session)
+ if has_more_events and dag.catchup:
+ # Historical events may have no ADRQ row. Queue the next ones
so
+ # the Dag is selected again until the backlog drains.
+ self._queue_next_backlog_events(dag_id=dag.dag_id,
session=session)
+
+ def _queue_next_backlog_events(self, *, dag_id: str, session: Session) ->
None:
+ """
+ Queue an ADRQ row for the earliest unconsumed event of each asset the
Dag schedules on.
+
+ This lets the Dag's asset condition be evaluated against its catchup
+ backlog. Events already queued for the Dag are skipped, so the insert
+ does not wait on a row locked elsewhere (e.g. by the API) on MySQL.
+ """
+ ranked = (
+ select(
+ AssetEvent.id.label("event_id"),
+ AssetEvent.asset_id.label("asset_id"),
+ func.row_number()
+ .over(
+ partition_by=AssetEvent.asset_id,
+ order_by=(AssetEvent.timestamp.asc(), AssetEvent.id.asc()),
+ )
+ .label("event_rank"),
)
+ .where(
+ _build_scheduled_asset_event_clause(dag_id),
+ ~build_consumed_by_dag_clause(dag_id, AssetEvent.id),
+ )
+ .subquery()
+ )
+ queued = (
+ select(AssetDagRunQueue.asset_event_id)
+ .where(
+ AssetDagRunQueue.target_dag_id == dag_id,
+ AssetDagRunQueue.asset_event_id == ranked.c.event_id,
+ )
+ .exists()
+ )
+ values = [
+ {"asset_id": asset_id, "target_dag_id": dag_id, "asset_event_id":
event_id}
+ for event_id, asset_id in session.execute(
+ select(ranked.c.event_id,
ranked.c.asset_id).where(ranked.c.event_rank == 1, ~queued)
+ )
+ ]
+ if values:
+ insert_asset_dag_run_queue_rows(values, session=session)
def _lock_backfills(self, dag_runs: Collection[DagRun], session: Session)
-> dict[int, Backfill]:
"""
diff --git a/airflow-core/src/airflow/models/asset.py
b/airflow-core/src/airflow/models/asset.py
index 47a7bf5b311..4ddb46b00cd 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -44,13 +44,15 @@ from airflow._shared.timezones import timezone
from airflow.configuration import conf as airflow_conf
from airflow.models.base import Base, StringID
from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.utils.sqlalchemy import UtcDateTime
+from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name
if TYPE_CHECKING:
- from collections.abc import Iterable
+ from collections.abc import Iterable, Sequence
from typing import Any, TypeAlias
+ from sqlalchemy import Insert
from sqlalchemy.orm import Session
+ from sqlalchemy.sql.elements import ColumnElement, SQLColumnExpression
from airflow.models.dag import DagModel
from airflow.models.trigger import Trigger
@@ -811,6 +813,44 @@ association_table = Table(
)
+def build_consumed_by_dag_clause(
+ dag_id: str | SQLColumnExpression[str], event_id: SQLColumnExpression[int]
+) -> ColumnElement[bool]:
+ """Check whether the asset event *event_id* has been consumed by a run of
the Dag *dag_id*."""
+ from airflow.models.dagrun import DagRun # Avoid circular import
+
+ return (
+ select(association_table.c.event_id)
+ .join(DagRun, DagRun.id == association_table.c.dag_run_id)
+ .where(DagRun.dag_id == dag_id, association_table.c.event_id ==
event_id)
+ .exists()
+ )
+
+
+def insert_asset_dag_run_queue_rows(values: Sequence[dict[str, Any]], *,
session: Session) -> None:
+ """Insert ADRQ rows, ignoring rows that are already queued."""
+ stmt: Insert
+ if (dialect_name := get_dialect_name(session)) == "mysql":
+ from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+ # The on-dup update is a no-op since we are always working on ADRQs
+ # for the asset. It may not hold if you change e.g. ADRQ schema.
+ stmt =
mysql_insert(AssetDagRunQueue).on_duplicate_key_update(asset_id=AssetDagRunQueue.asset_id)
+ elif dialect_name == "postgresql":
+ from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+ stmt = pg_insert(AssetDagRunQueue).on_conflict_do_nothing(
+ index_elements=["target_dag_id", "asset_event_id"]
+ )
+ else:
+ from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+ stmt = sqlite_insert(AssetDagRunQueue).on_conflict_do_nothing(
+ index_elements=["target_dag_id", "asset_event_id"]
+ )
+ session.execute(stmt, values)
+
+
class AssetEvent(Base):
"""
A table to store assets events.
diff --git a/airflow-core/src/airflow/models/dag.py
b/airflow-core/src/airflow/models/dag.py
index 814f6e74298..434c18c2af4 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -62,7 +62,7 @@ from airflow._shared.timezones import timezone
from airflow.assets.evaluation import AssetEvaluator
from airflow.configuration import conf as airflow_conf
from airflow.exceptions import AirflowException
-from airflow.models.asset import AssetDagRunQueue
+from airflow.models.asset import AssetDagRunQueue, build_consumed_by_dag_clause
from airflow.models.base import Base, StringID
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
@@ -737,8 +737,11 @@ class DagModel(Base):
# this loads all the ADRQ records.... may need to limit num dags
adrq_by_dag: dict[str, list[AssetDagRunQueue]] = defaultdict(list)
- for adrq in session.scalars(
- select(AssetDagRunQueue).options(
+ for adrq, consumed in session.execute(
+ select(
+ AssetDagRunQueue,
+ build_consumed_by_dag_clause(AssetDagRunQueue.target_dag_id,
AssetDagRunQueue.asset_event_id),
+ ).options(
joinedload(AssetDagRunQueue.dag_model),
joinedload(AssetDagRunQueue.asset),
)
@@ -748,7 +751,9 @@ class DagModel(Base):
# could happen if the dag DID depend on an asset at some point,
# but no longer does. Delete the stale adrq.
session.delete(adrq)
- else:
+ elif not consumed:
+ # A row whose event the Dag already consumed is cleaned up when
+ # the Dag is next processed. It must not count towards
readiness.
adrq_by_dag[adrq.target_dag_id].append(adrq)
if adrq_by_dag:
diff --git a/airflow-core/tests/unit/assets/test_manager.py
b/airflow-core/tests/unit/assets/test_manager.py
index f940f203912..8192db17b0a 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -302,7 +302,7 @@ class TestAssetManager:
session = mock.MagicMock(spec=Session)
event = AssetEvent(asset_id=1)
event.id = 99
- with mock.patch("airflow.assets.manager.get_dialect_name",
return_value=dialect_name):
+ with mock.patch("airflow.models.asset.get_dialect_name",
return_value=dialect_name):
AssetManager._queue_dagruns_nonpartitioned(
asset_id=1, dags_to_queue={dag}, event=event, session=session
)
@@ -323,7 +323,7 @@ class TestAssetManager:
session = mock.MagicMock(spec=Session)
event = AssetEvent(asset_id=1)
event.id = 99
- with mock.patch("airflow.assets.manager.get_dialect_name",
return_value="mysql"):
+ with mock.patch("airflow.models.asset.get_dialect_name",
return_value="mysql"):
AssetManager._queue_dagruns_nonpartitioned(
asset_id=1, dags_to_queue={dag}, event=event, session=session
)
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index c2636b32764..1df2086c4a9 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -7309,11 +7309,12 @@ class TestSchedulerJob:
# outside the bound and left queued for the next scheduling loop.
pytest.param(0, {0}, {1}, id="skipped-row-newer"),
# The skipped row has a lower event ID than the locked row, so it
is
- # within the bound, and its event is consumed and dequeued as well.
- pytest.param(1, {0, 1}, set(), id="skipped-row-older"),
+ # within the bound and its event is consumed, but the row is left
for
+ # whoever holds it, since the cleanup only deletes rows it can
lock.
+ pytest.param(1, {0, 1}, {0}, id="skipped-row-older"),
],
)
- def test_create_dag_runs_asset_triggered_deletes_only_selected_adrq_rows(
+ def test_create_dag_runs_asset_triggered_deletes_only_locked_adrq_rows(
self,
session: Session,
dag_maker,
@@ -7366,6 +7367,67 @@ class TestSchedulerJob:
).all()
assert set(queued_asset_ids) == {asset_ids[i] for i in
expected_queued_indexes}
+ @pytest.mark.need_serialized_dag
+ @pytest.mark.backend("postgres", "mysql")
+ @pytest.mark.usefixtures("disable_load_example")
+ @pytest.mark.parametrize("gated", [False, True], ids=["triggered",
"gated"])
+ def test_create_dag_runs_asset_does_not_wait_on_rows_locked_elsewhere(
+ self, session: Session, dag_maker, gated
+ ):
+ assets = [Asset("asset-locked-elsewhere-1"),
Asset("asset-locked-elsewhere-2")]
+ logical_date = pendulum.datetime(2026, 3, 29, 17, tz="UTC")
+ condition = assets[0] | assets[1]
+ schedule = (
+ AssetAndTimeSchedule(
+ timetable=CronTriggerTimetable("* * * * *", timezone="UTC"),
assets=condition
+ )
+ if gated
+ else condition
+ )
+ with dag_maker(dag_id="asset-consumer-locked-elsewhere",
schedule=schedule, session=session):
+ EmptyOperator(task_id="dummy_task")
+ dag_model = dag_maker.dag_model
+ dag_id = dag_model.dag_id
+ events = []
+ for asset in assets:
+ asset_id =
session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.name))
+ event = AssetEvent(asset_id=asset_id, timestamp=logical_date)
+ session.add(event)
+ session.flush()
+ session.add(AssetDagRunQueue(asset_id=asset_id,
target_dag_id=dag_id, asset_event_id=event.id))
+ events.append(event)
+ if gated:
+ dag_model.next_dagrun = logical_date
+ dag_model.next_dagrun_data_interval = (logical_date, logical_date)
+ dag_model.next_dagrun_create_after = logical_date
+ session.commit()
+ runner = SchedulerJobRunner(Job(),
executors=[MockExecutor(do_update=False)])
+
+ def create_dag_runs():
+ with create_session(scoped=False) as other:
+ runner._create_dagruns_for_dags(cast("CommitProhibitorGuard",
other), other)
+
+ with ThreadPoolExecutor(max_workers=1) as pool,
create_session(scoped=False) as locker:
+ # Hold the older row, as a queued event DELETE endpoint would.
+ locker.execute(
+ delete(AssetDagRunQueue).where(
+ AssetDagRunQueue.target_dag_id == dag_id,
+ AssetDagRunQueue.asset_event_id == events[0].id,
+ )
+ )
+ future = pool.submit(create_dag_runs)
+ try:
+ future.result(timeout=10)
+ finally:
+ locker.rollback()
+
+ assert
session.scalar(select(func.count()).select_from(DagRun).where(DagRun.dag_id ==
dag_id))
+ # The cleanup removed only the row it could lock.
+ queued_event_ids = session.scalars(
+
select(AssetDagRunQueue.asset_event_id).where(AssetDagRunQueue.target_dag_id ==
dag_id)
+ ).all()
+ assert queued_event_ids == [events[0].id]
+
@pytest.mark.need_serialized_dag
@pytest.mark.backend("postgres", "mysql")
@pytest.mark.parametrize("batch_asset_events", [True, False])
@@ -7658,52 +7720,46 @@ class TestSchedulerJob:
assert all(adrq.target_dag_id == "consumer" for adrq in adrqs)
@pytest.mark.need_serialized_dag
- def test_no_create_dag_runs_when_asset_event_already_consumed(self,
session: Session, dag_maker, caplog):
- asset = Asset(name="test_asset")
- with dag_maker(dag_id="consumer", schedule=asset, session=session):
+ @pytest.mark.parametrize("and_condition", [False, True],
ids=["single-asset", "and-condition"])
+ def test_no_create_dag_runs_when_asset_event_already_consumed(
+ self, session: Session, dag_maker, and_condition
+ ):
+ assets = [Asset(name="asset_a"), Asset(name="asset_b")] if
and_condition else [Asset(name="asset_a")]
+ with dag_maker(dag_id="consumer", schedule=assets, session=session):
pass
dag_model = dag_maker.dag_model
- asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri
== asset.uri))
-
- # An event that has already been consumed by an earlier DagRun of this
dag.
- event = AssetEvent(asset_id=asset_id, timestamp=timezone.utcnow())
- session.add(event)
+ events = [
+ AssetEvent(
+
asset_id=session.scalar(select(AssetModel.id).where(AssetModel.uri ==
asset.uri)),
+ timestamp=timezone.utcnow(),
+ )
+ for asset in assets
+ ]
+ session.add_all(events)
session.flush()
prior_run = dag_maker.create_dagrun(
run_id="prior-consuming-run",
logical_date=DEFAULT_DATE,
data_interval=(DEFAULT_DATE, DEFAULT_DATE),
)
- prior_run.consumed_asset_events.append(event)
- session.flush()
-
- # A stale ADRQ that still references the already-consumed event. It
should be cleaned up
- # even when no new DagRun is created, to prevent stale ADRQ rows from
accumulating and
- # causing infinite scheduler loops.
- session.add(
- AssetDagRunQueue(asset_id=asset_id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id)
+ # The first event was already consumed by an earlier run, but is still
queued.
+ prior_run.consumed_asset_events.append(events[0])
+ session.add_all(
+ AssetDagRunQueue(asset_id=event.asset_id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id)
+ for event in events
)
session.flush()
- with caplog.at_level("INFO"):
- scheduler_job = Job()
- self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[MockExecutor(do_update=False)])
- self.job_runner._create_dag_runs_asset_triggered(
- dag_models=[dag_model],
- session=session,
- )
+ self.job_runner = SchedulerJobRunner(job=Job(),
executors=[MockExecutor(do_update=False)])
+
self.job_runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
- assert "No DagRun created" in caplog.text
- # No *new* DagRun is created; only the pre-existing consuming run
remains.
runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)).all()
assert [r.run_id for r in runs] == ["prior-consuming-run"]
- _adrq = session.scalars(
- select(AssetDagRunQueue).where(
- AssetDagRunQueue.asset_id == asset_id,
AssetDagRunQueue.target_dag_id == dag_model.dag_id
- )
- ).one_or_none()
- # ADRQ is deleted even when no DagRun is created, to prevent stale
rows accumulating.
- assert _adrq is None
+ # The stale row is dequeued even though no run is created; a fresh one
stays.
+ queued = session.scalars(
+
select(AssetDagRunQueue.asset_event_id).where(AssetDagRunQueue.target_dag_id ==
dag_model.dag_id)
+ ).all()
+ assert queued == [event.id for event in events[1:]]
@time_machine.travel(DEFAULT_DATE + datetime.timedelta(days=1, seconds=9),
tick=False)
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
@@ -13082,9 +13138,9 @@ def _queue_many_asset_events(
@pytest.mark.need_serialized_dag
@pytest.mark.parametrize("batch_asset_events", [False, True])
def test_create_dagruns_asset_triggered_many_queued_events_no_literal_in(
- session: Session, dag_maker, batch_asset_events: bool
+ session: Session, dag_maker, monkeypatch, batch_asset_events: bool
):
- """Queued events must not be sent to the database as a literal list
(#75865)."""
+ """Queued events must not be sent to the database as a literal list."""
asset_1 = Asset(uri="test://many-events-1", name="many-events-1")
asset_2 = Asset(uri="test://many-events-2", name="many-events-2")
max_active_runs = 4 # Below the default per-loop run budget, so it is the
effective limit.
@@ -13106,6 +13162,8 @@ def
test_create_dagruns_asset_triggered_many_queued_events_no_literal_in(
session=session,
)
+ # Smaller than the number of queued events, so the cleanup spans several
chunks.
+
monkeypatch.setattr("airflow.jobs.scheduler_job_runner._ADRQ_DELETE_CHUNK_SIZE",
50)
runner = SchedulerJobRunner(job=Job(),
executors=[MockExecutor(do_update=False)])
with _capture_max_bind_params(session) as param_counts:
runner._create_dag_runs_asset_triggered(dag_models=[dag_model],
session=session)
@@ -13130,8 +13188,10 @@ def
test_create_dagruns_asset_triggered_many_queued_events_no_literal_in(
@pytest.mark.usefixtures("disable_load_example")
@pytest.mark.need_serialized_dag
-def test_create_dagruns_asset_gated_many_queued_events_no_literal_in(session:
Session, dag_maker):
- """Queued events must not be sent to the database as a literal list
(#75865)."""
+def test_create_dagruns_asset_gated_many_queued_events_no_literal_in(
+ session: Session, dag_maker, monkeypatch
+):
+ """Queued events must not be sent to the database as a literal list."""
asset_1 = Asset(uri="test://many-events-gated-1",
name="many-events-gated-1")
asset_2 = Asset(uri="test://many-events-gated-2",
name="many-events-gated-2")
logical_date = pendulum.datetime(2026, 3, 29, 17, tz="UTC")
@@ -13158,6 +13218,8 @@ def
test_create_dagruns_asset_gated_many_queued_events_no_literal_in(session: Se
session=session,
)
+ # Smaller than the number of queued events, so the cleanup spans several
chunks.
+
monkeypatch.setattr("airflow.jobs.scheduler_job_runner._ADRQ_DELETE_CHUNK_SIZE",
50)
runner = SchedulerJobRunner(job=Job(), executors=[MockExecutor()])
with _capture_max_bind_params(session) as param_counts:
runner._create_dagruns_for_dags(cast("CommitProhibitorGuard",
session), session)
@@ -13174,52 +13236,68 @@ def
test_create_dagruns_asset_gated_many_queued_events_no_literal_in(session: Se
@pytest.mark.usefixtures("disable_load_example")
@pytest.mark.need_serialized_dag
-def test_create_dagruns_asset_gated_dequeues_already_consumed_events(session:
Session, dag_maker):
- asset = Asset(uri="test://asset-gated-consumed",
name="asset-gated-consumed")
[email protected]("and_condition", [False, True], ids=["single-asset",
"and-condition"])
+def test_create_dagruns_asset_gated_dequeues_already_consumed_events(
+ session: Session, dag_maker, and_condition
+):
+ assets = [
+ Asset(uri=f"test://asset-gated-consumed-{suffix}",
name=f"asset-gated-consumed-{suffix}")
+ for suffix in (("a", "b") if and_condition else ("a",))
+ ]
logical_date = pendulum.datetime(2026, 3, 29, 17, tz="UTC")
with dag_maker(
dag_id="asset-gated-consumed",
schedule=AssetAndTimeSchedule(
timetable=CronTriggerTimetable("* * * * *", timezone="UTC"),
- assets=[asset],
+ assets=assets,
),
session=session,
):
EmptyOperator(task_id="dummy_task")
dag_model = dag_maker.dag_model
-
- # An event already consumed by an earlier run of this Dag...
- asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri ==
asset.uri))
- event = AssetEvent(asset_id=asset_id, timestamp=logical_date)
- session.add(event)
+ events = [
+ AssetEvent(
+ asset_id=session.scalar(select(AssetModel.id).where(AssetModel.uri
== asset.uri)),
+ timestamp=logical_date,
+ )
+ for asset in assets
+ ]
+ session.add_all(events)
session.flush()
prior_run = dag_maker.create_dagrun(
run_id="prior-consuming-run",
logical_date=logical_date - timedelta(hours=1),
data_interval=(logical_date - timedelta(hours=1), logical_date -
timedelta(hours=1)),
)
- prior_run.consumed_asset_events.append(event)
- # ... but still queued. This makes the asset condition look satisfied.
- session.add(AssetDagRunQueue(asset_id=asset_id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id))
+ # The first event was already consumed by an earlier run, but is still
queued.
+ prior_run.consumed_asset_events.append(events[0])
+ session.add_all(
+ AssetDagRunQueue(asset_id=event.asset_id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id)
+ for event in events
+ )
dag_model.next_dagrun = logical_date
dag_model.next_dagrun_data_interval = (logical_date, logical_date)
dag_model.next_dagrun_create_after = logical_date
session.flush()
- SchedulerJobRunner(job=Job(),
executors=[MockExecutor()])._create_dagruns_for_dags(
- cast("CommitProhibitorGuard", session), session
- )
+ runner = SchedulerJobRunner(job=Job(), executors=[MockExecutor()])
+ runner._create_dagruns_for_dags(cast("CommitProhibitorGuard", session),
session)
session.flush()
- # No run is created, the schedule slot is kept, and the stale row is
dequeued.
+ # No run is created and the schedule slot is kept. The stale row does not
+ # make the Dag ready, so it stays queued...
runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)).all()
assert [run.run_id for run in runs] == ["prior-consuming-run"]
session.refresh(dag_model)
assert dag_model.next_dagrun == logical_date
- assert (
-
session.scalar(select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
dag_model.dag_id))
- is None
- )
+ query =
select(AssetDagRunQueue.asset_event_id).where(AssetDagRunQueue.target_dag_id ==
dag_model.dag_id)
+ assert set(session.scalars(query)) == {event.id for event in events}
+
+ # ... until the Dag is processed. A fresh row stays.
+ serdag = runner._get_current_dag(dag_id=dag_model.dag_id, session=session)
+ assert serdag is not None
+ assert runner._collect_gated_asset_events(dag=serdag, session=session) is
None
+ assert set(session.scalars(query)) == {event.id for event in events[1:]}
@time_machine.travel("2026-03-29 19:00:00+00:00")
@@ -14343,11 +14421,74 @@ def
test_unbatched_asset_runs_drain_backlog_within_limits(
)
session.flush()
runs = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_id).order_by(DagRun.id)).all()
+ # Every queued row points at an event not consumed yet.
+ consumed_ids = {event.id for run in runs for event in
run.consumed_asset_events}
+ assert
consumed_ids.isdisjoint(session.scalars(select(AssetDagRunQueue.asset_event_id)))
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
+def
test_unbatched_catchup_backlog_drains_with_multi_asset_condition(dag_maker,
session, monkeypatch):
+ monkeypatch.setattr(DagModel, "NUM_DAGS_PER_DAGRUN_QUERY", 1)
+ # Asset A is scheduled on through an alias, so its backlog can only be
matched via the alias.
+ asset_a_model = AssetModel(name="backlog-a", uri="test://backlog-a",
group="asset")
+ alias_model = AssetAliasModel(name="backlog-a-alias", group="asset")
+ alias_model.assets.append(asset_a_model)
+ session.add_all([asset_a_model, alias_model])
+ session.flush()
+ asset_b = Asset(name="backlog-b")
+ dag_id = "backlog-and-consumer"
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=AssetTriggeredTimetable(
+ assets=AssetAlias(name=alias_model.name, group=alias_model.group)
& asset_b,
+ batch_asset_events=False,
+ ),
+ catchup=True,
+ session=session,
+ ):
+ EmptyOperator(task_id="consume")
+ session.commit()
+ a_id = asset_a_model.id
+ b_id = session.scalar(select(AssetModel.id).where(AssetModel.uri ==
asset_b.uri))
+ assert set(
+ session.scalars(
+
select(DagScheduleAssetReference.asset_id).where(DagScheduleAssetReference.dag_id
== dag_id)
+ )
+ ) == {b_id}
+ timestamp = timezone.parse("2026-10-01T00:00:00Z")
+ a1, b1, a2, b2 = events = [
+ AssetEvent(asset_id=asset_id, timestamp=timestamp +
timedelta(minutes=i))
+ for i, asset_id in enumerate([a_id, b_id, a_id, b_id])
+ ]
+ for event in (a1, a2):
+ event.source_aliases.append(alias_model)
+ session.add_all(events)
+ session.flush()
+ # Only the first event of each asset is queued; the rest is backlog.
+ session.add_all(
+ AssetDagRunQueue(asset_id=event.asset_id, target_dag_id=dag_id,
asset_event_id=event.id)
+ for event in (a1, b1)
+ )
+ session.flush()
+
+ runner = SchedulerJobRunner(job=Job(),
executors=[MockExecutor(do_update=False)])
+ for _ in range(4):
+ 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()
+
+ # The backlog drains one event per run while the condition holds. The last
+ # event of B waits for a new event of A, as it would without catchup.
+ 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] == [[a1], [b1], [a2]]
+ assert session.scalars(select(AssetDagRunQueue.asset_event_id)).all() ==
[b2.id]
+
+
@pytest.mark.need_serialized_dag
@mock.patch.object(DagModel, "NUM_DAGS_PER_DAGRUN_QUERY", 3)
@pytest.mark.parametrize("batch_first_dag", [False, True])
diff --git a/airflow-core/tests/unit/models/test_dag.py
b/airflow-core/tests/unit/models/test_dag.py
index 9c6b6754db6..0bb313664e9 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -2793,6 +2793,51 @@ class TestDagModel:
dag_models = query.all()
assert dag_models == [dag_model]
+ def test_dags_needing_dagruns_ignores_consumed_adrq(self, dag_maker,
session):
+ asset_a = Asset(uri="test://asset-a", group="test-group")
+ asset_b = Asset(uri="test://asset-b", group="test-group")
+ with dag_maker(
+ session=session,
+ dag_id="my_dag",
+ schedule=[asset_a, asset_b],
+ start_date=pendulum.now().add(days=-2),
+ ):
+ EmptyOperator(task_id="dummy")
+ dag_model = dag_maker.dag_model
+ a_id, b_id = (
+ session.scalar(select(AssetModel.id).where(AssetModel.uri ==
asset.uri))
+ for asset in (asset_a, asset_b)
+ )
+
+ def _queue_event(asset_id):
+ event = AssetEvent(asset_id=asset_id, timestamp=timezone.utcnow())
+ session.add(event)
+ session.flush()
+ session.add(
+ AssetDagRunQueue(asset_id=asset_id,
target_dag_id=dag_model.dag_id, asset_event_id=event.id)
+ )
+ session.flush()
+ return event
+
+ # A stale row for A (event already consumed) and a fresh event for B
+ # do not satisfy the condition.
+ prior_run = dag_maker.create_dagrun(
+ run_type=DagRunType.ASSET_TRIGGERED,
+ state=DagRunState.SUCCESS,
+ logical_date=pendulum.now("UTC"),
+ )
+ prior_run.consumed_asset_events.append(_queue_event(a_id))
+ _queue_event(b_id)
+ query, triggered_date_by_dag = DagModel.dags_needing_dagruns(session)
+ assert query.all() == []
+ assert triggered_date_by_dag == {}
+
+ # A fresh event for A does.
+ _queue_event(a_id)
+ query, triggered_date_by_dag = DagModel.dags_needing_dagruns(session)
+ assert query.all() == [dag_model]
+ assert set(triggered_date_by_dag) == {dag_model.dag_id}
+
def test_dags_needing_dagruns_skips_adrq_when_serialized_dag_missing(
self, session, caplog, testing_dag_bundle
):
