Skip to content

Reference asset events from asset_dag_run_queue #70972

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
uranusjr wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Reference asset events from asset_dag_run_queue #70972

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion airflow-core/docs/migrations-ref.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
Expand Down
128 changes: 21 additions & 107 deletions airflow-core/src/airflow/assets/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
from airflow.timetables.base import compute_rollup_fingerprint
from airflow.utils.helpers import is_container, prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import create_session
from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks

if TYPE_CHECKING:
Expand Down Expand Up @@ -120,38 +119,6 @@ def _lock_asset_model(
yield


def _create_asset_event(*, session: Session, **event_kwargs) -> AssetEvent:
"""
Persist an :class:`AssetEvent` row and return it, bound to *session*.
On SQLite the event is added directly to the caller's *session* and
flushed. SQLite serialises writes at the database-file level: opening
a second connection here would compete with any write locks the
caller's transaction already holds (for example, an UPDATE on
``dag_run`` flushed earlier in ``register_asset_changes_in_db``) and
deadlock with ``database is locked``.
On Postgres/MySQL a short-lived independent session is used so the
row is committed — and therefore visible to the scheduler's session
via MVCC — before the caller continues. The committed row is then
re-loaded into the caller's *session* so subsequent relationship
operations work correctly.
"""
if get_dialect_name(session) == "sqlite":
asset_event = AssetEvent(**event_kwargs)
session.add(asset_event)
session.flush()
return asset_event

with create_session(scoped=False) as ae_session:
asset_event = AssetEvent(**event_kwargs)
ae_session.add(asset_event)
ae_session.flush()
asset_event_id = asset_event.id

return session.get_one(AssetEvent, asset_event_id)


class AssetManager(LoggingMixin):
"""
A pluggable class that manages operations for assets.
Expand Down Expand Up @@ -374,7 +341,10 @@ def register_asset_change(
source_run_id=task_instance.run_id,
source_map_index=task_instance.map_index,
)
asset_event = _create_asset_event(session=session, **event_kwargs)

asset_event = AssetEvent(**event_kwargs)
session.add(asset_event)
session.flush()

dags_to_queue_from_asset = {ref.dag for ref in asset_model.scheduled_dags if not ref.dag.is_paused}

Expand Down Expand Up @@ -544,25 +514,7 @@ def _queue_dagruns(
if not non_partitioned_dags or partition_key is not None:
return None

# Possible race condition: if multiple dags or multiple (usually
# mapped) tasks update the same asset, this can fail with a unique
# constraint violation.
#
# Where the dialect supports a single-statement "insert, update on
# conflict" we use it; it is atomic, avoids the per-row SAVEPOINT churn,
# and holds locks for far less time (which on MySQL/InnoDB also makes the
# concurrent fan-out much less deadlock-prone). Otherwise we "fallback" to
# a nested transaction per row. Either way the rows are added in the same
# transaction where `ti.state` is changed.
dialect_name = get_dialect_name(session)
if TYPE_CHECKING:
assert dialect_name is not None
if dialect_name == "mysql":
return cls._queue_dagruns_nonpartitioned_mysql(asset_id, non_partitioned_dags, event, session)
# PostgreSQL and SQLite both support ON CONFLICT DO UPDATE.
return cls._queue_dagruns_nonpartitioned_conflict_update(
asset_id, non_partitioned_dags, event, session, dialect_name
)
return cls._queue_dagruns_nonpartitioned(asset_id, non_partitioned_dags, event, session)

@classmethod
def _queue_partitioned_dags(
Expand Down Expand Up @@ -827,71 +779,33 @@ def _get_or_create_apdr(
return apdr

@classmethod
def _queue_dagruns_nonpartitioned_slow_path(
def _queue_dagruns_nonpartitioned(
cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session
) -> None:
def _queue_dagrun_if_needed(dag: DagModel) -> str | None:
item = AssetDagRunQueue(target_dag_id=dag.dag_id, asset_id=asset_id, created_at=event.timestamp)
# Don't error whole transaction when a single RunQueue item conflicts.
# https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint
try:
with session.begin_nested():
existing = session.get(
AssetDagRunQueue, {"target_dag_id": dag.dag_id, "asset_id": asset_id}
)
if existing and existing.created_at >= event.timestamp:
cls.logger().debug("Skipping record %s due to newer timestamp", item)
return dag.dag_id # already queued with a newer timestamp
session.merge(item)
except exc.IntegrityError:
cls.logger().debug("Skipping record %s", item, exc_info=True)
return dag.dag_id

queued_results = (_queue_dagrun_if_needed(dag) for dag in dags_to_queue)
if queued_dag_ids := [r for r in queued_results if r is not None]:
cls.logger().debug("consuming dag ids %s", queued_dag_ids)

@classmethod
def _queue_dagruns_nonpartitioned_mysql(
cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent, session: Session
) -> None:
from sqlalchemy import case
from sqlalchemy.dialects.mysql import insert
if not dags_to_queue:
return
values = [
{"asset_id": asset_id, "target_dag_id": dag.dag_id, "asset_event_id": event.id}
for dag in dags_to_queue
]

values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue]
stmt = insert(AssetDagRunQueue).values(asset_id=asset_id, created_at=event.timestamp)
if (dialect_name := get_dialect_name(session)) == "mysql":
from sqlalchemy.dialects.mysql import insert as my_insert

update_stmt = stmt.on_duplicate_key_update(
created_at=case(
(stmt.inserted.created_at >= AssetDagRunQueue.created_at, stmt.inserted.created_at),
else_=AssetDagRunQueue.created_at,
)
)
session.execute(update_stmt, values)
session.execute(my_insert(AssetDagRunQueue).prefix_with("IGNORE"), values)

@kaxil kaxil Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

INSERT IGNORE downgrades more than duplicate-key errors to warnings on MySQL (FK violations, NOT NULL, truncation are all silently swallowed), while the postgres/sqlite path only ignores the PK conflict. A bad event id, for example, would raise on postgres but insert nothing here without a sound. on_duplicate_key_update(asset_id=stmt.inserted.asset_id) is the usual no-op trick that keeps the ignore scoped to duplicates.

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:
Expand Down
87 changes: 35 additions & 52 deletions airflow-core/src/airflow/jobs/scheduler_job_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from collections import Counter, defaultdict, deque
from collections.abc import Callable, Collection, Iterable, Iterator
from contextlib import ExitStack
from datetime import date, datetime, timedelta
from datetime import datetime, timedelta
from functools import lru_cache, partial
from itertools import groupby
from typing import TYPE_CHECKING, Any, cast
Expand Down Expand Up @@ -135,10 +135,10 @@
if TYPE_CHECKING:
from types import FrameType

from pendulum.datetime import DateTime
from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session
from sqlalchemy.orm.interfaces import LoaderOption
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.selectable import Subquery

from airflow._shared.logging.types import Logger
Expand Down Expand Up @@ -2641,9 +2641,7 @@ def _create_dag_runs_asset_triggered(

queued_adrqs = session.scalars(
with_row_locks(
select(AssetDagRunQueue)
.where(AssetDagRunQueue.target_dag_id == dag.dag_id)
.order_by(AssetDagRunQueue.created_at.desc()),
select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == dag.dag_id),
of=AssetDagRunQueue,
skip_locked=True,
key_share=False,
Expand All @@ -2658,53 +2656,31 @@ def _create_dag_runs_asset_triggered(
)
continue

triggered_date: DateTime = timezone.coerce_datetime(queued_adrqs[0].created_at)
self.log.debug(
"Creating asset-triggered DagRun for '%s': %d queued assets, triggered_date=%s",
dag.dag_id,
len(queued_adrqs),
triggered_date,
)
cte = (
select(func.max(DagRun.run_after).label("previous_dag_run_run_after"))
.where(
DagRun.dag_id == dag.dag_id,
DagRun.run_type == DagRunType.ASSET_TRIGGERED,
DagRun.run_after < triggered_date,
)
.cte()
)

# A first asset-triggered run has no previous run to floor the event window. With
# catchup off, floor it at when the Dag started scheduling on its assets so the
# backlog is skipped; with catchup on, only date.min applies and the backlog replays.
event_window_floor: list[Any] = [cte.c.previous_dag_run_run_after]
if not dag.catchup:
event_window_floor.append(
select(func.min(DagScheduleAssetReference.created_at))
.where(DagScheduleAssetReference.dag_id == dag.dag_id)
.scalar_subquery()
referenced_event_ids = {adrq.asset_event_id for adrq in queued_adrqs}
event_predicate: ColumnElement[bool] = AssetEvent.id.in_(referenced_event_ids)

@kaxil kaxil Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One behaviour delta worth confirming: the old select required the event's asset to still be in the dag's schedule (the DagScheduleAssetReference/alias join). Consuming by reference means a queued event for an asset that was since removed from the dag's schedule now gets consumed and shows up in the run's consumed_asset_events / triggering_asset_events, where before it was silently dropped along with its ADRQ row. Intended?

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)
Expand All @@ -2719,6 +2695,13 @@ def _create_dag_runs_asset_triggered(
)
)
if asset_events:
triggered_date = timezone.coerce_datetime(max(event.timestamp for event in asset_events))
self.log.debug(
"Creating asset-triggered DagRun for '%s': %d queued assets, triggered_date=%s",
dag.dag_id,
len(queued_adrqs),
triggered_date,
)
dag_run = dag.create_dagrun(
run_id=DagRun.generate_run_id(
run_type=DagRunType.ASSET_TRIGGERED, logical_date=None, run_after=triggered_date
Expand Down Expand Up @@ -2747,19 +2730,19 @@ def _create_dag_runs_asset_triggered(
)
else:
self.log.info(
"No DagRun created for '%s' at '%s' - asset events already consumed or none found",
"No DagRun created for '%s' - asset events already consumed or none found",
dag.dag_id,
triggered_date,
)
# Always delete ADRQ rows for this batch to prevent stale entries accumulating,
# including when all events were already consumed by a concurrent DagRun.
adrq_pks = [(record.asset_id, record.target_dag_id) for record in queued_adrqs]
result = cast(
"CursorResult",
session.execute(
delete(AssetDagRunQueue).where(
tuple_(AssetDagRunQueue.asset_id, AssetDagRunQueue.target_dag_id).in_(adrq_pks),
AssetDagRunQueue.created_at <= triggered_date,
tuple_(
AssetDagRunQueue.target_dag_id,
AssetDagRunQueue.asset_event_id,
).in_((adrq.target_dag_id, adrq.asset_event_id) for adrq in queued_adrqs)
)
),
)
Expand Down
Loading
Loading