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