This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch v3-4-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit b8216f347e4918cede5e8b08e2a3d176558545cf Author: Rahul Vats <[email protected]> AuthorDate: Fri Oct 30 20:40:07 2026 +0530 Fix stale queued asset rows in scheduler cleanup and readiness (#76002) (#76039) (cherry picked from commit 9e662c23fc02739c13215bf1edddbbe12dc1b97f) Co-authored-by: Tzu-ping Chung <[email protected]> --- 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 8a29b83c2e6..bd9212973ed 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 + def _eager_load_dag_run_for_validation() -> tuple[LoaderOption, LoaderOption]: """ @@ -303,6 +307,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 ( @@ -313,16 +329,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. @@ -2795,9 +2801,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) @@ -2832,12 +2838,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, ) @@ -2857,45 +2863,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, @@ -2903,87 +2893,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, @@ -3010,9 +3011,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( @@ -3020,6 +3019,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: @@ -3037,68 +3041,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 df845b53a9f..f6aac156609 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 @@ -833,6 +835,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 c700605c24f..a2007dadee1 100644 --- a/airflow-core/tests/unit/assets/test_manager.py +++ b/airflow-core/tests/unit/assets/test_manager.py @@ -310,7 +310,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 ) @@ -331,7 +331,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 e808b722674..268dbfb2ce2 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -7081,11 +7081,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, @@ -7138,6 +7139,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]) @@ -7430,52 +7492,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") @@ -12854,9 +12910,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. @@ -12878,6 +12934,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) @@ -12902,8 +12960,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") @@ -12930,6 +12990,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) @@ -12946,52 +13008,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") @@ -14115,11 +14193,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 54b3571bd5f..a46b1c034e7 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -2807,6 +2807,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 ):
