Skip to content

Fix duplicate pending Dag runs for the same asset partition key #71074

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
anmolxlight wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix duplicate pending Dag runs for the same asset partition key #71074

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4faf5ba
Fix duplicate pending Dag runs for the same asset partition key
Aug 4, 2026
7bc084a
Deduplicate pending asset-partition dag runs without locking the Dag …
Aug 4, 2026
ab4226c
Update migration reference doc and revision heads map for pending APD…
Aug 4, 2026
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). |
| ``ee86eed19e24`` (head) | ``7a98f1b7dbd3`` | ``3.4.0`` | Add pending_partition_key to asset_partition_dag_run and |
| | | | enforce single pending row per key. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``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
1 change: 1 addition & 0 deletions airflow-core/newsfragments/71074.bugfix.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix duplicate pending Dag runs for the same asset partition key
248 changes: 124 additions & 124 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 @@ -18,7 +18,6 @@
from __future__ import annotations

from collections.abc import Callable, Collection, Iterable
from contextlib import contextmanager
from functools import partial
from typing import TYPE_CHECKING

Expand Down Expand Up @@ -49,7 +48,7 @@
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
from airflow.utils.sqlalchemy import get_dialect_name

if TYPE_CHECKING:
from datetime import datetime
Expand All @@ -68,60 +67,6 @@
log = structlog.get_logger(__name__)


@contextmanager
def _lock_asset_model(
*,
session: Session,
asset_id: int,
max_retries: int = 10,
retry_delay: float = 0.1,
):
"""
Context manager to acquire a lock for AssetPartitionDagRun creation.

- SQLite: Use a no-op ORM update to trigger a write-transaction and acquire SQLite's global writer lock.
- Postgres/MySQL: uses row-level lock on AssetModel.
"""
if get_dialect_name(session) == "sqlite":
import time

from sqlalchemy import update

# no-op update
# This is used to acquire SQLite's global writer lock.
stmt = update(AssetModel).where(AssetModel.id == asset_id).values(id=AssetModel.id)
for _ in range(max_retries):
try:
session.execute(stmt)
session.flush()
except exc.OperationalError as err:
err_msg = str(err).lower()
if "locked" in err_msg or "busy" in err_msg:
session.rollback()
time.sleep(retry_delay)
continue

# lock acquired
yield
return

raise RuntimeError(f"Could not acquire SQLite AssetModel writer lock for asset_id={asset_id}")
else:
# Postgres/MySQL row-level lock
if (
session.scalar(
with_row_locks(
query=select(AssetModel.id).where(AssetModel.id == asset_id),
session=session,
key_share=True,
)
)
) is None:
raise RuntimeError(f"Asset {asset_id} does not exist – cannot lock.")

yield


def _create_asset_event(*, session: Session, **event_kwargs) -> AssetEvent:
"""
Persist an :class:`AssetEvent` row and return it, bound to *session*.
Expand Down Expand Up @@ -733,7 +678,6 @@ def _queue_partitioned_dags(
target_partition_date=target_partition_date,
target_dag=target_dag,
rollup_fingerprint=fingerprint,
asset_id=asset_id,
session=session,
)
log_record = PartitionedAssetKeyLog(
Expand All @@ -754,27 +698,124 @@ def _get_or_create_apdr(
target_partition_date: datetime | None,
target_dag: DagModel,
rollup_fingerprint: dict,
asset_id: int,
session: Session,
) -> AssetPartitionDagRun:
"""
Get or create an APDR.

If 2 processes invoke this method at the same time using the same (target_key, target_dag) pair,
they may both check the database and, finding no existing APDR, create separate instances.
This leads to the unintended outcome of having two APDRs created instead of one.
To resolve this, we add a mutex lock to AssetModel for PostgreSQL and MySQL and use
AssetPartitionDagRunMutexLock table for SQLite.
If 2 processes invoke this method at the same time using the same (target_key, target_dag)
pair, they may both check the database and, finding no existing APDR, attempt to create
separate instances. Rather than serializing this find-or-create behind a lock, a unique
constraint on (target_dag_id, pending_partition_key) — see the ``AssetPartitionDagRun``
docstring — makes the database itself reject the loser's INSERT. The loser catches that
``IntegrityError`` and re-selects, working on the winning row instead of raising, per the
model's "always work on the latest matching APDR record" contract. Optimistic and
lock-free, this scales with concurrent producer assets without contending on a Dag or
Asset row that has nothing to do with the (target_key, target_dag) pair being deduplicated.

