Use async DB session for Execution API task-instance heartbeat

[Dev-iL]

Depends on:

• Switch the default async Postgres driver from asyncpg to psycopg3 (core) #68496

Related:

• [POC][WIP] Async SQLAlchemy sessions in Airflow #36504
• Migrate API endpoints from sync to async DB sessions #67799
• Return 410 for stale id in set rtif API #68902

What & why

Converts the Execution API ti_heartbeat route from the synchronous SessionDep to the async AsyncSessionDep, adopting the async metadata engine that already ships in Airflow 3.x (create_async_engine, async_sessionmaker, AsyncSessionDep, create_session_async).

Heartbeat is the highest-QPS write path in the system (one call per running task per heartbeat interval, per worker), so it is a meaningful first production route to run on the async engine. The goal of this PR is behavioral parity with zero regression — not a throughput change. It follows the same low-blast-radius conversion pattern already merged for GET /execution/variables/keys.

What changed

• Route (execution_api/routes/task_instances.py): ti_heartbeat is now async def and takes AsyncSessionDep; the fast-path UPDATE, the slow-path SELECT … FOR UPDATE (.one()), the Task-Instance-History existence count, and the final UPDATE are awaited. On the slow-path miss the route awaits its own TI-History count query and then delegates the 410/404 decision to the shared _raise_ti_not_in_live_table helper (see next bullet). No explicit commit() is added — the async dependency commits on success and rolls back (releasing the FOR UPDATE row lock) on exception, mirroring the sync dependency exactly. synchronize_session=False and with_for_update() are unchanged.
• Shared helper (_raise_ti_not_in_live_table): Return 410 for stale id in set rtif API #68902 (merged into main after this branch was opened) refactored the inline 404/410 logic out of both ti_heartbeat and ti_put_rtif into this helper, which ran its own synchronous session.scalar(...). That is incompatible with the async heartbeat session — a scalar(...) on an AsyncSession returns a coroutine, and if <coroutine>: is always truthy, so calling it unchanged would return 410 on every miss. The helper is therefore made session-agnostic: it no longer touches the DB and takes a keyword-only archived_in_history: bool, raising 410 / 404 from that flag. Each caller runs its own count query in its own style — the async route awaits it, the sync ti_put_rtif route does not — so the log/message/response logic stays centralized in one place.
• rtif route (ti_put_rtif): updated to run the TI-History count query itself and pass the result as archived_in_history= to the helper. This is the one change outside the heartbeat route; it is a mechanical adaptation to the new helper signature, forced by rebasing onto Return 410 for stale id in set rtif API #68902, and does not alter the rtif contract.
• Tests (versions/head/test_task_instances.py):
  • The two tests that intercept the fast-path UPDATE now patch sqlalchemy.ext.asyncio.AsyncSession with an async interceptor (the route now awaits AsyncSession.execute); assertions are unchanged.
  • Added a reconfigure_async_db_engine autouse fixture to TestTIHealthEndpoint. The async engine binds its connection pool to the event loop that created it, while the test harness builds a fresh app and event loop per test; without this, a pooled connection from a prior test's closed loop is reused and fails (attached to a different loop). This is the same workaround already used by TestWaitDagRun.
• Docs / config: Switch the default async Postgres driver from asyncpg to psycopg3 #67801 (merged) added a dedicated "Async SQLAlchemy engine and the async driver" section to the Postgres setup guide that already documents psycopg3 as the PgBouncer-safe default and the asyncpg opt-in connect_args recipe. This PR therefore no longer adds a duplicate PgBouncer block to the setup guide (that block was dropped in the rebase); it keeps only a concise note on the sql_alchemy_connect_args_async config reference in config.yml, framed around the asyncpg opt-in and pointing operators at that section.
• Newsfragment (improvement): notes that the heartbeat endpoint now uses the async engine and that the API server may hold both the sync and async connection pools concurrently, so operators should budget DB max_connections for both.

