Fix duplicate pending Dag runs for the same asset partition key

[anmolxlight]

closes #71070

Summary

When two different producer assets map to the same downstream partition key (e.g. via IdentityMapper), each asset event could create its own AssetPartitionDagRun row. The scheduler then held two Pending Dag runs for the same partition key and never triggered the consumer Dag.

Root cause

AssetManager._get_or_create_apdr serialized APDR find-or-create with _lock_asset_model, which locks the producer AssetModel row. But APDR dedup is keyed on (target_dag_id, partition_key). Two events from different producer assets that resolve to the same target key therefore took locks on two different asset rows, neither blocked the other, and both observed "no existing APDR" and inserted a duplicate. The reporter's two rows created ~1ms apart with one carrying partition_date and the other empty match this exactly: the two events came from different producers, and only one carried a date.

Fix

Lock the target DagModel row instead (_lock_target_dag). All APDR find-or-create calls for a given consumer Dag now serialize on the same resource regardless of which producer asset triggered the event, so the second event finds the APDR created by the first and dedups onto it. SQLite still uses the global writer lock (unchanged semantics), and the existing "work on the latest matching APDR" fallback in the UI route and scheduler cleanup is preserved.

Test

Added test_queue_partitioned_dags_dedups_across_different_producer_assets: two register_asset_change calls from different producer assets with the same partition key must both lock the same target Dag and yield exactly one AssetPartitionDagRun.

Verified locally with breeze: 6 passed (including the new test and all existing _get_or_create_apdr partition-date tests), ruff clean, formatting clean.

[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)