Fix backfill completion race

[PrithviBadiga]

Fixes a backfill completion race where a backfill could be marked complete before all associated work had reached a terminal state.

Updates:

• make _mark_backfills_complete check unresolved BackfillDagRun associations rather than only unfinished DagRun.backfill_id rows
• preserve dag_run_id on BackfillDagRun rows when a matching DagRun already exists or is in-flight
• add a regression test covering an in-flight association linked to a queued DagRun

This keeps backfills open until all tracked backfill work has actually resolved.

This closes: #68721

[PrithviBadiga]

@ashb @XD-DENG @potiuk Can you check it out ?

[potiuk]

Thanks — this is a genuine bug and a subtle one: an IN_FLIGHT association with no dag_run_id yet points at work that hasn't finished, so the old ~exists over DagRun alone would happily complete a backfill out from under it. Driving the check off BackfillDagRun instead is the right pivot.

Two things I verified while reading, so you don't need to re-check them: State.unfinished_dr_states is exactly {QUEUED, RUNNING}, so swapping out the local tuple is behaviour-preserving; and _mark_backfills_complete runs on a 30-second call_regular_interval rather than every scheduler pass, so this is a periodic query rather than a hot-loop one. Dropping the redundant dag_run_id=None kwargs is fine too — that's the column default.

Also worth noting this is not a duplicate of #68125 despite the similar title — that one is about the creation race, this one about completion. Different bug, different code path.

Notes inline. And the same two asks I've been making today: this wants a second maintainer's review given it rewrites a scheduler predicate, and please do your own manual verification rather than relying on my read — my review was AI-assisted, and multi-level correlated subqueries are precisely where that goes wrong confidently. Seeing the actual generated SQL and an EXPLAIN from a real backfill would be more convincing than either of our analyses.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

[potiuk]

This is a correlated EXISTS containing a second correlated EXISTS, with the inner one referencing BackfillDagRun.logical_date / .partition_key from two levels out. SQLAlchemy's auto-correlation usually gets this right, but two-level correlation is a classic place for it to silently correlate against the wrong FROM and produce a subtly different predicate.

Could you paste the compiled SQL (print(stmt.compile(compile_kwargs={"literal_binds": True}))) into the PR description for at least Postgres? It'd let a reviewer confirm the correlation is what you intend without reconstructing it mentally.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

[potiuk]

The outer correlation is fine — BackfillDagRun has a unique constraint on (backfill_id, dag_run_id), so filtering by backfill_id is indexed.

The inner one I'm less sure about: it scans DagRun filtered by dag_id + state + either logical_date or partition_key. Is DagRun.partition_key indexed? If not, this runs every 30 seconds against what is typically the largest table in the deployment. Worth checking the plan on an instance with a large dag_run table before this lands.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

[potiuk]

The test covers the logical_date branch well, including the transition from "stays active" to "completes once the run succeeds". The partition_key branch — logical_date IS NULL with a non-null partition_key — has no coverage, and it's the harder of the two to reason about.

Since the partitioned path is the reason that branch exists, a companion test there would be worth adding.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting