From 7a0676394f9c934d32f14afe276e6d81c92c8e29 Mon Sep 17 00:00:00 2001 From: TP Date: Mon, 3 Aug 2026 14:22:37 +0800 Subject: [PATCH] Reference asset events from asset_dag_run_queue Asset-triggered scheduling had the scheduler reconstruct which AssetEvents fed a run from a created_at time window. The prior fix 9e5b7701f8 committed each AssetEvent in separate short-lived sessions so it was visible early, but that broke atomicity. This commit changes the ADRQ to directly reference the asset event that generates it, instead of using an implicit asset_id + timestamp combination. Atomicity is restored since we no longer needs to commit separately to get the timestamps right, while still being able to reference the correct asset events from the ADRQ. --- airflow-core/docs/migrations-ref.rst | 5 +- airflow-core/src/airflow/assets/manager.py | 128 ++-------- .../src/airflow/jobs/scheduler_job_runner.py | 87 +++---- .../0128_3_4_0_add_asset_event_id_to_adrq.py | 196 ++++++++++++++ airflow-core/src/airflow/models/asset.py | 13 +- airflow-core/src/airflow/utils/db.py | 2 +- .../core_api/routes/public/test_assets.py | 11 +- .../core_api/routes/ui/test_assets.py | 22 +- .../tests/unit/assets/test_manager.py | 90 +++++-- .../tests/unit/jobs/test_scheduler_job.py | 240 ++++++++++++------ airflow-core/tests/unit/models/test_dag.py | 57 ++++- .../tests/unit/models/test_taskinstance.py | 17 +- .../unit/timetables/test_assets_timetable.py | 5 +- 13 files changed, 584 insertions(+), 289 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index dfc96be79019e..e3b33cab18591 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +| ``b2f1a9c7d4e0`` (head) | ``7a98f1b7dbd3`` | ``3.4.0`` | Reference the asset event from asset_dag_run_queue (consume- | +| | | | by-reference). | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``7a98f1b7dbd3`` | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index d0f8470b17c94..67de46585ddfd 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -46,7 +46,6 @@ 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.session import create_session from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks if TYPE_CHECKING: @@ -120,38 +119,6 @@ def _lock_asset_model( yield -def _create_asset_event(*, session: Session, **event_kwargs) -> AssetEvent: - """ - Persist an :class:`AssetEvent` row and return it, bound to *session*. - - On SQLite the event is added directly to the caller's *session* and - flushed. SQLite serialises writes at the database-file level: opening - a second connection here would compete with any write locks the - caller's transaction already holds (for example, an UPDATE on - ``dag_run`` flushed earlier in ``register_asset_changes_in_db``) and - deadlock with ``database is locked``. - - On Postgres/MySQL a short-lived independent session is used so the - row is committed — and therefore visible to the scheduler's session - via MVCC — before the caller continues. The committed row is then - re-loaded into the caller's *session* so subsequent relationship - operations work correctly. - """ - if get_dialect_name(session) == "sqlite": - asset_event = AssetEvent(**event_kwargs) - session.add(asset_event) - session.flush() - return asset_event - - with create_session(scoped=False) as ae_session: - asset_event = AssetEvent(**event_kwargs) - ae_session.add(asset_event) - ae_session.flush() - asset_event_id = asset_event.id - - return session.get_one(AssetEvent, asset_event_id) - - class AssetManager(LoggingMixin): """ A pluggable class that manages operations for assets. @@ -374,7 +341,10 @@ def register_asset_change( source_run_id=task_instance.run_id, source_map_index=task_instance.map_index, ) - asset_event = _create_asset_event(session=session, **event_kwargs) + + asset_event = AssetEvent(**event_kwargs) + session.add(asset_event) + session.flush() dags_to_queue_from_asset = {ref.dag for ref in asset_model.scheduled_dags if not ref.dag.is_paused} @@ -544,25 +514,7 @@ def _queue_dagruns( if not non_partitioned_dags or partition_key is not None: return None - # Possible race condition: if multiple dags or multiple (usually - # mapped) tasks update the same asset, this can fail with a unique - # constraint violation. - # - # Where the dialect supports a single-statement "insert, update on - # conflict" we use it; it is atomic, avoids the per-row SAVEPOINT churn, - # and holds locks for far less time (which on MySQL/InnoDB also makes the - # concurrent fan-out much less deadlock-prone). Otherwise we "fallback" to - # a nested transaction per row. Either way the rows are added in the same - # transaction where `ti.state` is changed. - dialect_name = get_dialect_name(session) - if TYPE_CHECKING: - assert dialect_name is not None - if dialect_name == "mysql": - return cls._queue_dagruns_nonpartitioned_mysql(asset_id, non_partitioned_dags, event, session) - # PostgreSQL and SQLite both support ON CONFLICT DO UPDATE. - return cls._queue_dagruns_nonpartitioned_conflict_update( - asset_id, non_partitioned_dags, event, session, dialect_name - ) + return cls._queue_dagruns_nonpartitioned(asset_id, non_partitioned_dags, event, session) @classmethod def _queue_partitioned_dags( @@ -827,71 +779,33 @@ def _get_or_create_apdr( return apdr @classmethod - def _queue_dagruns_nonpartitioned_slow_path( + def _queue_dagruns_nonpartitioned( cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session ) -> None: - def _queue_dagrun_if_needed(dag: DagModel) -> str | None: - item = AssetDagRunQueue(target_dag_id=dag.dag_id, asset_id=asset_id, created_at=event.timestamp) - # Don't error whole transaction when a single RunQueue item conflicts. - # https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint - try: - with session.begin_nested(): - existing = session.get( - AssetDagRunQueue, {"target_dag_id": dag.dag_id, "asset_id": asset_id} - ) - if existing and existing.created_at >= event.timestamp: - cls.logger().debug("Skipping record %s due to newer timestamp", item) - return dag.dag_id # already queued with a newer timestamp - session.merge(item) - except exc.IntegrityError: - cls.logger().debug("Skipping record %s", item, exc_info=True) - return dag.dag_id - - queued_results = (_queue_dagrun_if_needed(dag) for dag in dags_to_queue) - if queued_dag_ids := [r for r in queued_results if r is not None]: - cls.logger().debug("consuming dag ids %s", queued_dag_ids) - - @classmethod - def _queue_dagruns_nonpartitioned_mysql( - cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session - ) -> None: - from sqlalchemy import case - from sqlalchemy.dialects.mysql import insert + if not dags_to_queue: + return + values = [ + {"asset_id": asset_id, "target_dag_id": dag.dag_id, "asset_event_id": event.id} + for dag in dags_to_queue + ] - values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] - stmt = insert(AssetDagRunQueue).values(asset_id=asset_id, created_at=event.timestamp) + if (dialect_name := get_dialect_name(session)) == "mysql": + from sqlalchemy.dialects.mysql import insert as my_insert - update_stmt = stmt.on_duplicate_key_update( - created_at=case( - (stmt.inserted.created_at >= AssetDagRunQueue.created_at, stmt.inserted.created_at), - else_=AssetDagRunQueue.created_at, - ) - ) - session.execute(update_stmt, values) + session.execute(my_insert(AssetDagRunQueue).prefix_with("IGNORE"), values) + return - @classmethod - def _queue_dagruns_nonpartitioned_conflict_update( - cls, - asset_id: int, - dags_to_queue: set[DagModel], - event: AssetEvent, - session: Session, - dialect_name: str, - ) -> None: - """Handle ON CONFLICT DO UPDATE upsert for dialects that support it (postgresql, sqlite).""" if dialect_name == "postgresql": from sqlalchemy.dialects.postgresql import insert else: from sqlalchemy.dialects.sqlite import insert # type: ignore[assignment] - values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue] - stmt = insert(AssetDagRunQueue).values(asset_id=asset_id, created_at=event.timestamp) - update_stmt = stmt.on_conflict_do_update( - index_elements=["asset_id", "target_dag_id"], - set_={"created_at": stmt.excluded.created_at}, - where=(AssetDagRunQueue.created_at < stmt.excluded.created_at), + session.execute( + insert(AssetDagRunQueue).on_conflict_do_nothing( + index_elements=["target_dag_id", "asset_event_id"] + ), + values, ) - session.execute(update_stmt, values) 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 b6a29995d5937..40d1651505794 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -28,7 +28,7 @@ from collections import Counter, defaultdict, deque from collections.abc import Callable, Collection, Iterable, Iterator from contextlib import ExitStack -from datetime import date, datetime, timedelta +from datetime import datetime, timedelta from functools import lru_cache, partial from itertools import groupby from typing import TYPE_CHECKING, Any, cast @@ -135,10 +135,10 @@ if TYPE_CHECKING: from types import FrameType - from pendulum.datetime import DateTime from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from sqlalchemy.orm.interfaces import LoaderOption + from sqlalchemy.sql.elements import ColumnElement from sqlalchemy.sql.selectable import Subquery from airflow._shared.logging.types import Logger @@ -2641,9 +2641,7 @@ def _create_dag_runs_asset_triggered( queued_adrqs = session.scalars( with_row_locks( - select(AssetDagRunQueue) - .where(AssetDagRunQueue.target_dag_id == dag.dag_id) - .order_by(AssetDagRunQueue.created_at.desc()), + select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == dag.dag_id), of=AssetDagRunQueue, skip_locked=True, key_share=False, @@ -2658,53 +2656,31 @@ def _create_dag_runs_asset_triggered( ) continue - triggered_date: DateTime = timezone.coerce_datetime(queued_adrqs[0].created_at) - self.log.debug( - "Creating asset-triggered DagRun for '%s': %d queued assets, triggered_date=%s", - dag.dag_id, - len(queued_adrqs), - triggered_date, - ) - cte = ( - select(func.max(DagRun.run_after).label("previous_dag_run_run_after")) - .where( - DagRun.dag_id == dag.dag_id, - DagRun.run_type == DagRunType.ASSET_TRIGGERED, - DagRun.run_after < triggered_date, - ) - .cte() - ) - - # A first asset-triggered run has no previous run to floor the event window. With - # catchup off, floor it at when the Dag started scheduling on its assets so the - # backlog is skipped; with catchup on, only date.min applies and the backlog replays. - event_window_floor: list[Any] = [cte.c.previous_dag_run_run_after] - if not dag.catchup: - event_window_floor.append( - select(func.min(DagScheduleAssetReference.created_at)) - .where(DagScheduleAssetReference.dag_id == dag.dag_id) - .scalar_subquery() + referenced_event_ids = {adrq.asset_event_id for adrq in queued_adrqs} + event_predicate: ColumnElement[bool] = AssetEvent.id.in_(referenced_event_ids) + if dag.catchup: + # With catchup on, also consume events recorded before the Dag started + # scheduling on its assets/aliases, not just those with a queue row. (With catchup + # off only queued events are consumed.) The not-consumed filter below dedupes + # across runs, so no event window is needed. + 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_window_floor.append(date.min) - asset_events = list( session.scalars( select(AssetEvent) .where( - or_( - 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 - ) - ), - ), - AssetEvent.timestamp > func.coalesce(*event_window_floor), - AssetEvent.timestamp <= triggered_date, + event_predicate, ~( select(association_table.c.event_id) .join(DagRun, DagRun.id == association_table.c.dag_run_id) @@ -2719,6 +2695,13 @@ def _create_dag_runs_asset_triggered( ) ) 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, + ) dag_run = dag.create_dagrun( run_id=DagRun.generate_run_id( run_type=DagRunType.ASSET_TRIGGERED, logical_date=None, run_after=triggered_date @@ -2747,19 +2730,19 @@ def _create_dag_runs_asset_triggered( ) else: self.log.info( - "No DagRun created for '%s' at '%s' - asset events already consumed or none found", + "No DagRun created for '%s' - asset events already consumed or none found", dag.dag_id, - triggered_date, ) # Always delete ADRQ rows for this batch to prevent stale entries accumulating, # including when all events were already consumed by a concurrent DagRun. - adrq_pks = [(record.asset_id, record.target_dag_id) for record in queued_adrqs] result = cast( "CursorResult", session.execute( delete(AssetDagRunQueue).where( - tuple_(AssetDagRunQueue.asset_id, AssetDagRunQueue.target_dag_id).in_(adrq_pks), - AssetDagRunQueue.created_at <= triggered_date, + tuple_( + AssetDagRunQueue.target_dag_id, + AssetDagRunQueue.asset_event_id, + ).in_((adrq.target_dag_id, adrq.asset_event_id) for adrq in queued_adrqs) ) ), ) diff --git a/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py new file mode 100644 index 0000000000000..47fd9258e9fec --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py @@ -0,0 +1,196 @@ +# +# 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. + +""" +Reference the asset event from asset_dag_run_queue (consume-by-reference). + +Table ``asset_dag_run_queue`` gets a new column ``asset_event_id``, and the +primary key is made ``(target_dag_id, asset_event_id)``, so the scheduler +consumes queued asset events by reference instead of by a ``created_at`` time +window. ``asset_id`` is kept as a denormalized column. + +Existing rows are coalesced (one per ``(asset_id, target_dag_id)``) and carry no +event reference. Rather than dropping them (which would silently skip pending +asset-triggered Dag runs), the pre-migration scheduler's own consumption window is +replayed per dag to rebuild per-event rows: + + triggered_date = + MAX(asset_dag_run_queue.created_at) (per dag) + floor = + MAX(dag_run.run_after) for asset-triggered runs of the dag with + run_after < triggered_date (per dag) + contributing events = + the queued asset's events with + floor < asset_event.timestamp <= triggered_date + +The expansion is staged in a side table, the queue is cleared, the (now empty) +table is reshaped, and the staged rows are inserted back. + +Revision ID: b2f1a9c7d4e0 +Revises: 7a98f1b7dbd3 +Create Date: 2026-08-03 12:00:00.000000 +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from textwrap import dedent + +import sqlalchemy as sa +from alembic import context, op + +# revision identifiers, used by Alembic. +revision = "b2f1a9c7d4e0" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +_STAGING = "_adrq_migration_staging" + +_STAGE_SQL = f""" +CREATE TABLE {_STAGING} AS +SELECT DISTINCT + adrq.target_dag_id AS target_dag_id, + adrq.asset_id AS asset_id, + ae.id AS asset_event_id +FROM asset_dag_run_queue adrq +JOIN ( + SELECT t.target_dag_id, + t.triggered_date, + ( + SELECT MAX(dr.run_after) + FROM dag_run dr + WHERE dr.dag_id = t.target_dag_id + AND dr.run_type = 'asset_triggered' + AND dr.run_after < t.triggered_date + ) AS floor_date + FROM ( + SELECT target_dag_id, MAX(created_at) AS triggered_date + FROM asset_dag_run_queue + GROUP BY target_dag_id + ) t +) td ON td.target_dag_id = adrq.target_dag_id +JOIN asset_event ae + ON ae.asset_id = adrq.asset_id + AND ae.timestamp <= td.triggered_date + AND ae.timestamp > COALESCE(td.floor_date, :floor_min) +""" + + +def upgrade(): + """Directly reference the asset event from asset_dag_run_queue.""" + # 1. Add the reference column, nullable for now; it is made NOT NULL after the rebuild. + op.add_column("asset_dag_run_queue", sa.Column("asset_event_id", sa.Integer(), nullable=True)) + + # 2. Stage the faithful per-event expansion, then clear the coalesced rows. + if context.is_offline_mode(): + print( + dedent(""" + ------------ + -- WARNING: asset_dag_run_queue cannot be rebuilt in offline mode; + -- any pending (unprocessed) queued asset events will be dropped. + ------------ + """) + ) + op.execute("DELETE FROM asset_dag_run_queue") + else: + conn = op.get_bind() + floor_min = datetime(1970, 1, 1, tzinfo=timezone.utc) + conn.execute(sa.text(_STAGE_SQL), {"floor_min": floor_min}) + conn.execute(sa.text("DELETE FROM asset_dag_run_queue")) + + # 3. Make the reference NOT NULL, move the PK, add the FKs. + with op.batch_alter_table("asset_dag_run_queue") as batch_op: + batch_op.alter_column("asset_event_id", existing_type=sa.Integer(), nullable=False) + batch_op.drop_constraint("adrq_asset_fkey", type_="foreignkey") + batch_op.drop_constraint("assetdagrunqueue_pkey", type_="primary") + batch_op.create_primary_key("assetdagrunqueue_pkey", ["target_dag_id", "asset_event_id"]) + batch_op.create_foreign_key("adrq_asset_fkey", "asset", ["asset_id"], ["id"], ondelete="CASCADE") + batch_op.create_foreign_key( + "adrq_asset_event_fkey", "asset_event", ["asset_event_id"], ["id"], ondelete="CASCADE" + ) + + # 4. Repopulate the per-event rows from staging. + if not context.is_offline_mode(): + conn = op.get_bind() + now = datetime.now(timezone.utc) + conn.execute( + sa.text( + f""" + INSERT INTO asset_dag_run_queue (asset_id, target_dag_id, asset_event_id, created_at) + SELECT asset_id, target_dag_id, asset_event_id, :now + FROM {_STAGING} + """ + ), + {"now": now}, + ) + op.drop_table(_STAGING) + + +def downgrade(): + """Revert reference to asset_id from asset_dag_run_queue.""" + # 1. Rebuild the coalesced rows: collapse per-event rows to one row per + # (asset_id, target_dag_id) (created_at = max referenced event timestamp), then clear. + if context.is_offline_mode(): + print( + dedent(""" + ------------ + -- WARNING: asset_dag_run_queue cannot be rebuilt in offline mode; + -- any pending (unprocessed) queued asset events will be dropped. + ------------ + """) + ) + op.execute("DELETE FROM asset_dag_run_queue") + else: + conn = op.get_bind() + conn.execute( + sa.text( + f""" + CREATE TABLE {_STAGING} AS + SELECT adrq.asset_id AS asset_id, + adrq.target_dag_id AS target_dag_id, + MAX(ae.timestamp) AS created_at + FROM asset_dag_run_queue adrq + JOIN asset_event ae ON ae.id = adrq.asset_event_id + GROUP BY adrq.asset_id, adrq.target_dag_id + """ + ) + ) + conn.execute(sa.text("DELETE FROM asset_dag_run_queue")) + + # 2. Drop the reference + FK, restore the old primary key. + with op.batch_alter_table("asset_dag_run_queue") as batch_op: + batch_op.drop_constraint("adrq_asset_event_fkey", type_="foreignkey") + batch_op.drop_constraint("assetdagrunqueue_pkey", type_="primary") + batch_op.create_primary_key("assetdagrunqueue_pkey", ["asset_id", "target_dag_id"]) + batch_op.drop_column("asset_event_id") + + # 3. Repopulate the coalesced rows. + if not context.is_offline_mode(): + conn = op.get_bind() + conn.execute( + sa.text( + f""" + INSERT INTO asset_dag_run_queue (asset_id, target_dag_id, created_at) + SELECT asset_id, target_dag_id, created_at + FROM {_STAGING} + """ + ) + ) + op.drop_table(_STAGING) diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 7aafddb3b9490..8a8122db0c93b 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -749,21 +749,28 @@ def __repr__(self): class AssetDagRunQueue(Base): """Model for storing asset events that need processing.""" - asset_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) target_dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False) + asset_event_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) + asset_id: Mapped[int] = mapped_column(Integer, nullable=False) created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) asset: Mapped[AssetModel] = relationship("AssetModel", viewonly=True) dag_model: Mapped[DagModel] = relationship("DagModel", viewonly=True) __tablename__ = "asset_dag_run_queue" __table_args__ = ( - PrimaryKeyConstraint(asset_id, target_dag_id, name="assetdagrunqueue_pkey"), + PrimaryKeyConstraint(target_dag_id, asset_event_id, name="assetdagrunqueue_pkey"), ForeignKeyConstraint( (asset_id,), ["asset.id"], name="adrq_asset_fkey", ondelete="CASCADE", ), + ForeignKeyConstraint( + (asset_event_id,), + ["asset_event.id"], + name="adrq_asset_event_fkey", + ondelete="CASCADE", + ), ForeignKeyConstraint( (target_dag_id,), ["dag.dag_id"], @@ -775,7 +782,7 @@ class AssetDagRunQueue(Base): def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self.asset_id == other.asset_id and self.target_dag_id == other.target_dag_id + return self.target_dag_id == other.target_dag_id and self.asset_event_id == other.asset_event_id return NotImplemented def __hash__(self): diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 615ecb66c6799..21ccd5ea3bc67 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "7a98f1b7dbd3", + "3.4.0": "b2f1a9c7d4e0", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py index f0e60462bf34c..139b7ad4a1a5e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py @@ -1504,7 +1504,10 @@ class TestQueuedEventEndpoint(TestAssets): def _create_asset_dag_run_queues(self, dag_id, asset_id, session): session.execute(delete(AssetDagRunQueue)) session.flush() - adrq = AssetDagRunQueue(target_dag_id=dag_id, asset_id=asset_id) + event = AssetEvent(asset_id=asset_id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() + adrq = AssetDagRunQueue(target_dag_id=dag_id, asset_id=asset_id, asset_event_id=event.id) session.add(adrq) session.commit() return adrq @@ -2186,12 +2189,12 @@ def test_should_respond_204(self, test_client, session, create_dummy_dag): dag, _ = create_dummy_dag() dag_id = dag.dag_id (asset,) = self.create_assets(session=session, num=1) - self._create_asset_dag_run_queues(dag_id, asset.id, session) + adrq = self._create_asset_dag_run_queues(dag_id, asset.id, session) - assert session.get(AssetDagRunQueue, (asset.id, dag_id)) is not None + assert session.get(AssetDagRunQueue, (dag_id, adrq.asset_event_id)) is not None response = test_client.delete(f"/assets/{asset.id}/queuedEvents") assert response.status_code == 204 - assert session.get(AssetDagRunQueue, (asset.id, dag_id)) is None + assert session.get(AssetDagRunQueue, (dag_id, adrq.asset_event_id)) is None check_last_log(session, dag_id=None, event="delete_asset_queued_events", logical_date=None) def test_should_respond_401(self, unauthenticated_test_client): diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py index 05735154c5f6c..2427729035fea 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py @@ -137,9 +137,15 @@ def test_should_set_last_update_only_for_queued_and_hide_flag(self, test_client, ) } # Queue and add an event only for A - session.add(AssetDagRunQueue(asset_id=assets["s3://bucket/A"].id, target_dag_id="two_assets_equal")) + event = AssetEvent(asset_id=assets["s3://bucket/A"].id, timestamp=dr.logical_date or pendulum.now()) + session.add(event) + session.flush() session.add( - AssetEvent(asset_id=assets["s3://bucket/A"].id, timestamp=dr.logical_date or pendulum.now()) + AssetDagRunQueue( + asset_id=assets["s3://bucket/A"].id, + target_dag_id="two_assets_equal", + asset_event_id=event.id, + ) ) session.commit() @@ -210,12 +216,16 @@ def test_last_update_respects_latest_run_filter(self, test_client, dag_maker, se dag_maker.sync_dagbag_to_db() asset = session.scalars(select(AssetModel).where(AssetModel.uri == "s3://bucket/F")).one() - session.add(AssetDagRunQueue(asset_id=asset.id, target_dag_id="filter_run")) - # event before latest_run should be ignored ts_base = dr.logical_date or pendulum.now() - session.add(AssetEvent(asset_id=asset.id, timestamp=ts_base.subtract(minutes=10))) + # event before latest_run should be ignored + event_before = AssetEvent(asset_id=asset.id, timestamp=ts_base.subtract(minutes=10)) # event after latest_run counts - session.add(AssetEvent(asset_id=asset.id, timestamp=ts_base.add(minutes=10))) + event_after = AssetEvent(asset_id=asset.id, timestamp=ts_base.add(minutes=10)) + session.add_all([event_before, event_after]) + session.flush() + session.add( + AssetDagRunQueue(asset_id=asset.id, target_dag_id="filter_run", asset_event_id=event_after.id) + ) session.commit() resp = test_client.get("/next_run_assets/filter_run") diff --git a/airflow-core/tests/unit/assets/test_manager.py b/airflow-core/tests/unit/assets/test_manager.py index bc034cd751540..4dd2d92621588 100644 --- a/airflow-core/tests/unit/assets/test_manager.py +++ b/airflow-core/tests/unit/assets/test_manager.py @@ -26,7 +26,7 @@ import pytest from sqlalchemy import delete, func, select -from sqlalchemy.dialects import mysql +from sqlalchemy.dialects import mysql, postgresql, sqlite from sqlalchemy.orm import Session from airflow import settings @@ -212,23 +212,43 @@ def test_register_asset_change_no_downstreams(self, session, mock_task_instance) ) assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 0 - @pytest.mark.parametrize( - ("dialect_name", "expected_helper"), - [ - ("postgresql", "_queue_dagruns_nonpartitioned_conflict_update"), - ("mysql", "_queue_dagruns_nonpartitioned_mysql"), - ("sqlite", "_queue_dagruns_nonpartitioned_conflict_update"), - ], - ) - def test_queue_dagruns_routes_by_dialect(self, dialect_name, expected_helper): - """Test that _queue_dagruns routes to the dialect-appropriate queue helper.""" + def test_register_asset_change_is_atomic_on_caller_session(self, session, mock_task_instance): + """The AssetEvent is written on the caller's session: visible before commit, gone after rollback. + + Under the old side-session behaviour the event was committed independently and would be + orphaned if the caller's transaction rolled back. Registration is now atomic with the caller. + """ + asset_manager = AssetManager() + + asset = Asset(uri="test://atomic1", name="atomic_asset") + asm = AssetModel(uri="test://atomic1/", name="atomic_asset", group="asset") + session.add(asm) + session.execute(delete(AssetDagRunQueue)) + session.flush() + + event = asset_manager.register_asset_change( + task_instance=mock_task_instance, asset=asset, session=session + ) + session.flush() + + # The event is written on the caller's session and carries a real id before any commit. + assert event is not None + assert event.id is not None + event_id = event.id + assert session.get(AssetEvent, event_id) is not None + + # Rolling back the caller's transaction discards the event -> no orphan row. + session.rollback() + assert session.get(AssetEvent, event_id) is None + + def test_queue_dagruns_calls_nonpartitioned_helper(self): + """`_queue_dagruns` delegates non-partitioned dags to the single non-partitioned helper.""" dag = DagModel(dag_id="dag1") session = mock.MagicMock(spec=Session) event = mock.MagicMock() with ( - mock.patch("airflow.assets.manager.get_dialect_name", return_value=dialect_name), mock.patch.object(AssetManager, "_queue_partitioned_dags"), - mock.patch.object(AssetManager, expected_helper) as mock_helper, + mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned") as mock_helper, ): AssetManager._queue_dagruns( asset_id=1, @@ -239,26 +259,44 @@ def test_queue_dagruns_routes_by_dialect(self, dialect_name, expected_helper): task_instance=None, session=session, ) - if expected_helper == "_queue_dagruns_nonpartitioned_conflict_update": - mock_helper.assert_called_once_with(1, {dag}, event, session, dialect_name) - elif expected_helper == "_queue_dagruns_nonpartitioned_mysql": - mock_helper.assert_called_once_with(1, {dag}, event, session) - else: - raise AssertionError(f"Unexpected expected_helper: {expected_helper}") + mock_helper.assert_called_once_with(1, {dag}, event, session) - def test_queue_dagruns_nonpartitioned_mysql_builds_upsert(self): - """Test that the MySQL queue path emits an INSERT ... ON DUPLICATE KEY UPDATE.""" + @pytest.mark.parametrize("dialect_name", ["postgresql", "sqlite"]) + def test_queue_dagruns_nonpartitioned_insert_or_ignore(self, dialect_name): + """On postgres/sqlite the helper emits INSERT ... ON CONFLICT DO NOTHING referencing the event.""" dag = DagModel(dag_id="dag1") session = mock.MagicMock(spec=Session) event = AssetEvent(asset_id=1) - AssetManager._queue_dagruns_nonpartitioned_mysql( - asset_id=1, dags_to_queue={dag}, event=event, session=session - ) + event.id = 99 + with mock.patch("airflow.assets.manager.get_dialect_name", return_value=dialect_name): + AssetManager._queue_dagruns_nonpartitioned( + asset_id=1, dags_to_queue={dag}, event=event, session=session + ) + + stmt, values = session.execute.call_args.args + dialect = {"postgresql": postgresql.dialect(), "sqlite": sqlite.dialect()}[dialect_name] + compiled = str(stmt.compile(dialect=dialect)).upper() + assert "ON CONFLICT" in compiled + assert "DO NOTHING" in compiled + # One ADRQ row is inserted per (dag, event), carrying the denormalized asset_id + # plus the referenced asset_event_id. + assert values == [{"asset_id": 1, "target_dag_id": "dag1", "asset_event_id": 99}] + + def test_queue_dagruns_nonpartitioned_insert_ignore_on_mysql(self): + """On MySQL the helper emits INSERT IGNORE referencing the triggering event.""" + dag = DagModel(dag_id="dag1") + 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"): + AssetManager._queue_dagruns_nonpartitioned( + asset_id=1, dags_to_queue={dag}, event=event, session=session + ) stmt, values = session.execute.call_args.args compiled = str(stmt.compile(dialect=mysql.dialect())).upper() - assert "ON DUPLICATE KEY UPDATE" in compiled - assert values == [{"target_dag_id": "dag1"}] + assert "INSERT IGNORE" in compiled + assert values == [{"asset_id": 1, "target_dag_id": "dag1", "asset_event_id": 99}] def test_register_asset_change_notifies_asset_listener( self, session, mock_task_instance, testing_dag_bundle, listener_manager diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 53ea5bc933f3f..c83d9039d8f6e 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -47,7 +47,7 @@ from airflow._shared.observability.metrics.base_stats_logger import StatsLogger from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator -from airflow.assets.manager import AssetManager, _create_asset_event +from airflow.assets.manager import AssetManager from airflow.callbacks.callback_requests import ( DagCallbackRequest, DagRunContext, @@ -5707,16 +5707,14 @@ def test_create_dag_runs_assets(self, session, dag_maker): timestamp=base + timedelta(seconds=2), ) session.add_all([event1, event2]) + session.flush() # assign event ids so the ADRQ rows can reference them session = dag_maker.session session.add_all( [ - AssetDagRunQueue( - asset_id=asset1_id, target_dag_id=dag2.dag_id, created_at=base + timedelta(hours=1) - ), - AssetDagRunQueue( - asset_id=asset1_id, target_dag_id=dag3.dag_id, created_at=base + timedelta(hours=1) - ), + AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag2.dag_id, asset_event_id=event1.id), + AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag3.dag_id, asset_event_id=event1.id), + AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag3.dag_id, asset_event_id=event2.id), ] ) session.flush() @@ -5773,28 +5771,21 @@ def dict_from_obj(obj): def test_new_asset_triggered_dag_backlog_gated_by_catchup( self, catchup, expects_old_event, session, dag_maker ): - """Reproduces #39456: catchup gates whether a new asset-triggered Dag replays the - pre-creation backlog. With catchup off (the default) it only consumes events after it - started scheduling on the asset; with catchup on it replays the full history.""" + """catchup gates whether a newly-subscribed asset-triggered Dag consumes its backlog. + + With catchup off (the default) the Dag consumes only events with a queue row (i.e. those + emitted after it began scheduling on the asset). With catchup on, the first triggered run + also consumes the pre-subscription backlog -- every not-yet-consumed event for the Dag's + assets -- selected directly by the scheduler (no queue row required, no time window). + """ asset = Asset(uri="test://asset-historical", name="hist_asset", group="test_group") - # Producer Dag + run that the asset events are sourced from. + # Producer + a historical event that exists BEFORE any consumer subscribes. with dag_maker(dag_id="historical-producer", start_date=timezone.utcnow(), session=session): BashOperator(task_id="task", bash_command="echo 1", outlets=[asset]) producer_run = dag_maker.create_dagrun(run_id="producer-run") - asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) - # Consumer Dag created now; its schedule reference's created_at is the cut-off. - with dag_maker(dag_id="historical-consumer", schedule=[asset], catchup=catchup): - pass - consumer_dag = dag_maker.dag - reference_created_at = session.scalar( - select(DagScheduleAssetReference.created_at).where( - DagScheduleAssetReference.dag_id == consumer_dag.dag_id - ) - ) - def _make_event(timestamp): return AssetEvent( asset_id=asset_id, @@ -5805,18 +5796,28 @@ def _make_event(timestamp): timestamp=timestamp, ) - old_event = _make_event(reference_created_at - timedelta(days=1)) - new_event = _make_event(reference_created_at + timedelta(seconds=1)) - session.add_all([old_event, new_event]) - # Trigger time after both events so neither is excluded by the upper bound. + old_event = _make_event(timezone.utcnow() - timedelta(days=1)) + session.add(old_event) + session.commit() + + # Consumer subscribes now: catchup=True backfills a queue row for old_event (the + # pre-subscription backlog); catchup=False backfills nothing. + with dag_maker(dag_id="historical-consumer", schedule=[asset], catchup=catchup, session=session): + pass + consumer_dag = dag_maker.dag + + # A post-subscription event always gets its own queue row. + new_event = _make_event(timezone.utcnow()) + session.add(new_event) + session.flush() session.add( AssetDagRunQueue( asset_id=asset_id, target_dag_id=consumer_dag.dag_id, - created_at=reference_created_at + timedelta(hours=1), + asset_event_id=new_event.id, ) ) - session.flush() + session.commit() scheduler_job = Job() self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) @@ -5828,6 +5829,78 @@ def _make_event(timestamp): expected = {new_event.id} | ({old_event.id} if expects_old_event else set()) assert {e.id for e in created_run.consumed_asset_events} == expected + @pytest.mark.need_serialized_dag + def test_asset_events_out_of_order_are_both_consumed(self, session, dag_maker): + """Regression test for GH-54659. + + Two events for the same asset can become visible out of timestamp order (for example a + long-running producer commits an "older" event after a "newer" one). Under the old + created_at/timestamp watermark the older event could be stranded below the watermark and + never consumed. With consume-by-reference every event referenced by an ADRQ row is + consumed, regardless of the order in which the timestamps became visible. + """ + asset = Asset(uri="test://asset-ooo", name="ooo_asset", group="test_group") + + with dag_maker(dag_id="ooo-producer", start_date=timezone.utcnow(), session=session): + BashOperator(task_id="task", bash_command="echo 1", outlets=[asset]) + producer_run = dag_maker.create_dagrun(run_id="producer-run") + + asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) + + with dag_maker(dag_id="ooo-consumer", schedule=[asset]): + pass + consumer_dag = dag_maker.dag + + base = timezone.utcnow() + + def _make_event(timestamp): + return AssetEvent( + asset_id=asset_id, + source_task_id="task", + source_dag_id=producer_run.dag_id, + source_run_id=producer_run.run_id, + source_map_index=-1, + timestamp=timestamp, + ) + + # The "newer" event is registered (and gets its lower id) BEFORE the "older" one, so + # insertion order and timestamp order disagree. + newer_event = _make_event(base + timedelta(seconds=10)) + session.add(newer_event) + session.flush() + older_event = _make_event(base + timedelta(seconds=1)) + session.add(older_event) + session.flush() + + session.add_all( + [ + AssetDagRunQueue( + asset_id=asset_id, target_dag_id=consumer_dag.dag_id, asset_event_id=newer_event.id + ), + AssetDagRunQueue( + asset_id=asset_id, target_dag_id=consumer_dag.dag_id, asset_event_id=older_event.id + ), + ] + ) + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + with create_session() as session: + self.job_runner._create_dagruns_for_dags(session, session) + + created_run = session.scalars(select(DagRun).where(DagRun.dag_id == consumer_dag.dag_id)).one() + assert created_run.state == State.QUEUED + # Neither event is stranded: both are consumed by the single run. + assert {e.id for e in created_run.consumed_asset_events} == {newer_event.id, older_event.id} + # All ADRQ rows for the dag are cleared once consumed. + assert ( + session.scalars( + select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == consumer_dag.dag_id) + ).all() + == [] + ) + @pytest.mark.need_serialized_dag def test_create_dag_runs_asset_triggered_skips_stale_triggered_date(self, session, dag_maker): asset = Asset(uri="test://asset-for-stale-trigger-date", name="asset-for-stale-trigger-date") @@ -5836,8 +5909,12 @@ def test_create_dag_runs_asset_triggered_skips_stale_triggered_date(self, sessio dag_model = dag_maker.dag_model asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) - queued_at = timezone.utcnow() - session.add(AssetDagRunQueue(target_dag_id=dag_model.dag_id, asset_id=asset_id, created_at=queued_at)) + event = AssetEvent(asset_id=asset_id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() + session.add( + AssetDagRunQueue(target_dag_id=dag_model.dag_id, asset_id=asset_id, asset_event_id=event.id) + ) session.flush() # Simulate another scheduler consuming ADRQ rows after we computed triggered_date_by_dag. @@ -5865,21 +5942,20 @@ def test_create_dag_runs_asset_triggered_deletes_only_selected_adrq_rows( dag_model = dag_maker.dag_model asset_1_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset_1.name)) asset_2_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset_2.name)) + event_1 = AssetEvent(asset_id=asset_1_id, timestamp=timezone.utcnow()) + event_2 = AssetEvent(asset_id=asset_2_id, timestamp=timezone.utcnow()) + session.add_all([event_1, event_2]) + session.flush() session.add_all( [ - AssetEvent( - asset_id=asset_1_id, - timestamp=timezone.utcnow(), - ), - # The ADRQ that should triggers the Dag run creation + # The ADRQ that should trigger the Dag run creation AssetDagRunQueue( - asset_id=asset_1_id, target_dag_id=dag_model.dag_id, created_at=timezone.utcnow() + asset_id=asset_1_id, target_dag_id=dag_model.dag_id, asset_event_id=event_1.id ), - AssetEvent(asset_id=asset_2_id, timestamp=timezone.utcnow()), - # The ADRQ that arrives after the Dag run creation but before ADRQ clean up - # This situation is simulated by _lock_only_selected_asset below + # The ADRQ that arrives after the Dag run creation but before ADRQ clean up. + # This situation is simulated by _lock_only_selected_asset below. AssetDagRunQueue( - asset_id=asset_2_id, target_dag_id=dag_model.dag_id, created_at=timezone.utcnow() + asset_id=asset_2_id, target_dag_id=dag_model.dag_id, asset_event_id=event_2.id ), ] ) @@ -5946,23 +6022,23 @@ def create_asset_events(sleep): dag = session.get(DagModel, consumer_dag_id) now = timezone.utcnow() asset_manager = AssetManager() - asset_event = _create_asset_event(session=session, asset_id=asset_id, timestamp=now) + # The event is now created inline on the caller's session (atomic), rather than in + # a side session, so build it directly here. + asset_event = AssetEvent(asset_id=asset_id, timestamp=now) + session.add(asset_event) + session.flush() + event_id = asset_event.id time.sleep(sleep) # widen the race window between event creation and queueing - dialect_name = inspect(session.get_bind()).dialect.name - if dialect_name in ("postgresql", "sqlite"): - asset_manager._queue_dagruns_nonpartitioned_conflict_update( - asset_id=asset_id, - dags_to_queue=[dag], - event=asset_event, - session=session, - dialect_name=dialect_name, - ) - elif dialect_name == "mysql": - asset_manager._queue_dagruns_nonpartitioned_mysql( - asset_id=asset_id, dags_to_queue=[dag], event=asset_event, session=session - ) + # A single dialect-agnostic helper now performs insert-or-ignore keyed on + # (target_dag_id, asset_event_id). + asset_manager._queue_dagruns_nonpartitioned( + asset_id=asset_id, + dags_to_queue={dag}, + event=asset_event, + session=session, + ) - return asset_event.id, now.isoformat() + return event_id, now.isoformat() with ( ThreadPoolExecutor(max_workers=3) as executor, @@ -6045,7 +6121,9 @@ def test_create_dag_runs_asset_alias_with_asset_event_attached(self, session, da session = dag_maker.session session.add_all( [ - AssetDagRunQueue(asset_id=asset1_id, target_dag_id=consumer_dag.dag_id), + AssetDagRunQueue( + asset_id=asset1_id, target_dag_id=consumer_dag.dag_id, asset_event_id=event.id + ), ] ) session.flush() @@ -6119,9 +6197,11 @@ def test_asset_triggered_dagruns_respects_team_name( source_map_index=-1, ) session.add(event) - # flush here to ensure event timestamp is before the ADRQ created_at timestamp + # flush here to assign the event id referenced by the ADRQ row session.flush() - session.add(AssetDagRunQueue(asset_id=asset_id, target_dag_id=f"consumer_{suffix}")) + session.add( + AssetDagRunQueue(asset_id=asset_id, target_dag_id=f"consumer_{suffix}", asset_event_id=event.id) + ) session.flush() with conf_vars({("core", "multi_team"): multi_team}): @@ -6187,26 +6267,42 @@ def test_no_create_dag_runs_when_dag_disabled(self, session, dag_maker, disable, ) session.flush() assert [e.source_run_id for e in session.scalars(ase_q)] == [dr1.run_id, dr2.run_id] - assert len(session.scalars(adrq_q).all()) == 1 - assert session.scalars(adrq_q).one().target_dag_id == "consumer" + # ADRQ rows are per asset event now. A stale dag still enqueues while disabled (asserted + # above), so both events stay queued; a paused dag never enqueued the first event, so only + # the second remains. + expected_adrqs = 2 if "is_stale" in disable else 1 + adrqs = session.scalars(adrq_q).all() + assert len(adrqs) == expected_adrqs + assert all(adrq.target_dag_id == "consumer" for adrq in adrqs) @pytest.mark.need_serialized_dag - def test_no_create_dag_runs_when_no_asset_event(self, session: Session, dag_maker, caplog): + 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): pass dag_model = dag_maker.dag_model asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) - # Simulate an ADRQ row whose matching asset events were already consumed by an earlier DagRun. - # The ADRQ should be cleaned up even when no new DagRun is created, to prevent stale ADRQ - # rows from accumulating and causing infinite scheduler loops. - adrq = AssetDagRunQueue( - asset_id=asset_id, target_dag_id=dag_model.dag_id, created_at=timezone.utcnow() + + # 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) + session.flush() + prior_run = dag_maker.create_dagrun( + run_id="prior-consuming-run", + logical_date=DEFAULT_DATE, + data_interval=(DEFAULT_DATE, DEFAULT_DATE), ) - session.add(adrq) + prior_run.consumed_asset_events.append(event) session.flush() - adrq.created_at = timezone.utcnow() + timedelta(seconds=1) - session.merge(adrq) + + # 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) + ) + session.flush() + with caplog.at_level("INFO"): scheduler_job = Job() self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[MockExecutor(do_update=False)]) @@ -6214,9 +6310,11 @@ def test_no_create_dag_runs_when_no_asset_event(self, session: Session, dag_make dag_models=[dag_model], session=session, ) - dr = session.scalars(select(DagRun).where(DagRun.dag_id == dag_model.dag_id)).one_or_none() + assert "No DagRun created" in caplog.text - assert dr is None + # 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 diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index ead765b31716d..242bfc25e591f 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -2585,7 +2585,12 @@ def test_dags_needing_dagruns_assets(self, dag_maker, session): # add queue records so we'll need a run dag_model = session.scalar(select(DagModel).where(DagModel.dag_id == dag.dag_id)) asset_model: AssetModel = dag_model.schedule_assets[0] - session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id)) + event = AssetEvent(asset_id=asset_model.id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() + session.add( + AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id, asset_event_id=event.id) + ) session.flush() query, _ = DagModel.dags_needing_dagruns(session) dag_models = query.all() @@ -2639,7 +2644,10 @@ def test_dags_needing_dagruns_skips_adrq_when_serialized_dag_missing( session.add(dag_model) session.flush() - session.add(AssetDagRunQueue(asset_id=asset_id, target_dag_id=orphan_dag_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=orphan_dag_id, asset_event_id=event.id)) session.flush() with caplog.at_level(logging.DEBUG, logger="airflow.models.dag"): @@ -2707,10 +2715,14 @@ def test_dags_needing_dagruns_missing_serialized_debug_lists_sorted_dag_ids( ) session.flush() + event_z = AssetEvent(asset_id=id_z, timestamp=timezone.utcnow()) + event_a = AssetEvent(asset_id=id_a, timestamp=timezone.utcnow()) + session.add_all([event_z, event_a]) + session.flush() session.add_all( [ - AssetDagRunQueue(asset_id=id_z, target_dag_id="ghost_z"), - AssetDagRunQueue(asset_id=id_a, target_dag_id="ghost_a"), + AssetDagRunQueue(asset_id=id_z, target_dag_id="ghost_z", asset_event_id=event_z.id), + AssetDagRunQueue(asset_id=id_a, target_dag_id="ghost_a", asset_event_id=event_a.id), ] ) session.flush() @@ -2754,7 +2766,14 @@ def test_dags_needing_dagruns_query_count(self, dag_maker, session): asset_models = dag_model.schedule_assets assert len(asset_models) == num_assets for asset_model in asset_models: - session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id)) + event = AssetEvent(asset_id=asset_model.id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() + session.add( + AssetDagRunQueue( + asset_id=asset_model.id, target_dag_id=dag_model.dag_id, asset_event_id=event.id + ) + ) session.flush() # Clear identity map so N+1 on adrq.asset is exposed @@ -2788,7 +2807,12 @@ def test_dags_needing_dagruns_asset_aliases(self, dag_maker, session): # add queue records so we'll need a run dag_model = dag_maker.dag_model - session.add(AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id)) + event = AssetEvent(asset_id=asset_model.id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() + session.add( + AssetDagRunQueue(asset_id=asset_model.id, target_dag_id=dag_model.dag_id, asset_event_id=event.id) + ) session.flush() query, _ = DagModel.dags_needing_dagruns(session) dag_models = query.all() @@ -3046,12 +3070,24 @@ def test_dags_needing_dagruns_triggered_date_by_dag_queued_times(self, session, pass session.flush() + asset_event_ids = { + e.asset_id: e.id + for e in session.scalars( + select(AssetEvent).where(AssetEvent.asset_id.in_([asset1_id, asset2_id])) + ) + } session.add_all( [ - AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag.dag_id, created_at=DEFAULT_DATE), + AssetDagRunQueue( + asset_id=asset1_id, + target_dag_id=dag.dag_id, + asset_event_id=asset_event_ids[asset1_id], + created_at=DEFAULT_DATE, + ), AssetDagRunQueue( asset_id=asset2_id, target_dag_id=dag.dag_id, + asset_event_id=asset_event_ids[asset2_id], created_at=DEFAULT_DATE + timedelta(hours=1), ), ] @@ -3824,10 +3860,13 @@ def test_get_asset_triggered_next_run_info(dag_maker, clear_assets): session = dag_maker.session asset1_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset1.uri)) + event = AssetEvent(asset_id=asset1_id, timestamp=timezone.utcnow()) + session.add(event) + session.flush() session.bulk_save_objects( [ - AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag2.dag_id), - AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag3.dag_id), + AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag2.dag_id, asset_event_id=event.id), + AssetDagRunQueue(asset_id=asset1_id, target_dag_id=dag3.dag_id, asset_event_id=event.id), ] ) session.flush() diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 53fa99aa22460..6a348a3f2ca26 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -3877,14 +3877,15 @@ def test_runtime_partition_key_does_not_backfill_dag_run_when_none(dag_maker, se @pytest.mark.backend("sqlite") def test_runtime_partition_key_backfill_does_not_deadlock_on_sqlite(dag_maker, session): - """Regression test for the SQLite ``database is locked`` deadlock between the - writes in ``register_asset_changes_in_db`` and the second connection that - ``_create_asset_event`` used to open. - - On file-based SQLite (the default ``-b sqlite`` test backend) the two - connections compete for the same RESERVED lock; the SQLite branch of - ``_create_asset_event`` must add the event directly to the caller's session - instead of opening a second connection. + """Regression test for the SQLite ``database is locked`` deadlock. + + This happens when a second connection is used to trigger while + ``register_asset_changes_in_db`` was writing. + + The asset event is now created inline on the caller's session (see + ``AssetManager.register_asset_change``) instead of opening a side session, so + on file-based SQLite (the default ``-b sqlite`` test backend) there is no + longer a second connection competing for the same RESERVED lock. """ asset = Asset(name="hello") with dag_maker(dag_id="rt_pk_backfill_sqlite", schedule=PartitionedAtRuntime()) as dag: diff --git a/airflow-core/tests/unit/timetables/test_assets_timetable.py b/airflow-core/tests/unit/timetables/test_assets_timetable.py index 3c12c886f283f..ccdc395ae410d 100644 --- a/airflow-core/tests/unit/timetables/test_assets_timetable.py +++ b/airflow-core/tests/unit/timetables/test_assets_timetable.py @@ -295,7 +295,10 @@ def test_asset_dag_run_queue_processing(self, session, dag_maker, create_test_as # Add AssetDagRunQueue entries to simulate asset event processing for am in asset_models: - session.add(AssetDagRunQueue(asset_id=am.id, target_dag_id=dag.dag_id)) + event = AssetEvent(asset_id=am.id) + session.add(event) + session.flush() + session.add(AssetDagRunQueue(asset_id=am.id, target_dag_id=dag.dag_id, asset_event_id=event.id)) session.commit() # Fetch and evaluate asset triggers for all DAGs affected by asset events