This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 44acf383e47 Reference asset events from asset_dag_run_queue (#70972)
44acf383e47 is described below
commit 44acf383e47892465c7e14f55f08b914ae150085
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Wed Aug 5 09:13:54 2026 +0800
Reference asset events from asset_dag_run_queue (#70972)
---
airflow-core/docs/howto/usage-cli.rst | 6 +
airflow-core/docs/migrations-ref.rst | 5 +-
airflow-core/newsfragments/70972.significant.rst | 13 ++
airflow-core/src/airflow/assets/manager.py | 135 +++---------
.../src/airflow/jobs/scheduler_job_runner.py | 87 +++-----
.../0128_3_4_0_add_asset_event_id_to_adrq.py | 205 ++++++++++++++++++
airflow-core/src/airflow/models/asset.py | 15 +-
airflow-core/src/airflow/models/dag.py | 55 +----
airflow-core/src/airflow/utils/db.py | 2 +-
.../core_api/routes/public/test_assets.py | 9 +-
.../api_fastapi/core_api/routes/ui/test_assets.py | 22 +-
airflow-core/tests/unit/assets/test_manager.py | 122 +++++++----
airflow-core/tests/unit/jobs/test_scheduler_job.py | 240 +++++++++++++++------
airflow-core/tests/unit/models/test_dag.py | 116 ++++------
.../tests/unit/models/test_taskinstance.py | 17 +-
.../tests/unit/timetables/test_assets_timetable.py | 5 +-
16 files changed, 627 insertions(+), 427 deletions(-)
diff --git a/airflow-core/docs/howto/usage-cli.rst
b/airflow-core/docs/howto/usage-cli.rst
index d96cf5331f9..21b37747e68 100644
--- a/airflow-core/docs/howto/usage-cli.rst
+++ b/airflow-core/docs/howto/usage-cli.rst
@@ -223,6 +223,12 @@ The ``db clean`` command works by deleting from each table
the records older tha
You can optionally provide a list of tables to perform deletes on. If no list
of tables is supplied, all tables will be included.
+.. note::
+
+ Cleaning the ``asset_event`` table also drops any
queued-but-not-yet-consumed asset events. This mainly
+ affects Dags waiting on a multi-asset condition, where a pending event can
be purged before the condition
+ is met, meaning the Dag will not be triggered by it.
+
You can filter cleanup to specific DAGs using ``--dag-ids`` (comma-separated
list), or exclude specific DAGs using ``--exclude-dag-ids`` (comma-separated
list). These options allow you to target or avoid cleanup for particular DAGs.
You can use the ``--dry-run`` option to print the row counts in the primary
tables to be cleaned.
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index dfc96be7901..e3b33cab185 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/newsfragments/70972.significant.rst
b/airflow-core/newsfragments/70972.significant.rst
new file mode 100644
index 00000000000..bae78066f12
--- /dev/null
+++ b/airflow-core/newsfragments/70972.significant.rst
@@ -0,0 +1,13 @@
+Changes to how pending asset events trigger Dag runs
+
+Deleting asset event records with ``airflow db clean`` now also removes any
+asset events that are still pending for a downstream Dag. This mainly affects
+Dags waiting on a multi-asset condition, where a pending event can be purged
+before the condition is met. Deleting an asset event now consistently means it
+will not trigger a run; previously such a pending trigger could still fire
after
+the event had been purged.
+
+Relatedly, an asset event that was queued for a Dag while it was scheduled on
+that asset is still used to trigger the run, and appears among the run's
+triggering events, even if the asset is later removed from the Dag's schedule.
+Previously, they were silently dropped.
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index 9330c6071db..a417ad96eae 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -48,7 +48,6 @@ from airflow.models.log import Log
from airflow.timetables.base import compute_rollup_fingerprint
from airflow.utils.helpers import is_container, prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.session import create_session
from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks
if TYPE_CHECKING:
@@ -122,38 +121,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.
@@ -379,7 +346,10 @@ class AssetManager(LoggingMixin):
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}
@@ -561,25 +531,7 @@ class AssetManager(LoggingMixin):
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(
@@ -844,71 +796,38 @@ class AssetManager(LoggingMixin):
return apdr
@classmethod
- def _queue_dagruns_nonpartitioned_slow_path(
- 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(
+ def _queue_dagruns_nonpartitioned(
cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent,
session: Session
) -> None:
- from sqlalchemy import case
- from sqlalchemy.dialects.mysql import insert
-
- 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_duplicate_key_update(
- created_at=case(
- (stmt.inserted.created_at >= AssetDagRunQueue.created_at,
stmt.inserted.created_at),
- else_=AssetDagRunQueue.created_at,
+ 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
+ ]
+
+ if (dialect_name := get_dialect_name(session)) == "mysql":
+ from sqlalchemy.dialects.mysql import insert as my_insert
+
+ session.execute(
+ # The on-dup update a no-op since we are always working on
ADRQs
+ # for the asset. It may not hold if you change e.g. ADRQ
schema.
+
my_insert(AssetDagRunQueue).on_duplicate_key_update(asset_id=AssetDagRunQueue.asset_id),
+ values,
)
- )
- session.execute(update_stmt, 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 b6a29995d59..40d16515057 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 @@ import time
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 @@ from airflow.utils.types import DagRunTriggeredByType,
DagRunType
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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
)
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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
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 00000000000..ff2a2282181
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_id_to_adrq.py
@@ -0,0 +1,205 @@
+#
+# 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 = COALESCE(
+ MAX(dag_run.run_after) for asset-triggered runs of the dag with
+ run_after < triggered_date, (per dag)
+ dag_schedule_asset_reference.created_at, (per dag + asset)
+ )
+ 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
+LEFT JOIN dag_schedule_asset_reference dsar
+ ON dsar.dag_id = adrq.target_dag_id
+ AND dsar.asset_id = adrq.asset_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, dsar.created_at, :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, adjust FKs and indexes.
+ 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_index("idx_adrq_asset_id", ["asset_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"
+ )
+ batch_op.drop_index("idx_asset_dag_run_queue_target_dag_id")
+
+ # 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 and indexes.
+ with op.batch_alter_table("asset_dag_run_queue") as batch_op:
+ batch_op.drop_constraint("adrq_asset_event_fkey", type_="foreignkey")
+ batch_op.create_index("idx_asset_dag_run_queue_target_dag_id",
["target_dag_id"])
+ batch_op.drop_constraint("assetdagrunqueue_pkey", type_="primary")
+ batch_op.create_primary_key("assetdagrunqueue_pkey", ["asset_id",
"target_dag_id"])
+ batch_op.drop_index("idx_adrq_asset_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 7aafddb3b94..750aac8d2b7 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -749,33 +749,40 @@ class TaskInletAssetReference(Base):
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"],
name="adrq_dag_fkey",
ondelete="CASCADE",
),
- Index("idx_asset_dag_run_queue_target_dag_id", target_dag_id),
+ Index("idx_adrq_asset_id", asset_id),
)
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/models/dag.py
b/airflow-core/src/airflow/models/dag.py
index 29407e53bb0..d6dfab06a27 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -36,7 +36,6 @@ from sqlalchemy import (
Integer,
String,
Text,
- and_,
case,
func,
or_,
@@ -60,7 +59,7 @@ from airflow._shared.timezones import timezone
from airflow.assets.evaluation import AssetEvaluator
from airflow.configuration import conf as airflow_conf
from airflow.exceptions import AirflowException
-from airflow.models.asset import AssetDagRunQueue, AssetModel
+from airflow.models.asset import AssetDagRunQueue
from airflow.models.base import Base, StringID
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
@@ -263,47 +262,6 @@ def get_last_dagrun(dag_id: str, session: Session,
include_manually_triggered: b
return session.scalar(query.limit(1))
-def get_asset_triggered_next_run_info(
- dag_ids: list[str], *, session: Session
-) -> dict[str, dict[str, int | str]]:
- """
- Get next run info for a list of dag_ids.
-
- Given a list of dag_ids, get string representing how close any that are
asset triggered are to
- their next run, e.g. "1 of 2 assets updated".
- """
- from airflow.models.asset import AssetDagRunQueue as ADRQ,
DagScheduleAssetReference
-
- return {
- x.dag_id: {
- "uri": x.uri,
- "ready": x.ready,
- "total": x.total,
- }
- for x in session.execute(
- select(
- DagScheduleAssetReference.dag_id,
- # This is a dirty hack to workaround group by requiring an
aggregate,
- # since grouping by asset is not what we want to do here...but
it works
- case((func.count() == 1, func.max(AssetModel.uri)),
else_="").label("uri"),
- func.count().label("total"),
- func.sum(case((ADRQ.target_dag_id.is_not(None), 1),
else_=0)).label("ready"),
- )
- .join(
- ADRQ,
- and_(
- ADRQ.asset_id == DagScheduleAssetReference.asset_id,
- ADRQ.target_dag_id == DagScheduleAssetReference.dag_id,
- ),
- isouter=True,
- )
- .join(AssetModel, AssetModel.id ==
DagScheduleAssetReference.asset_id)
- .group_by(DagScheduleAssetReference.dag_id)
- .where(DagScheduleAssetReference.dag_id.in_(dag_ids))
- ).all()
- }
-
-
class DagTag(Base):
"""A tag name per dag, to allow quick filtering in the DAG view."""
@@ -855,17 +813,6 @@ class DagModel(Base):
next_dagrun_partition_date=str(self.next_dagrun_partition_date),
)
- @provide_session
- def get_asset_triggered_next_run_info(
- self, *, session: Session = NEW_SESSION
- ) -> dict[str, int | str] | None:
- if self.asset_expression is None:
- return None
-
- # When an asset alias does not resolve into assets,
get_asset_triggered_next_run_info returns
- # an empty dict as there's no asset info to get. This method should
thus return None.
- return get_asset_triggered_next_run_info([self.dag_id],
session=session).get(self.dag_id, None)
-
@staticmethod
@cached(_team_name_cache, key=lambda dag_id, **_: dag_id,
lock=_team_name_cache_lock)
@provide_session
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 615ecb66c67..21ccd5ea3bc 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -117,7 +117,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"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 39054a28d70..462d1605af0 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
@@ -1505,7 +1505,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
@@ -2226,10 +2229,10 @@ class
TestDeleteAssetQueuedEvents(TestQueuedEventEndpoint):
(asset,) = self.create_assets(session=session, num=1)
self._create_asset_dag_run_queues(dag_id, asset.id, session)
- assert session.get(AssetDagRunQueue, (asset.id, dag_id)) is not None
+ assert session.scalars(select(AssetDagRunQueue)).all()
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.scalars(select(AssetDagRunQueue)).all() == []
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 05735154c5f..2427729035f 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 @@ class TestNextRunAssets:
)
}
# 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 @@ class TestNextRunAssets:
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 030a0d96c8b..3d9d75ff12a 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -25,8 +25,8 @@ from typing import TYPE_CHECKING
from unittest import mock
import pytest
-from sqlalchemy import delete, func, select
-from sqlalchemy.dialects import mysql
+from sqlalchemy import func, select
+from sqlalchemy.dialects import mysql, postgresql, sqlite
from sqlalchemy.orm import Session
from airflow import settings
@@ -67,7 +67,7 @@ pytestmark = pytest.mark.db_test
pytest.importorskip("pydantic", minversion="2.0.0")
[email protected]
[email protected](autouse=True)
def clear_assets():
from tests_common.test_utils.db import clear_db_assets
@@ -137,7 +137,6 @@ class TestAssetManager:
asm = AssetModel(uri="test://asset1/", name="test_asset_uri",
group="asset")
session.add(asm)
asm.scheduled_dags = [DagScheduleAssetReference(dag_id=dag.dag_id) for
dag in (dag1, dag2)]
- session.execute(delete(AssetDagRunQueue))
session.flush()
asset_manager.register_asset_change(task_instance=mock_task_instance,
asset=asset, session=session)
@@ -150,7 +149,6 @@ class TestAssetManager:
)
assert
session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2
- @pytest.mark.usefixtures("clear_assets")
def test_register_asset_change_with_alias(
self, session, dag_maker, mock_task_instance, testing_dag_bundle
):
@@ -173,7 +171,6 @@ class TestAssetManager:
DagScheduleAssetAliasReference(alias_id=asam.id, dag_id=dag.dag_id)
for dag in (consumer_dag_1, consumer_dag_2)
]
- session.execute(delete(AssetDagRunQueue))
session.flush()
asset = Asset(uri="test://asset1", name="test_asset_uri")
@@ -199,7 +196,6 @@ class TestAssetManager:
asset = Asset(uri="test://asset1", name="never_consumed")
asm = AssetModel(uri="test://asset1/", name="never_consumed",
group="asset")
session.add(asm)
- session.execute(delete(AssetDagRunQueue))
session.flush()
asset_manager.register_asset_change(task_instance=mock_task_instance,
asset=asset, session=session)
@@ -212,23 +208,42 @@ class TestAssetManager:
)
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.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 +254,45 @@ class TestAssetManager:
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_upsert_noop_on_mysql(self):
+ """On MySQL the helper upserts with a no-op ON DUPLICATE KEY UPDATE
(scoped to duplicate
+ keys, unlike 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 "ON DUPLICATE KEY UPDATE ASSET_ID =
ASSET_DAG_RUN_QUEUE.ASSET_ID" 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
@@ -381,7 +415,7 @@ class TestAssetManager:
assert len(set(ids)) == 1
assert
session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) == 1
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_suppresses_conflicting_partition_date(self,
session):
"""Two events resolving the same target key to different dates →
suppress to None.
@@ -416,7 +450,7 @@ class TestAssetManager:
assert second.id == first.id # same pending APDR
assert second.partition_date is None # conflict suppressed,
deterministic
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_keeps_agreeing_partition_date(self, session):
"""A later event carrying the same (or no) date does not trip the
conflict suppression."""
asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
@@ -444,7 +478,7 @@ class TestAssetManager:
assert with_none.id == first.id
assert with_none.partition_date == source_date
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_adopts_date_when_existing_is_none(self,
session):
"""An APDR created with no date adopts a later event's carried date
(not dropped)."""
asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
@@ -469,7 +503,7 @@ class TestAssetManager:
assert adopted.id == first.id
assert adopted.partition_date == source_date
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_recovers_after_conflict(self, session):
"""Once a conflict has suppressed the date to None, a later event
re-adopts a date."""
asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
@@ -497,7 +531,7 @@ class TestAssetManager:
assert recovered.id == first.id
assert recovered.partition_date == date_2
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_carry_partition_date_failure_degrades_to_none(self, session,
dag_maker, mock_task_instance):
"""A mapper whose carry_partition_date raises must not abort the write.
@@ -634,7 +668,6 @@ class TestAssetManager:
# Link the Stale Dag to the Asset
asm.scheduled_dags =
[DagScheduleAssetReference(dag_id=stale_dag.dag_id)]
- session.execute(delete(AssetDagRunQueue))
session.flush()
# Register the asset change
@@ -649,7 +682,7 @@ class TestAssetManager:
queued_id = session.scalar(select(AssetDagRunQueue.target_dag_id))
assert queued_id == "stale_dag"
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def
test_partitioned_asset_event_does_not_trigger_non_partitioned_dag(self,
session, mock_task_instance):
"""partitioned asset events (events with partition key) must not queue
non-partition-aware Dags."""
asm = AssetModel(uri="test://asset/", name="test_asset", group="asset")
@@ -659,7 +692,6 @@ class TestAssetManager:
)
session.add(dag)
asm.scheduled_dags = [DagScheduleAssetReference(dag_id=dag.dag_id)]
- session.execute(delete(AssetDagRunQueue))
session.flush()
AssetManager.register_asset_change(
@@ -683,7 +715,7 @@ class TestAssetManager:
pytest.param(6, True, id="one_over_cap_trips"),
],
)
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_partition_fan_out_cap(self, session, dag_maker,
mock_task_instance, cap, expect_trip):
"""The ``[scheduler] partition_mapper_max_downstream_keys`` cap gates
fan-out.
@@ -740,7 +772,7 @@ class TestAssetManager:
assert error_call.kwargs["cap_source"] == f"[scheduler]
partition_mapper_max_downstream_keys={cap}"
@conf_vars({("scheduler", "partition_mapper_max_downstream_keys"): "100"})
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_partition_fanout_per_mapper_override_stricter_than_global_trips(
self, session, dag_maker, mock_task_instance
):
@@ -787,7 +819,7 @@ class TestAssetManager:
assert error_call.kwargs["cap_source"] == "max_downstream_keys=3"
@conf_vars({("scheduler", "partition_mapper_max_downstream_keys"): "3"})
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_partition_fanout_per_mapper_override_looser_than_global_permits(
self, session, dag_maker, mock_task_instance
):
@@ -831,7 +863,7 @@ class TestAssetManager:
)
@conf_vars({("scheduler", "partition_mapper_max_downstream_keys"): "1"})
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_partition_fanout_per_mapper_at_cap_is_allowed(self, session,
dag_maker, mock_task_instance):
"""Per-mapper max_downstream_keys=7 with a 7-key fanout: exactly at
cap is allowed.
@@ -870,7 +902,7 @@ class TestAssetManager:
)
@conf_vars({("scheduler", "partition_mapper_max_downstream_keys"): "1"})
- @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle")
+ @pytest.mark.usefixtures("testing_dag_bundle")
def test_partition_fanout_per_mapper_one_over_cap_trips(self, session,
dag_maker, mock_task_instance):
"""Per-mapper max_downstream_keys=6 with a 7-key fanout: one over cap
trips the guard.
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 53ea5bc933f..c83d9039d8f 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.module_loading import qualname
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
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 @@ class TestSchedulerJob:
)
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 @@ class TestSchedulerJob:
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 ead765b3171..17fdac7b7ba 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -55,7 +55,6 @@ from airflow.models.dag import (
DagOwnerAttributes,
DagTag,
clear_team_name_cache,
- get_asset_triggered_next_run_info,
get_next_data_interval,
get_run_data_interval,
)
@@ -2585,7 +2584,12 @@ class TestDagModel:
# 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 +2643,10 @@ class TestDagModel:
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 +2714,14 @@ class TestDagModel:
)
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 +2765,14 @@ class TestDagModel:
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 +2806,12 @@ class TestDagModel:
# 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 +3069,24 @@ class TestDagModel:
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),
),
]
@@ -3806,71 +3841,6 @@ def test__time_restriction(dag_maker, dag_date,
tasks_date, catchup, restrict):
assert dag._time_restriction == restrict
-def test_get_asset_triggered_next_run_info(dag_maker, clear_assets):
- asset1 = Asset(uri="test://asset1", name="test_asset1", group="test-group")
- asset2 = Asset(uri="test://asset2", group="test-group")
- asset3 = Asset(uri="test://asset3", group="test-group")
- with dag_maker(dag_id="assets-1", schedule=[asset2]):
- pass
- dag1 = dag_maker.dag
-
- with dag_maker(dag_id="assets-2", schedule=[asset1, asset2]):
- pass
- dag2 = dag_maker.dag
-
- with dag_maker(dag_id="assets-3", schedule=[asset1, asset2, asset3]):
- pass
- dag3 = dag_maker.dag
-
- session = dag_maker.session
- asset1_id = session.scalar(select(AssetModel.id).where(AssetModel.uri ==
asset1.uri))
- 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),
- ]
- )
- session.flush()
-
- assets =
session.execute(select(AssetModel.uri).order_by(AssetModel.id)).all()
-
- info = get_asset_triggered_next_run_info([dag1.dag_id], session=session)
- assert info[dag1.dag_id] == {
- "ready": 0,
- "total": 1,
- "uri": assets[0].uri,
- }
-
- # This time, check both dag2 and dag3 at the same time (tests filtering)
- info = get_asset_triggered_next_run_info([dag2.dag_id, dag3.dag_id],
session=session)
- assert info[dag2.dag_id] == {
- "ready": 1,
- "total": 2,
- "uri": "",
- }
- assert info[dag3.dag_id] == {
- "ready": 1,
- "total": 3,
- "uri": "",
- }
-
-
[email protected]_serialized_dag
-def
test_get_asset_triggered_next_run_info_with_unresolved_asset_alias(dag_maker,
clear_assets):
- asset_alias1 = AssetAlias(name="alias")
- with dag_maker(dag_id="dag-1", schedule=[asset_alias1]):
- pass
- dag1 = dag_maker.dag
- session = dag_maker.session
- session.flush()
-
- info = get_asset_triggered_next_run_info([dag1.dag_id], session=session)
- assert info == {}
-
- dag1_model = DagModel.get_dagmodel(dag1.dag_id)
- assert dag1_model.get_asset_triggered_next_run_info(session=session) is
None
-
-
@pytest.mark.parametrize(
"run_id_type",
[DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED,
DagRunType.ASSET_TRIGGERED],
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 53fa99aa224..6a348a3f2ca 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 3c12c886f28..ccdc395ae41 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 @@ class TestAssetConditionWithTimetable:
# 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
