This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 9e5b7701f8d Fix asset event scheduling race condition and consistency
(#62501)
9e5b7701f8d is described below
commit 9e5b7701f8ded5bca18d665ecc23ae1294a4f1b1
Author: Stanley Law <[email protected]>
AuthorDate: Wed Jul 29 16:00:44 2026 +0800
Fix asset event scheduling race condition and consistency (#62501)
Co-authored-by: Kaxil Naik <[email protected]>
Co-authored-by: Copilot <[email protected]>
---
airflow-core/newsfragments/62501.bugfix.rst | 1 +
airflow-core/src/airflow/assets/manager.py | 107 +++++++---
.../src/airflow/jobs/scheduler_job_runner.py | 80 +++++---
airflow-core/tests/unit/assets/test_manager.py | 20 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 220 ++++++++++++++++++++-
.../tests/unit/models/test_taskinstance.py | 33 ++++
6 files changed, 391 insertions(+), 70 deletions(-)
diff --git a/airflow-core/newsfragments/62501.bugfix.rst
b/airflow-core/newsfragments/62501.bugfix.rst
new file mode 100644
index 00000000000..bb2a30511d4
--- /dev/null
+++ b/airflow-core/newsfragments/62501.bugfix.rst
@@ -0,0 +1 @@
+Fix a race condition that could cause asset-triggered Dags to be skipped when
asset events were produced concurrently. The scheduler could miss a newly
produced asset event because there was a visibility gap between when the record
was flushed and when the producing task's transaction committed. The asset
event is now committed ahead of the producing task's remaining bookkeeping to
close that gap; in rare cases this can surface as a duplicate or orphaned asset
event if that task fails im [...]
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index 1dd26379b9d..d0f8470b17c 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -46,6 +46,7 @@ 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:
@@ -119,6 +120,38 @@ 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.
@@ -341,10 +374,7 @@ class AssetManager(LoggingMixin):
source_run_id=task_instance.run_id,
source_map_index=task_instance.map_index,
)
-
- asset_event = AssetEvent(**event_kwargs)
- session.add(asset_event)
- session.flush() # Ensure the event is written earlier than ADRQ
entries below.
+ asset_event = _create_asset_event(session=session, **event_kwargs)
dags_to_queue_from_asset = {ref.dag for ref in
asset_model.scheduled_dags if not ref.dag.is_paused}
@@ -518,18 +548,21 @@ class AssetManager(LoggingMixin):
# mapped) tasks update the same asset, this can fail with a unique
# constraint violation.
#
- # Where the dialect supports a single-statement "insert, ignore on
+ # 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 dialect_name == "postgresql":
- return cls._queue_dagruns_nonpartitioned_postgres(asset_id,
non_partitioned_dags, 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, session)
- return cls._queue_dagruns_nonpartitioned_slow_path(asset_id,
non_partitioned_dags, session)
+ 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
+ )
@classmethod
def _queue_partitioned_dags(
@@ -795,14 +828,20 @@ class AssetManager(LoggingMixin):
@classmethod
def _queue_dagruns_nonpartitioned_slow_path(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ 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)
+ 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)
@@ -813,28 +852,46 @@ class AssetManager(LoggingMixin):
cls.logger().debug("consuming dag ids %s", queued_dag_ids)
@classmethod
- def _queue_dagruns_nonpartitioned_postgres(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ def _queue_dagruns_nonpartitioned_mysql(
+ cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent,
session: Session
) -> None:
- from sqlalchemy.dialects.postgresql import insert
+ 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).on_conflict_do_nothing()
- session.execute(stmt, values)
+ 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,
+ )
+ )
+ session.execute(update_stmt, values)
@classmethod
- def _queue_dagruns_nonpartitioned_mysql(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ def _queue_dagruns_nonpartitioned_conflict_update(
+ cls,
+ asset_id: int,
+ dags_to_queue: set[DagModel],
+ event: AssetEvent,
+ session: Session,
+ dialect_name: str,
) -> None:
- from sqlalchemy.dialects.mysql import insert
+ """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)
- # MySQL has no "ON CONFLICT DO NOTHING"; a no-op ON DUPLICATE KEY
UPDATE turns a
- # conflicting (asset_id, target_dag_id) row into a no-op rather than
an error,
- # matching the Postgres path.
- stmt =
stmt.on_duplicate_key_update(target_dag_id=stmt.inserted.target_dag_id)
- session.execute(stmt, values)
+ 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(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 daef05b2b7c..5f6d17ff360 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -85,6 +85,7 @@ from airflow.models.asset import (
PartitionedAssetKeyLog,
TaskInletAssetReference,
TaskOutletAssetReference,
+ association_table,
)
from airflow.models.asset_state_store import AssetStateStoreModel
from airflow.models.backfill import Backfill, BackfillDagRun
@@ -2702,48 +2703,63 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
),
),
- AssetEvent.timestamp <= triggered_date,
AssetEvent.timestamp >
func.coalesce(*event_window_floor),
+ AssetEvent.timestamp <= triggered_date,
+ ~(
+ select(association_table.c.event_id)
+ .join(DagRun, DagRun.id ==
association_table.c.dag_run_id)
+ .where(
+ DagRun.dag_id == dag.dag_id,
+ association_table.c.event_id == AssetEvent.id,
+ )
+ .exists()
+ ),
)
.order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
)
)
-
- dag_run = dag.create_dagrun(
- run_id=DagRun.generate_run_id(
- run_type=DagRunType.ASSET_TRIGGERED, logical_date=None,
run_after=triggered_date
- ),
- logical_date=None,
- data_interval=None,
- run_after=triggered_date,
- run_type=DagRunType.ASSET_TRIGGERED,
- triggered_by=DagRunTriggeredByType.ASSET,
- state=DagRunState.QUEUED,
- creating_job_id=self.job.id,
- session=session,
- )
- team_name = (
- self._get_team_names_for_dag_ids([dag.dag_id],
session).get(dag.dag_id)
- if self._multi_team
- else None
- )
- stats.incr("asset.triggered_dagruns",
tags=prune_dict({"team_name": team_name}))
- dag_run.consumed_asset_events.extend(asset_events)
- self.log.info(
- "Created asset-triggered DagRun for '%s': run_id=%s, consumed
%d asset events",
- dag.dag_id,
- dag_run.run_id,
- len(asset_events),
- )
-
- # Delete only consumed ADRQ rows to avoid dropping newly queued
events
- # (e.g. DagRun triggered by asset A while a new event for asset B
arrives).
+ if asset_events:
+ dag_run = dag.create_dagrun(
+ run_id=DagRun.generate_run_id(
+ run_type=DagRunType.ASSET_TRIGGERED,
logical_date=None, run_after=triggered_date
+ ),
+ logical_date=None,
+ data_interval=None,
+ run_after=triggered_date,
+ run_type=DagRunType.ASSET_TRIGGERED,
+ triggered_by=DagRunTriggeredByType.ASSET,
+ state=DagRunState.QUEUED,
+ creating_job_id=self.job.id,
+ session=session,
+ )
+ team_name = (
+ self._get_team_names_for_dag_ids([dag.dag_id],
session).get(dag.dag_id)
+ if self._multi_team
+ else None
+ )
+ stats.incr("asset.triggered_dagruns",
tags=prune_dict({"team_name": team_name}))
+ dag_run.consumed_asset_events.extend(asset_events)
+ self.log.info(
+ "Created asset-triggered DagRun for '%s': run_id=%s,
consumed %d asset events",
+ dag.dag_id,
+ dag_run.run_id,
+ len(asset_events),
+ )
+ else:
+ self.log.info(
+ "No DagRun created for '%s' at '%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)
+ tuple_(AssetDagRunQueue.asset_id,
AssetDagRunQueue.target_dag_id).in_(adrq_pks),
+ AssetDagRunQueue.created_at <= triggered_date,
)
),
)
diff --git a/airflow-core/tests/unit/assets/test_manager.py
b/airflow-core/tests/unit/assets/test_manager.py
index c290d7f5a63..bc034cd7515 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -215,15 +215,16 @@ class TestAssetManager:
@pytest.mark.parametrize(
("dialect_name", "expected_helper"),
[
- ("postgresql", "_queue_dagruns_nonpartitioned_postgres"),
+ ("postgresql", "_queue_dagruns_nonpartitioned_conflict_update"),
("mysql", "_queue_dagruns_nonpartitioned_mysql"),
- ("sqlite", "_queue_dagruns_nonpartitioned_slow_path"),
+ ("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."""
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"),
@@ -234,18 +235,25 @@ class TestAssetManager:
dags_to_queue={dag},
partition_key=None,
partition_date=None,
- event=mock.MagicMock(),
+ event=event,
task_instance=None,
session=session,
)
- mock_helper.assert_called_once_with(1, {dag}, 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}")
def test_queue_dagruns_nonpartitioned_mysql_builds_upsert(self):
"""Test that the MySQL queue path emits an INSERT ... ON DUPLICATE KEY
UPDATE."""
dag = DagModel(dag_id="dag1")
session = mock.MagicMock(spec=Session)
-
- AssetManager._queue_dagruns_nonpartitioned_mysql(asset_id=1,
dags_to_queue={dag}, session=session)
+ event = AssetEvent(asset_id=1)
+ AssetManager._queue_dagruns_nonpartitioned_mysql(
+ 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()
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 4add965a6ab..1fb72d1df13 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -22,8 +22,10 @@ import datetime
import logging
import os
import re
+import time
from collections import Counter, deque
from collections.abc import Callable, Generator, Iterator
+from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import ExitStack, contextmanager
from datetime import timedelta
from pathlib import Path
@@ -45,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
+from airflow.assets.manager import AssetManager, _create_asset_event
from airflow.callbacks.callback_requests import (
DagCallbackRequest,
DagRunContext,
@@ -5788,6 +5790,175 @@ 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_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")
+ with dag_maker(dag_id="asset-consumer-stale-trigger-date",
schedule=[asset], session=session):
+ pass
+ 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))
+ session.flush()
+
+ # Simulate another scheduler consuming ADRQ rows after we computed
triggered_date_by_dag.
+
session.execute(delete(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id
== dag_model.dag_id))
+ session.flush()
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[self.null_exec])
+ self.job_runner._create_dag_runs_asset_triggered(
+ dag_models=[dag_model],
+ session=session,
+ )
+
+ # We do not create a new DagRun since the ADRQ has already been
consumed
+ assert session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)).one_or_none() is None
+
+ @pytest.mark.need_serialized_dag
+ def test_create_dag_runs_asset_triggered_deletes_only_selected_adrq_rows(
+ self, session: Session, dag_maker
+ ):
+ asset_1 = Asset("ready-to-trigger-a-Dag-run")
+ asset_2 = Asset("should-still-exist-after-a-Dag-run-created")
+ with dag_maker(dag_id="asset-consumer-delete-selected",
schedule=asset_1 | asset_2, session=session):
+ pass
+ 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))
+ session.add_all(
+ [
+ AssetEvent(
+ asset_id=asset_1_id,
+ timestamp=timezone.utcnow(),
+ ),
+ # The ADRQ that should triggers the Dag run creation
+ AssetDagRunQueue(
+ asset_id=asset_1_id, target_dag_id=dag_model.dag_id,
created_at=timezone.utcnow()
+ ),
+ 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
+ AssetDagRunQueue(
+ asset_id=asset_2_id, target_dag_id=dag_model.dag_id,
created_at=timezone.utcnow()
+ ),
+ ]
+ )
+ session.flush()
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[MockExecutor(do_update=False)])
+
+ def _lock_only_selected_asset(query, **_):
+ # Simulate SKIP LOCKED behavior where this scheduler can only
consume one ADRQ row.
+ return query.where(AssetDagRunQueue.asset_id == asset_1_id)
+
+ with patch("airflow.jobs.scheduler_job_runner.with_row_locks",
side_effect=_lock_only_selected_asset):
+ self.job_runner._create_dag_runs_asset_triggered(
+ dag_models=[dag_model],
+ session=session,
+ )
+
+ dr = session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)).one_or_none()
+ assert dr is not None
+
+ adrq_1 = session.scalars(
+ select(AssetDagRunQueue).where(
+ AssetDagRunQueue.target_dag_id == dag_model.dag_id,
+ AssetDagRunQueue.asset_id == asset_1_id,
+ )
+ ).one_or_none()
+ assert adrq_1 is None
+ adrq_2 = session.scalars(
+ select(AssetDagRunQueue).where(
+ AssetDagRunQueue.target_dag_id == dag_model.dag_id,
+ AssetDagRunQueue.asset_id == asset_2_id,
+ )
+ ).one_or_none()
+ assert adrq_2 is not None
+
+ @pytest.mark.need_serialized_dag
+ @pytest.mark.backend("postgres", "mysql")
+ def test_create_dag_runs_when_concurrent_asset_events_created(self,
session: Session, dag_maker, caplog):
+
+ ASSET_EVENT_COUNT = 30
+ asset = Asset(name="test_asset")
+ with dag_maker(dag_id="consumer", schedule=asset, session=session):
+ pass
+ dag_model = dag_maker.dag_model
+ # Capture the dag_id as a plain string in the main thread. The worker
threads must not
+ # touch this ORM object: it is bound to the main thread's session,
which the loop below
+ # commits (and thus expires) concurrently, so any attribute access
from a worker would
+ # load against another thread's session.
+ consumer_dag_id = dag_model.dag_id
+ with dag_maker(dag_id="asset-producer", start_date=timezone.utcnow(),
session=session):
+ BashOperator(task_id="simulate-asset-outlet", bash_command="echo
1")
+ dag_maker.create_dagrun(run_id="asset-producer-run")
+ asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri
== asset.uri))
+ futures = []
+ consumed_asset_events = []
+ asset_event_metadata: list[tuple[int, datetime.datetime]] = []
+
+ def create_asset_events(sleep):
+
+ with create_session() as session:
+ # Re-fetch the DagModel in this thread's own session so all
ORM access stays
+ # thread-local.
+ 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)
+ 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
+ )
+
+ return asset_event.id, now.isoformat()
+
+ with (
+ ThreadPoolExecutor(max_workers=3) as executor,
+ caplog.at_level(
+ "WARNING",
+ logger="airflow.jobs.scheduler_job_runner",
+ ),
+ ):
+ for i in range(ASSET_EVENT_COUNT):
+ # Deterministically alternate between fast (0s) and slow (2s)
workers so the
+ # test reliably exercises both code paths without relying on
RNG.
+ future = executor.submit(create_asset_events, i % 3)
+ futures.append(future)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[MockExecutor(do_update=False)])
+ seen_dr_ids: set[int] = set()
+ for future in as_completed(futures, timeout=120):
+ asset_event_metadata.append(future.result())
+ self.job_runner._create_dag_runs_asset_triggered(
+ dag_models=[dag_model],
+ session=session,
+ )
+ session.commit()
+ all_drs = session.scalars(select(DagRun).where(DagRun.dag_id
== dag_model.dag_id)).all()
+ for dr in all_drs:
+ if dr.id not in seen_dr_ids:
+ seen_dr_ids.add(dr.id)
+ consumed_asset_events += dr.consumed_asset_events
+ total_consumed_asset_events = len(consumed_asset_events)
+ assert total_consumed_asset_events == ASSET_EVENT_COUNT
+ assert len({event.id for event in consumed_asset_events}) ==
total_consumed_asset_events, (
+ "Expected no duplicated Asset event consumed"
+ )
+
@pytest.mark.need_serialized_dag
def test_create_dag_runs_asset_alias_with_asset_event_attached(self,
session, dag_maker):
"""
@@ -5897,6 +6068,10 @@ class TestSchedulerJob:
BashOperator(task_id="task", bash_command="echo 1",
outlets=[asset])
dr = dag_maker.create_dagrun()
+ with dag_maker(
+ dag_id=f"consumer_{suffix}", schedule=[asset],
bundle_name=bundle_name, session=session
+ ):
+ pass
asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri
== asset.uri))
event = AssetEvent(
asset_id=asset_id,
@@ -5906,12 +6081,8 @@ class TestSchedulerJob:
source_map_index=-1,
)
session.add(event)
-
- with dag_maker(
- dag_id=f"consumer_{suffix}", schedule=[asset],
bundle_name=bundle_name, session=session
- ):
- pass
-
+ # flush here to ensure event timestamp is before the ADRQ created_at
timestamp
+ session.flush()
session.add(AssetDagRunQueue(asset_id=asset_id,
target_dag_id=f"consumer_{suffix}"))
session.flush()
@@ -5981,6 +6152,41 @@ class TestSchedulerJob:
assert len(session.scalars(adrq_q).all()) == 1
assert session.scalars(adrq_q).one().target_dag_id == "consumer"
+ @pytest.mark.need_serialized_dag
+ def test_no_create_dag_runs_when_no_asset_event(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()
+ )
+ session.add(adrq)
+ session.flush()
+ adrq.created_at = timezone.utcnow() + timedelta(seconds=1)
+ session.merge(adrq)
+ with caplog.at_level("INFO"):
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[MockExecutor(do_update=False)])
+ self.job_runner._create_dag_runs_asset_triggered(
+ dag_models=[dag_model],
+ session=session,
+ )
+ 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
+ _adrq = session.scalars(
+ select(AssetDagRunQueue).where(
+ AssetDagRunQueue.asset_id == asset_id,
AssetDagRunQueue.target_dag_id == dag_model.dag_id
+ )
+ ).one_or_none()
+ # ADRQ is deleted even when no DagRun is created, to prevent stale
rows accumulating.
+ assert _adrq is None
+
@time_machine.travel(DEFAULT_DATE + datetime.timedelta(days=1, seconds=9),
tick=False)
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
def test_start_dagruns(self, mock_get_backend, dag_maker, session):
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index e0d4a1063e9..53fa99aa224 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -3875,6 +3875,39 @@ def
test_runtime_partition_key_does_not_backfill_dag_run_when_none(dag_maker, se
assert dr.partition_key is None
[email protected]("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.
+ """
+ asset = Asset(name="hello")
+ with dag_maker(dag_id="rt_pk_backfill_sqlite",
schedule=PartitionedAtRuntime()) as dag:
+ EmptyOperator(task_id="hi", outlets=[asset])
+ dr = dag_maker.create_dagrun(session=session)
+ assert dr.partition_key is None
+ [ti] = dr.get_task_instances(session=session)
+
+ # Must not raise "sqlite3.OperationalError: database is locked."
+ TaskInstance.register_asset_changes_in_db(
+ ti=ti,
+ task_outlets=[ensure_serialized_asset(asset).asprofile()],
+ outlet_events=[
+ {"dest_asset_key": {"name": "hello", "uri": "hello"}, "extra": {},
"partition_key": "us"},
+ ],
+ session=session,
+ )
+ event = session.scalar(select(AssetEvent).where(AssetEvent.source_dag_id
== dag.dag_id))
+ assert event.partition_key == "us"
+ session.refresh(dr)
+ assert dr.partition_key is None
+
+
def
test_runtime_partition_key_does_not_overwrite_scheduler_partition(dag_maker,
session):
"""Task-emitted key lands on the AssetEvent but does NOT overwrite a
scheduler-set DagRun.partition_key."""
asset = Asset(name="hello")
