diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index dfc96be79019e..e2cf3d9c6ee70 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). | +| ``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. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/newsfragments/71074.bugfix.rst b/airflow-core/newsfragments/71074.bugfix.rst new file mode 100644 index 0000000000000..b20efc4bdac06 --- /dev/null +++ b/airflow-core/newsfragments/71074.bugfix.rst @@ -0,0 +1 @@ +Fix duplicate pending Dag runs for the same asset partition key diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index 9330c6071dbee..85465b5bc220d 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -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 @@ -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 @@ -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*. @@ -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( @@ -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 @@ -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( diff --git a/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_pending_partition_key_to_apdr.py b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_pending_partition_key_to_apdr.py new file mode 100644 index 0000000000000..86a9756d9d5e5 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_pending_partition_key_to_apdr.py @@ -0,0 +1,116 @@ +# +# 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. + +""" +Add pending_partition_key to asset_partition_dag_run and enforce single pending row per key. + +Two asset events from different producer assets that resolve to the same downstream +partition key could each create their own AssetPartitionDagRun (APDR), leaving two +pending rows that could never both be satisfied (apache/airflow#71070). This is now +prevented with a unique constraint on (target_dag_id, pending_partition_key). + +A partial/filtered unique index on (target_dag_id, partition_key) WHERE +created_dag_run_id IS NULL would be the more direct fix, but MySQL supports neither +partial nor filtered indexes. ``pending_partition_key`` is the portable equivalent: it +mirrors ``partition_key`` only while ``created_dag_run_id`` is null, and is null once +the dag run is created. A unique index treats null as distinct from every other value +on all three supported backends, so completed rows never collide with each other while +pending rows for the same key do. + +Pre-existing duplicate pending rows can never both be satisfied -- the asset events +that would complete them are necessarily split across the duplicates -- so before the +constraint is created, all but the latest (highest id) pending row per +(target_dag_id, partition_key) is dropped, along with its PartitionedAssetKeyLog rows. +This mirrors the scheduler's stale-APDR cleanup and the model docstring's "always work +on the latest matching APDR record" fallback. + +Revision ID: ee86eed19e24 +Revises: 7a98f1b7dbd3 +Create Date: 2026-08-04 00:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import context, op + +from airflow.migrations.db_types import StringID +from airflow.migrations.utils import disable_sqlite_fkeys + +revision = "ee86eed19e24" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + +_TABLE = "asset_partition_dag_run" +_LOG_TABLE = "partitioned_asset_key_log" +_UQ_NAME = "apdr_target_dag_id_pending_partition_key_uq" + + +def _drop_stale_duplicate_pending_apdrs(conn) -> None: + """Collapse pre-existing duplicate pending APDR rows down to the latest one per key.""" + stale_ids = [ + row[0] + for row in conn.execute( + sa.text( + f"SELECT id FROM {_TABLE} WHERE created_dag_run_id IS NULL AND id NOT IN (" + f" SELECT MAX(id) FROM {_TABLE} " + " WHERE created_dag_run_id IS NULL " + " GROUP BY target_dag_id, partition_key" + ")" + ) + ).fetchall() + ] + if not stale_ids: + return + id_list = ", ".join(str(i) for i in stale_ids) + conn.execute(sa.text(f"DELETE FROM {_LOG_TABLE} WHERE asset_partition_dag_run_id IN ({id_list})")) + conn.execute(sa.text(f"DELETE FROM {_TABLE} WHERE id IN ({id_list})")) + + +def upgrade(): + """Add pending_partition_key to asset_partition_dag_run and enforce single pending row per key.""" + with disable_sqlite_fkeys(op): + with op.batch_alter_table(_TABLE, schema=None) as batch_op: + batch_op.add_column(sa.Column("pending_partition_key", StringID(), nullable=True)) + + conn = op.get_bind() + # Duplicate resolution requires reading actual data, which offline (SQL-script) + # mode has no connection for; the constraint below still applies to whatever data + # is present when the generated script is eventually run against a live database. + if not context.is_offline_mode(): + _drop_stale_duplicate_pending_apdrs(conn) + + conn.execute( + sa.text( + f"UPDATE {_TABLE} SET pending_partition_key = partition_key WHERE created_dag_run_id IS NULL" + ) + ) + + with op.batch_alter_table(_TABLE, schema=None) as batch_op: + batch_op.create_unique_constraint(_UQ_NAME, ["target_dag_id", "pending_partition_key"]) + + +def downgrade(): + """Drop the pending-partition uniqueness guard and pending_partition_key column.""" + with disable_sqlite_fkeys(op): + with op.batch_alter_table(_TABLE, schema=None) as batch_op: + batch_op.drop_constraint(_UQ_NAME, type_="unique") + batch_op.drop_column("pending_partition_key") diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 7aafddb3b9490..1731e1edd1d5f 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -32,11 +32,12 @@ PrimaryKeyConstraint, String, Table, + UniqueConstraint, delete, select, ) from sqlalchemy.ext.associationproxy import association_proxy -from sqlalchemy.orm import Mapped, mapped_column, relationship +from sqlalchemy.orm import Mapped, mapped_column, relationship, validates from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID @@ -920,15 +921,27 @@ class AssetPartitionDagRun(Base): Where created_dag_run_id is null, the dag run has not yet been created. We should not allow more than one row with the same target_dag_id / - partition_key where created_dag_run_id is null, and this is what the - `_lock_asset_model` mutex control is for. In case a duplicate somehow - gets created, we always work on the latest matching APDR record. + partition_key where created_dag_run_id is null. This is enforced by a + unique constraint on (target_dag_id, pending_partition_key) rather than a + partial/filtered unique index, because MySQL supports neither: + pending_partition_key mirrors partition_key only while created_dag_run_id + is null (see the `_sync_pending_partition_key` validator below), and is + null once the dag run is created. A unique index treats null as distinct + from every other value on all three supported backends, so completed rows + never collide with each other while pending rows for the same key do. In + the rare case a duplicate pending row still exists (e.g. rows created + before this constraint was added), we always work on the latest matching + APDR record. """ id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) target_dag_id: Mapped[str] = mapped_column(StringID(), nullable=False) created_dag_run_id: Mapped[int | None] = mapped_column(Integer(), nullable=True) partition_key: Mapped[str] = mapped_column(StringID(), nullable=False) + # Mirrors partition_key while created_dag_run_id is null, and is cleared to null once + # the dag run is created; see the class docstring for why this backs the uniqueness + # guard instead of a partial index. + pending_partition_key: Mapped[str | None] = mapped_column(StringID(), nullable=True) partition_date: Mapped[datetime | None] = mapped_column(UtcDateTime, nullable=True) # Serialized snapshot of the rollup definition (mapper + window for every # partitioned asset in the timetable) at the time this APDR was created. @@ -953,8 +966,26 @@ class AssetPartitionDagRun(Base): name="apdr_created_dag_run_id_fkey", ondelete="CASCADE", ), + UniqueConstraint( + "target_dag_id", + "pending_partition_key", + name="apdr_target_dag_id_pending_partition_key_uq", + ), ) + @validates("created_dag_run_id") + def _sync_pending_partition_key(self, key, value): + """ + Keep pending_partition_key in lockstep with created_dag_run_id. + + Runs whenever created_dag_run_id is assigned (including at construction), so + every call site — not just _get_or_create_apdr — automatically clears the + pending marker once a dag run is created, without needing to remember to. + """ + if value is not None: + self.pending_partition_key = None + return value + class PartitionedAssetKeyLog(Base): """ diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 615ecb66c6799..759b6b66af917 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "7a98f1b7dbd3", + "3.4.0": "ee86eed19e24", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/assets/test_manager.py b/airflow-core/tests/unit/assets/test_manager.py index 030a0d96c8b0c..0d8280736c127 100644 --- a/airflow-core/tests/unit/assets/test_manager.py +++ b/airflow-core/tests/unit/assets/test_manager.py @@ -27,6 +27,7 @@ import pytest from sqlalchemy import delete, func, select from sqlalchemy.dialects import mysql +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from airflow import settings @@ -361,7 +362,6 @@ def _get_or_create_apdr(): target_partition_date=None, target_dag=testing_dag, rollup_fingerprint=rollup_fingerprint, - asset_id=asm.id, session=_session, ).id finally: @@ -373,10 +373,21 @@ def _get_or_create_apdr(): with concurrent.futures.ThreadPoolExecutor(max_workers=thread_count) as pool: ids = pool.map(lambda _: _get_or_create_apdr(), [None] * thread_count) - assert Counter(r.msg for r in caplog.records) == { - "Existing APDR found for key test_partition_key dag_id testing_dag": thread_count - 1, - "No existing APDR found. Create APDR for key test_partition_key dag_id testing_dag": 1, - } + # The find-or-create is lock-free: a thread may either see the row via its initial + # SELECT ("Existing APDR found"), or race straight to the INSERT and lose to the + # unique constraint on (target_dag_id, pending_partition_key), catching the + # IntegrityError and re-selecting the winner ("Lost race creating APDR"). Which of + # the two happens per thread is a timing detail; the invariant that matters is that + # exactly one thread ever wins the INSERT and every thread converges on that row. + msg_counts = Counter(r.msg for r in caplog.records) + create_msg = "No existing APDR found. Create APDR for key test_partition_key dag_id testing_dag" + found_msg = "Existing APDR found for key test_partition_key dag_id testing_dag" + lost_race_msg = ( + "Lost race creating APDR for key test_partition_key dag_id testing_dag; " + "using the winning APDR instead" + ) + assert msg_counts[create_msg] == 1 + assert msg_counts[found_msg] + msg_counts[lost_race_msg] == thread_count - 1 assert len(set(ids)) == 1 assert session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) == 1 @@ -399,7 +410,6 @@ def test_get_or_create_apdr_suppresses_conflicting_partition_date(self, session) target_partition_date=timezone.parse("2026-05-20T00:00:00"), target_dag=testing_dag, rollup_fingerprint=fp, - asset_id=asm.id, session=session, ) assert first.partition_date == timezone.parse("2026-05-20T00:00:00") @@ -410,7 +420,6 @@ def test_get_or_create_apdr_suppresses_conflicting_partition_date(self, session) target_partition_date=timezone.parse("2026-05-21T00:00:00"), target_dag=testing_dag, rollup_fingerprint=fp, - asset_id=asm.id, session=session, ) assert second.id == first.id # same pending APDR @@ -430,7 +439,6 @@ def test_get_or_create_apdr_keeps_agreeing_partition_date(self, session): target_key="2026-05-20", target_dag=testing_dag, rollup_fingerprint=fp, - asset_id=asm.id, session=session, ) first = AssetManager._get_or_create_apdr(target_partition_date=source_date, **kwargs) @@ -458,7 +466,6 @@ def test_get_or_create_apdr_adopts_date_when_existing_is_none(self, session): target_key="2026-05-20", target_dag=testing_dag, rollup_fingerprint=fp, - asset_id=asm.id, session=session, ) # First event carries no date (e.g. producer had no partition_date). @@ -484,7 +491,6 @@ def test_get_or_create_apdr_recovers_after_conflict(self, session): target_key="2026-05-20", target_dag=testing_dag, rollup_fingerprint=fp, - asset_id=asm.id, session=session, ) first = AssetManager._get_or_create_apdr(target_partition_date=date_1, **kwargs) @@ -539,6 +545,99 @@ def test_carry_partition_date_failure_degrades_to_none(self, session, dag_maker, assert apdr.partition_date is None mock_log.exception.assert_called_once() + @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle") + def test_queue_partitioned_dags_dedups_across_different_producer_assets( + self, session, dag_maker, mock_task_instance + ): + """Regression test for apache/airflow#71070. + + Two DIFFERENT producer assets that both feed a partitioned consumer Dag and + resolve (via IdentityMapper) to the SAME downstream partition key must + collapse into a single AssetPartitionDagRun. The find-or-create step used to + lock the *producer* AssetModel row instead of the target Dag; concurrent + events from two different producer assets therefore took two different row + locks, could each observe "no existing APDR", and inserted a duplicate row + for the same (target_dag_id, partition_key) — leaving neither APDR ever + satisfied by all required events. + + The fix is a unique constraint on (target_dag_id, pending_partition_key) rather + than a lock on any row (see ``AssetPartitionDagRun.pending_partition_key``), so + this asserts the outcome directly: both events, from different producer assets, + collapse into a single row without either insert raising. + """ + _clear_partition_db() + + asset_1 = Asset(name="asset-71070-a") + asset_2 = Asset(name="asset-71070-b") + dag_id = "asset-event-consumer-71070" + with dag_maker( + dag_id=dag_id, + schedule=PartitionedAssetTimetable( + assets=(asset_1 & asset_2), + default_partition_mapper=IdentityMapper(), + ), + serialized=True, + ): + EmptyOperator(task_id="hi") + dag_maker.create_dagrun() + dag_maker.sync_dagbag_to_db() + + AssetManager.register_asset_change( + task_instance=mock_task_instance, + asset=asset_1, + session=session, + partition_key="shared-key", + ) + session.flush() + AssetManager.register_asset_change( + task_instance=mock_task_instance, + asset=asset_2, + session=session, + partition_key="shared-key", + ) + session.flush() + + assert session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) == 1 + + @pytest.mark.usefixtures("clear_assets", "testing_dag_bundle") + def test_pending_partition_key_unique_constraint_blocks_duplicate_pending_rows(self, session): + """DB-level guard backing _get_or_create_apdr's IntegrityError-and-reselect path. + + A second pending row for the same (target_dag_id, partition_key), inserted directly + without going through _get_or_create_apdr, must be rejected by the database itself — + this is what makes the concurrent-insert race in _get_or_create_apdr resolve to one row. + """ + testing_dag = DagModel(dag_id="testing_dag_uq", is_stale=False, bundle_name="testing") + session.add(testing_dag) + session.commit() + + first = AssetPartitionDagRun( + target_dag_id="testing_dag_uq", partition_key="k", pending_partition_key="k" + ) + session.add(first) + session.commit() + + second = AssetPartitionDagRun( + target_dag_id="testing_dag_uq", partition_key="k", pending_partition_key="k" + ) + session.add(second) + with pytest.raises(IntegrityError): + session.commit() + session.rollback() + + def test_created_dag_run_id_assignment_clears_pending_partition_key(self): + """Setting created_dag_run_id must clear pending_partition_key, not just at creation. + + This is what lets a new pending APDR be created for the same (target_dag_id, + partition_key) once the previous one is fulfilled, without ever needing every call + site that marks an APDR as fulfilled to remember to also clear the marker. + """ + apdr = AssetPartitionDagRun(target_dag_id="d", partition_key="k", pending_partition_key="k") + assert apdr.pending_partition_key == "k" + + apdr.created_dag_run_id = 123 + assert apdr.pending_partition_key is None + @pytest.mark.need_serialized_dag @pytest.mark.usefixtures("testing_dag_bundle") def test_queue_partitioned_dags_stamps_rollup_fingerprint(self, session, dag_maker):