[anmolxlight]

closes: #71070

Two producer assets that emit events for the same partition at the same moment can each create a pending AssetPartitionDagRun. The events end up split across two rows, so a consumer scheduled on asset_a & asset_b waits forever. The old find-or-create locked the producer's AssetModel row, but two different producers lock two different rows.

This adds a unique constraint on (target_dag_id, pending_partition_key). When two writers race, the second INSERT fails inside a savepoint and that writer reuses the row the first one created. If the scheduler already consumed that row before the loser re-selects, its key is released and the lookup finds nothing, so the loser INSERTs again (two attempts in total) instead of failing the task. #58919 already ruled out locking the Dag row, and an advisory lock would serialize unrelated Dags too. I added a test that drives two sessions through register_asset_change the way the race happens in production. It leaves two pending rows on main and one row with both events here, on Postgres and MySQL.

pending_partition_key exists because MySQL has no partial unique index. It holds the partition key while the row is pending, and a @validates hook clears it when created_dag_run_id is set. Unique indexes treat NULLs as distinct on all three backends, so consumed rows never collide.

Since #68517 made batching opt-in, a Dag without batch_asset_events is supposed to get one APDR per event. Those rows are written with a NULL pending key and the constraint ignores them. The reuse lookup still matches any pending row for the key, so a Dag that switches batching on keeps adding to its latest pending row.

Migration 0148 (3.4.0) adds the column and the constraint, and sets pending_partition_key on the latest pending row of each (Dag, key). Older duplicate pending rows are left alone: a Dag without batch_asset_events legitimately has one row per event, and merging them would collapse several Dag runs into one. Their NULL key never collides with the constraint, and offline mode runs the same SQL. The 3.5.0 migration becomes 0149, the same reshuffle #74837 did. A development database already at the old 3.5.0 head (c71065a1b14f) skips 0148, so downgrade it past 0149 and upgrade again.

Known issues

An INSERT that hits another transaction's uncommitted row waits for it to finish. If two producer transactions each create new pending rows for the same two (Dag, key) pairs in opposite orders, they deadlock. Sorting target Dags by dag_id rules that out for the Dags of a single event. It can still happen across a task's outlets, or when two mappers emit the same keys in a different order. ti_update_state marks the task FAILED on an unexpected exception today, deadlocks included. I'll send a follow-up so a deadlock there returns a 5xx and the worker retries.

This needs to go to v3-4-test before 3.4.0 final, since the migration is stamped 3.4.0.

[Lee-W]

According to previous discussion #58919 (comment), I don't think we would like to have Dag table locked

[anmolxlight]

Pushed a rework that drops the Dag-row lock entirely — no _lock_target_dag, no SQLite writer-lock retry loop, no Postgres/MySQL SELECT ... FOR UPDATE on DagModel.

The new approach: a unique constraint on (target_dag_id, pending_partition_key), where pending_partition_key mirrors partition_key only while created_dag_run_id IS NULL (cleared via a validates hook the moment a dag run is created). This stays in the design space outlined in #58919 (comment) — no Dag table lock, and no table-wide/advisory lock either. The reason it's a mirrored column instead of the more direct partial/filtered unique index (WHERE created_dag_run_id IS NULL) is portability: MySQL supports neither partial nor filtered indexes, so the mirrored-null-column trick is what's actually usable across sqlite/Postgres/MySQL — a plain unique index already treats NULL as distinct from every other value on all three.

Concurrency is now optimistic instead of pessimistic: _get_or_create_apdr selects the latest pending APDR, and if none exists, inserts inside a SAVEPOINT (session.begin_nested()). If two threads race, the loser's INSERT hits the unique constraint and raises IntegrityError, which is caught to re-select and work on the winner's row instead of raising — consistent with the model's existing "always work on the latest matching APDR record" contract. Migration 0128_3_4_0_add_pending_partition_key_to_apdr.py adds the column, backfills it, collapses any pre-existing duplicate pending rows down to the latest one per key (mirroring the scheduler's own stale-APDR cleanup), then creates the constraint.