Behavioral parity

No change to the endpoint contract: same 204 success, same 404 / 409 (running-elsewhere / not-running) / 410 (cleared-and-archived) responses with identical detail payloads, same request signature. No Execution API version bump. The 404/410 branch now routes through the shared helper introduced by #68902 rather than an inline block, but the emitted responses are identical; the heartbeat and rtif tests (including the 410 cases) pass and are the parity proof.

Testing

Heartbeat + rtif suite passes on SQLite (aiosqlite) after the rebase onto #68902, including both 410 paths (test_ti_heartbeat_cleared_task_returns_410, test_ti_put_rtif_archived_ti_returns_410)

Operational notes

• Default driver is PgBouncer-safe. Switch the default async Postgres driver from asyncpg to psycopg3 #67801 switched the default async Postgres driver to psycopg3, which is safe under transaction-mode PgBouncer with no extra configuration. For the common deployment nothing is required from operators.
• asyncpg opt-in. asyncpg remains supported as a high-throughput opt-in (sql_alchemy_conn_async = postgresql+asyncpg://…). It uses named server-side prepared statements, which break under transaction-mode PgBouncer, so it needs sql_alchemy_connect_args_async = {"statement_cache_size": 0, "prepared_statement_cache_size": 0}. This PR documents that recipe.

Note to reviewers

• This branch was rebased onto Return 410 for stale id in set rtif API #68902 (stale-id 410 for the rtif API), which introduced the shared _raise_ti_not_in_live_table helper. Resolving that rebase is why this heartbeat PR also touches ti_put_rtif and the helper's signature — see the "Shared helper" / "rtif route" bullets above. No rtif behavior changes.
• A cosmetic Event loop is closed message is logged during session teardown when the async engine is disposed (dispose_orm closing the async engine synchronously). It appears on all three async drivers, is pre-existing engine disposal behavior unrelated to this change, and every run exits 0. Tidying async-engine disposal is out of scope for this PR.

---

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

• Yes (please specify the tool below)

Generated-by: Claude Opus 4.8 following the guidelines

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[ashb]

Nice one.

[ashb]

• Connection budget: heartbeat moving to the async engine means a busy API server can hold both the sync pool (for the remaining sync routes) and the async pool concurrently. Size DB max_connections accordingly (called out in the newsfragment).
• PgBouncer (transaction pooling): asyncpg uses server-side prepared statements, which break under transaction-mode PgBouncer. This is documented as a deployment-configuration requirement here; shipping pgbouncer-safe async-engine defaults is deferred and tracked in #.

I think the agent made these up :)

[Dev-iL]

• Connection budget: heartbeat moving to the async engine means a busy API server can hold both the sync pool (for the remaining sync routes) and the async pool concurrently. Size DB max_connections accordingly (called out in the newsfragment).
• PgBouncer (transaction pooling): asyncpg uses server-side prepared statements, which break under transaction-mode PgBouncer. This is documented as a deployment-configuration requirement here; shipping pgbouncer-safe async-engine defaults is deferred and tracked in #.

I think the agent made these up :)

I failed to edit the description after opening the mentioned tracking issue - here it is: #67801

[Dev-iL]

@ashb Could you please take another look and see if your concerns were addressed?

[ashb]

@Dev-iL I feel very uncomfortable merging this without the swap away from asyncpg first as it would essentially break all task execution for anyone using pgbouncer unless they happen to notice the change in the docs.

Either we need to change the asyncpg migration to disable the prepared statments by default, or do #67801 first.

Given 3.3.0 is about to enter the Beta and RC stage we need to be careful about not merging "half-breaking" changes to main.

[Dev-iL]

Either we need to change the asyncpg migration to disable the prepared statments by default, or do #67801 first.

• Switch the default async Postgres driver from asyncpg to psycopg3 (core) #68496

[Dev-iL]

@ashb How about now?

[ashb]

One tiny nit, LGTM