anmolxlight commented on code in PR #71074:
URL: https://github.com/apache/airflow/pull/71074#discussion_r3734689141
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -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()
Review Comment:
Verified this empirically rather than just from the code. The scenario does
not reproduce: `Session.begin_nested()` itself flushes all pending ORM state
into the parent transaction *before* the SAVEPOINT is opened. In SQLAlchemy
2.0.51, `SessionTransaction.__init__` calls `_take_snapshot()`, and for a
`BEGIN_NESTED` origin `is_begin` is False, so it runs `self.session.flush()`
before resetting `self._new` to an empty snapshot. By the time `with
session.begin_nested():` in `_get_or_create_apdr` is entered, any unflushed
`Log`/`PartitionedAssetKeyLog` rows accumulated earlier in
`_queue_partitioned_dags`'s loop are already persistent, not in `session._new`.
So `_restore_snapshot()`'s `to_expunge = set(self._new) | set(session._new)` on
rollback only ever contains the savepoint-local `apdr` object, which the
`except IntegrityError` block already re-selects around.
I confirmed this with a standalone repro (unflushed rows added before
`begin_nested()`, forced `IntegrityError` inside the savepoint, rollback, then
flush) and all pre-savepoint rows persisted. I also added a comment explaining
this in `_get_or_create_apdr` (manager.py) and a regression test,
`test_lost_race_does_not_expunge_unflushed_log_rows` in `test_manager.py`, that
forces a deterministic lost race and asserts an unflushed log row added earlier
survives and persists. See 09fcebfadc3.
##########
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})"))
Review Comment:
Agreed, and confirmed there is no constraint blocking the UPDATE.
`_drop_stale_duplicate_pending_apdrs` (now
`_heal_stale_duplicate_pending_apdrs`) computes the loser-to-winner id mapping
per (target_dag_id, partition_key) group, runs `UPDATE
partitioned_asset_key_log SET asset_partition_dag_run_id = <winner_id> WHERE
asset_partition_dag_run_id = <loser_id>` for each loser before dropping the
loser APDR rows. The module docstring is updated to describe the re-pointing
instead of the drop; its first line (the one docs/migrations-ref.rst pulls
from) is unchanged. Added a regression test,
test_0130_pending_partition_key_migration.py, seeding a duplicate pending pair
with log rows and asserting the loser's log rows survive re-pointed onto the
winner and the loser APDR is gone. See 5c9a229d7bb.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