Test evidence: test_get_or_create_apdr_is_idempotent_under_concurrent_calls now asserts convergence on exactly one row without assuming which of the two code paths (initial SELECT hit vs. losing the INSERT race) each thread takes, since that's a timing detail, not the invariant. test_register_asset_change_from_different_producer_assets_does_not_duplicate_apdr (the #71070 regression test) now exercises the real path with no lock spy. Added test_pending_partition_key_unique_constraint_blocks_duplicate_pending_rows (DB rejects a second pending row directly) and test_created_dag_run_id_assignment_clears_pending_partition_key (validator coverage). Full airflow-core/tests/unit/assets/test_manager.py suite: 56 passed.

---

Drafted-by: Claude Code (Sonnet 5)

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

[nathadfield]

Hello!

We are encountering this issue and have been able to reproduce it on current main with PostgreSQL 14 and a single scheduler.

The reproduction has two parallel producer tasks emitting different assets for the same partition, and a consumer using PartitionedAssetTimetable with AssetAll(A, B). Both tasks succeeded about 0.3 ms apart, but their events ended up on separate pending APDRs, each with one asset. No consumer run was created. The source and mapped partition keys match exactly, so this isn’t a partition-key mismatch or the multiple-scheduler duplicate-run issue.

We also reproduced it deterministically using two real transactions calling register_asset_change() , synchronized so both APDR lookups finish before either inserts. The trace confirms they lock different asset rows. Sequential A/B registration and concurrent same-asset registration both converge on one APDR.

Is there still a plan to refresh this PR? Happy to share the reproduction and help validate the updated approach rather than open a duplicate.

[Lee-W]

I think we'll need rebasing before the next review round. I can take a look next week

[nathadfield]

A quick update: we tested the current PR head in Breeze with PostgreSQL 14, and it fixes the race we were seeing.

We used two real transactions emitting different assets for the same partition, forcing both to look for an APDR before either inserts one. The same test fails on unfixed code with two pending APDRs. With this PR, all five rounds produced one APDR containing both events, and the scheduler created exactly one consumer run for each pair. Sending another pair for the same partition after consumption also worked.

The existing manager and migration-helper tests passed alongside our new scenarios. We haven’t tested MySQL concurrency or a full upgrade from an existing database yet, so this isn’t a complete sign-off, but it’s encouraging evidence for the core fix.

Let me know if I can assist anymore with this.

[Lee-W]

create a new PR and do a similiar fix is something we can do since the author might not be intersted in coming back to this.

[nathadfield]

Ok, I'll get something up soon-ish.

[kaxil]

@nathadfield Thanks for the reproduction and for testing the earlier head. I've rebased this onto main and pushed to this branch, so no need to open a new PR.

Main moved under it: #68517 made asset event batching opt-in, so the unique constraint now only covers Dags that batch events (an AssetAll condition always does). I also added a two-session race test along the lines of your reproduction; it fails on main with two pending APDRs and passes on Postgres and MySQL with the fix.

Could you re-run your Postgres reproduction against the new head (4a121a9)? If you can also try a MySQL run or an upgrade from an existing 3.3 database with stuck duplicate rows, that would cover the two gaps you mentioned.

[kaxil]

@Lee-W Could you take this PR over from here? It's rebased and green locally, but it still needs a review round .

Could you also pick up the follow-up from the Known issues section? A deadlock in ti_update_state should return a 5xx so the worker retries, instead of marking the task FAILED.

[Lee-W]

Took this over and did another of review. To get this issue merged faster, I make the ti_update_state deadlock follow-up as a separate PR. Will post it here once finished. Gonna merge this one and backport it

[github-actions]

Backport failed to create: v3-4-test. View the failure log Run details

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
❌	v3-4-test

You can attempt to backport this manually by running:

cherry_picker bf022cc v3-4-test

This should apply the commit to the v3-4-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.

After you have resolved the conflicts, you can continue the backport process by running:

cherry_picker --continue

If you don't have cherry-picker installed, see the installation guide.