``rollup_fingerprint`` is the serialized mapper / window definition for all partitioned
assets in the timetable at creation time; the scheduler discards APDRs whose stamp no
longer matches the current timetable's fingerprint (mapper / window may have changed).
"""
latest_apdr = cls._get_latest_pending_apdr(
target_key=target_key, target_dag_id=target_dag.dag_id, session=session
)
if latest_apdr is not None:
cls._reconcile_partition_date(
apdr=latest_apdr,
target_partition_date=target_partition_date,
target_dag_id=target_dag.dag_id,
target_key=target_key,
session=session,
)
cls.logger().debug(
"Existing APDR found for key %s dag_id %s",
target_key,
target_dag.dag_id,
exc_info=True,
)
return latest_apdr

apdr = AssetPartitionDagRun(
target_dag_id=target_dag.dag_id,
created_dag_run_id=None,
partition_key=target_key,
pending_partition_key=target_key,
partition_date=target_partition_date,
rollup_fingerprint=rollup_fingerprint,
)
try:
# A SAVEPOINT scopes the potential IntegrityError so only this INSERT is rolled
# back on conflict; the caller's surrounding transaction (with any other work
# already flushed in this scheduler tick) stays intact.
with session.begin_nested():
session.add(apdr)
session.flush()
except exc.IntegrityError:
cls.logger().debug(
"Lost race creating APDR for key %s dag_id %s; using the winning APDR instead",
target_key,
target_dag.dag_id,
exc_info=True,
)
winner = cls._get_latest_pending_apdr(
target_key=target_key, target_dag_id=target_dag.dag_id, session=session
)
if winner is None:
raise RuntimeError(
f"APDR insert for target_dag_id={target_dag.dag_id!r} "
f"partition_key={target_key!r} failed with an integrity error, but no "
"existing pending APDR was found on re-select."
)
cls._reconcile_partition_date(
apdr=winner,
target_partition_date=target_partition_date,
target_dag_id=target_dag.dag_id,
target_key=target_key,
session=session,
)
return winner

cls.logger().debug(
"No existing APDR found. Create APDR for key %s dag_id %s",
target_key,
target_dag.dag_id,
exc_info=True,
)
return apdr

Reconciling the carried ``partition_date`` on an existing pending APDR is best-effort:
a partitioned consumer's feeding assets are expected to agree on the partition's
datetime. The carry only matters for ``IdentityMapper`` (whose key the scheduler
cannot decode); temporal/composite feeds re-derive the date from the key at run
creation regardless of what is stored here. Within that contract:
@classmethod
def _get_latest_pending_apdr(
cls, *, target_key: str, target_dag_id: str, session: Session
) -> AssetPartitionDagRun | None:
"""Return the latest still-pending APDR for (target_dag_id, target_key), if any."""
return session.scalar(
select(AssetPartitionDagRun)
.where(
AssetPartitionDagRun.pending_partition_key == target_key,
AssetPartitionDagRun.target_dag_id == target_dag_id,
)
.order_by(AssetPartitionDagRun.id.desc())
.limit(1)
)

@classmethod
def _reconcile_partition_date(
cls,
*,
apdr: AssetPartitionDagRun,
target_partition_date: datetime | None,
target_dag_id: str,
target_key: str,
session: Session,
) -> None:
"""
Reconcile the carried ``partition_date`` on an existing pending APDR.

This is best-effort: a partitioned consumer's feeding assets are expected to agree
on the partition's datetime. The carry only matters for ``IdentityMapper`` (whose key
the scheduler cannot decode); temporal/composite feeds re-derive the date from the key
at run creation regardless of what is stored here. Within that contract:

- If the APDR carries no date yet (``None`` — created by an event that carried none),
adopt the incoming date when this event carries one. There is nothing to conflict
Expand All @@ -784,64 +825,23 @@ def _get_or_create_apdr(
carried date is suppressed to ``None`` (and re-adoptable by a later event).
- Otherwise (the dates agree, or this event carries none) the existing value is kept.
"""
with _lock_asset_model(session=session, asset_id=asset_id):
latest_apdr: AssetPartitionDagRun | None = session.scalar(
select(AssetPartitionDagRun)
.where(
AssetPartitionDagRun.partition_key == target_key,
AssetPartitionDagRun.target_dag_id == target_dag.dag_id,
)
.order_by(AssetPartitionDagRun.id.desc())
.limit(1)
)
if latest_apdr and latest_apdr.created_dag_run_id is None:
existing_partition_date = latest_apdr.partition_date
if existing_partition_date is None:
# No carried date yet; adopt the incoming one if present (no conflict
# to resolve). Keeps a later identity event's date from being dropped.
if target_partition_date is not None:
latest_apdr.partition_date = target_partition_date
session.flush()
elif target_partition_date is not None and existing_partition_date != target_partition_date:
# Two contributing events carry conflicting partition_dates for the same
# (target_key, target_dag). Choosing one would be order-dependent, so
# suppress: the consumer DagRun gets partition_date=None rather than a
# wrong, unstable value.
log.warning(
"Conflicting partition_date carried for the same target key; "
"suppressing it so the consumer DagRun's partition_date is None. "
"The producing assets likely disagree on the partition's datetime.",
target_dag_id=target_dag.dag_id,
target_key=target_key,
existing_partition_date=existing_partition_date,
incoming_partition_date=target_partition_date,
)
latest_apdr.partition_date = None
session.flush()
cls.logger().debug(
"Existing APDR found for key %s dag_id %s",
target_key,
target_dag.dag_id,
exc_info=True,
)
return latest_apdr

apdr = AssetPartitionDagRun(
target_dag_id=target_dag.dag_id,
created_dag_run_id=None,
partition_key=target_key,
partition_date=target_partition_date,
rollup_fingerprint=rollup_fingerprint,
existing_partition_date = apdr.partition_date
if existing_partition_date is None:
if target_partition_date is not None:
apdr.partition_date = target_partition_date
session.flush()
elif target_partition_date is not None and existing_partition_date != target_partition_date:
log.warning(
"Conflicting partition_date carried for the same target key; "
"suppressing it so the consumer DagRun's partition_date is None. "
"The producing assets likely disagree on the partition's datetime.",
target_dag_id=target_dag_id,
target_key=target_key,
existing_partition_date=existing_partition_date,
incoming_partition_date=target_partition_date,
)
session.add(apdr)
apdr.partition_date = None
session.flush()
cls.logger().debug(
"No existing APDR found. Create APDR for key %s dag_id %s",
target_key,
target_dag.dag_id,
exc_info=True,
)
return apdr

@classmethod
def _queue_dagruns_nonpartitioned_slow_path(
Expand Down
Loading
Loading