Pass the scheduler's session to set_state when processing executor events

[namanjain24-sudo]

Opened fresh from #73213, which was closed under the new open-PR-limit policy and then failed to reopen (both gh pr reopen and the UI button errored with a generic "Could not open the pull request") — following @potiuk's guidance on Slack to push/open fresh from the branch when reopening doesn't work.

SchedulerJobRunner.process_executor_events calls ti.set_state() without session in two places: when a cleared (RESTARTING) task instance is reported as successfully terminated, and when the Dag of a finished task instance cannot be loaded. TaskInstance.set_state is @provide_session and settings.Session is scoped, so create_session() returns the scheduler's own session and commits and closes it on exit. This is the same mechanism as #67850 and #71968.

process_executor_events runs inside with create_session() in _run_scheduler_loop, not under prohibit_commit, so nothing raises. Instead, in the middle of an executor-event batch:

• the scheduler's transaction is committed early, which releases the FOR UPDATE SKIP LOCKED row locks taken on the batch's task instances;
• close() detaches the task instances loaded for the batch, so changes made to them afterwards are never written. In a local run with one RESTARTING → SUCCESS event and four QUEUED events carrying an external executor id, none of the four external_executor_id values reached the database without this change, and all four did with it (SQLite and PostgreSQL 16).

This passes the scheduler's session at both call sites, as _enqueue_task_instances_with_queued_state already does for its own ti.set_state() call.

Tests:

• test_process_executor_events_sets_state_in_callers_transaction covers both paths: the state change has to roll back with the caller's transaction. On main both cases fail (the task instance is already committed as None / failed); with this change they pass.
• The rest of test_scheduler_job.py and prek, including mypy for airflow-core, pass locally on SQLite and PostgreSQL 16.

Not covered here. In the same loop, executor.send_callback() reaches DatabaseCallbackSink.send, which is also @provide_session and gets no session. So when the task instance has an on_failure_callback / on_retry_callback (reproduced locally with this change applied) or email configured, the scheduler's transaction is still committed at that call. _maybe_requeue_stuck_ti and _purge_task_instances_without_heartbeats call send_callback() the same way. Fixing that means either adding a session argument to BaseExecutor.send_callback and BaseCallbackSink.send, which executors that override send_callback would then have to accept, or handing the session to the sink directly, as Trigger already does with DatabaseCallbackSink().send(callback=request, session=session). I kept this PR to the set_state calls and am happy to follow up with whichever approach you prefer.

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Generated-by: a Gen-AI coding assistant, following the guidelines. I reviewed the change and ran the checks above.

[luc-pimentel]

After #73916, this ti1.key event is dropped before it reaches set_state, so the test passes even without the fix...

Should we rebase and use TaskInstanceUuid(ti1.id) like the other tests do now? With that change and adding an assertion that the state is FAILED before the rollback, the test fails again without the fix (which is the correct behavior).

[namanjain24-sudo]

Good catch — addressed in cf2bac6: the test now keys the executor event by TaskInstanceUuid(ti1.id) (matching the other tests in this file) and asserts ti1.state == TaskInstanceState.FAILED right after _process_executor_events and before the session.rollback(), so it fails again without the session fix.

---

Drafted-by: Claude Code (Sonnet 5); reviewed by @namanjain24-sudo before posting