From 0c6afdb18d49306bda6cc9974321d3d706165210 Mon Sep 17 00:00:00 2001
From: Dev-iL <6509619+Dev-iL@users.noreply.github.com>
Date: Sun, 31 May 2026 09:23:36 +0300
Subject: [PATCH] Use async DB session for Execution API task-instance
 heartbeat

Convert ti_heartbeat from the synchronous SessionDep to the async AsyncSessionDep, adopting the async metadata engine that already ships in Airflow 3.x. The route's behavior is unchanged: same 204/404/409/410 responses, same fast-path UPDATE / SELECT ... FOR UPDATE fallback, same last_heartbeat_at semantics, no version bump.

Make AsyncSessionDep function-scoped so the async session commits before the response is sent, matching the synchronous SessionDep. FastAPI's default "request" scope for yield-dependencies commits after the response is sent, so a failed commit would roll back after the worker already received a 204 success. A regression test forces a commit failure and asserts the route surfaces an error rather than a silent 204.

Heartbeat's async writes exposed a test-harness issue on Postgres: the async engine binds its connection pool to the event loop that created it, while the test harness builds a fresh app and loop per test, so a pooled connection from a prior test's closed loop was reused. Add a reconfigure_async_db_engine autouse fixture to the heartbeat tests that rebuilds the async session per test, mirroring the existing workaround in TestWaitDagRun.

Document the asyncpg + transaction-mode pgbouncer prepared-statement caveat in the sql_alchemy_connect_args_async config reference and the PGBouncer setup guide (set statement_cache_size=0 and prepared_statement_cache_size=0); switching the default async driver to psycopg3 is tracked separately.
---
 .../airflow/api_fastapi/common/db/common.py   |  2 +-
 .../execution_api/routes/task_instances.py    | 32 +++++----
 .../src/airflow/config_templates/config.yml   |  7 +-
 airflow-core/src/airflow/settings.py          | 34 +++++++--
 .../versions/head/test_task_instances.py      | 71 ++++++++++++++++---
 airflow-core/tests/unit/core/test_settings.py | 57 ++++++++++++++-
 6 files changed, 174 insertions(+), 29 deletions(-)

diff --git a/airflow-core/src/airflow/api_fastapi/common/db/common.py b/airflow-core/src/airflow/api_fastapi/common/db/common.py
index ff1e3f7043e6a..a940a7991330e 100644
--- a/airflow-core/src/airflow/api_fastapi/common/db/common.py
+++ b/airflow-core/src/airflow/api_fastapi/common/db/common.py
@@ -64,7 +64,7 @@ async def _get_async_session() -> AsyncGenerator[AsyncSession, None]:
         yield session
 
 
-AsyncSessionDep = Annotated[AsyncSession, Depends(_get_async_session)]
+AsyncSessionDep = Annotated[AsyncSession, Depends(_get_async_session, scope="function")]
 
 
 @overload
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index 34c3dc35406f8..0febf650c5e2e 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -33,7 +33,7 @@
 from opentelemetry.trace import StatusCode
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from pydantic import JsonValue
-from sqlalchemy import and_, func, or_, tuple_, update
+from sqlalchemy import and_, exists, func, or_, tuple_, update
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
 from sqlalchemy.orm import joinedload
@@ -45,7 +45,7 @@
 from airflow._shared.timezones import timezone
 from airflow.api_fastapi.auth.tokens import JWTGenerator
 from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
-from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.common.db.common import AsyncSessionDep, SessionDep
 from airflow.api_fastapi.common.types import UtcDateTime
 from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
@@ -856,11 +856,9 @@ def ti_skip_downstream(
     log.info("Downstream tasks skipped", tasks_skipped=getattr(result, "rowcount", 0))
 
 
-def _raise_ti_not_in_live_table(task_instance_id: UUID, session: SessionDep) -> NoReturn:
+def _raise_ti_not_in_live_table(task_instance_id: UUID, *, archived_in_history: bool) -> NoReturn:
     """Raise 410 Gone if the missing TI id was archived to history, else 404 Not Found."""
-    if session.scalar(
-        select(func.count(TIH.task_instance_id)).where(TIH.task_instance_id == task_instance_id)
-    ):
+    if archived_in_history:
         log.error("TaskInstance not in live table but archived in history", ti_id=str(task_instance_id))
         raise HTTPException(
             status_code=status.HTTP_410_GONE,
@@ -897,10 +895,10 @@ def _raise_ti_not_in_live_table(task_instance_id: UUID, session: SessionDep) ->
         ]
     ),
 )
-def ti_heartbeat(
+async def ti_heartbeat(
     task_instance_id: UUID,
     ti_payload: TIHeartbeatInfo,
-    session: SessionDep,
+    session: AsyncSessionDep,
 ):
     """Update the heartbeat of a TaskInstance to mark it as alive & still running."""
     bind_contextvars(ti_id=str(task_instance_id))
@@ -910,7 +908,7 @@ def ti_heartbeat(
     # so we can update last_heartbeat_at directly without first taking a row lock.
     fast_path_result = cast(
         "CursorResult[Any]",
-        session.execute(
+        await session.execute(
             update(TI)
             .where(
                 TI.id == task_instance_id,
@@ -931,7 +929,7 @@ def ti_heartbeat(
     old = select(TI.state, TI.hostname, TI.pid).where(TI.id == task_instance_id).with_for_update()
 
     try:
-        (previous_state, hostname, pid) = session.execute(old).one()
+        (previous_state, hostname, pid) = (await session.execute(old)).one()
         log.debug(
             "Retrieved current task state", state=previous_state, current_hostname=hostname, current_pid=pid
         )
@@ -939,7 +937,10 @@ def ti_heartbeat(
         # Check if the TI exists in the Task Instance History table.
         # If it does, it was likely cleared while running, so return 410 Gone
         # instead of 404 Not Found to give the client a more specific signal.
-        _raise_ti_not_in_live_table(task_instance_id, session)
+        archived_in_history = bool(
+            await session.scalar(select(exists().where(TIH.task_instance_id == task_instance_id)))
+        )
+        _raise_ti_not_in_live_table(task_instance_id, archived_in_history=archived_in_history)
 
     if hostname != ti_payload.hostname or pid != ti_payload.pid:
         log.warning(
@@ -971,7 +972,9 @@ def ti_heartbeat(
         )
 
     # Update the last heartbeat time!
-    session.execute(update(TI).where(TI.id == task_instance_id).values(last_heartbeat_at=timezone.utcnow()))
+    await session.execute(
+        update(TI).where(TI.id == task_instance_id).values(last_heartbeat_at=timezone.utcnow())
+    )
     log.debug("Heartbeat updated", state=previous_state)
 
 
@@ -1009,7 +1012,10 @@ def ti_put_rtif(
     task_instance = session.scalar(select(TI).where(TI.id == task_instance_id))
     if not task_instance:
         # On retry/clear, the server regenerates the TI id. Return 410 for the stale id.
-        _raise_ti_not_in_live_table(task_instance_id, session)
+        archived_in_history = bool(
+            session.scalar(select(exists().where(TIH.task_instance_id == task_instance_id)))
+        )
+        _raise_ti_not_in_live_table(task_instance_id, archived_in_history=archived_in_history)
     task_instance.update_rtif(put_rtif_payload, session=session)
     log.debug("RenderedTaskInstanceFields updated successfully")
 
diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml
index b3dc3f133d229..2922ac020968c 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -717,9 +717,12 @@ database:
     sql_alchemy_connect_args_async:
       description: |
         Import path for connect args in SQLAlchemy. Defaults to an empty dict.
-        This is similar to ``sql_alchemy_connect_args``, but only for async connections.
+        This is similar to ``sql_alchemy_connect_args``, but only for async connections and must be a dict.
 
-        This configuration is only applied to async engines, such as asyncpg.
+        This is applied to the async engine only. The default async Postgres driver (psycopg3) needs
+        no extra args. If you opt in to the ``asyncpg`` driver behind transaction-mode pgbouncer, set
+        ``statement_cache_size`` and ``prepared_statement_cache_size`` to ``0`` here to avoid "prepared
+        statement does not exist" errors. See the async driver section of the Postgres setup guide.
       version_added: 3.1.0
       type: string
       example: 'airflow_local_settings.connect_args_async'
diff --git a/airflow-core/src/airflow/settings.py b/airflow-core/src/airflow/settings.py
index 0a0e1c6619726..c7e0f4baa11f9 100644
--- a/airflow-core/src/airflow/settings.py
+++ b/airflow-core/src/airflow/settings.py
@@ -240,8 +240,28 @@ def load_policy_plugins(pm: pluggy.PluginManager):
     pm.load_setuptools_entrypoints("airflow.policy")
 
 
+def _translate_asyncpg_sslmode(async_uri: str) -> str:
+    """
+    Rename the libpq ``sslmode`` query param to asyncpg's ``ssl`` equivalent.
+
+    asyncpg has no ``sslmode`` connect arg -- SQLAlchemy's asyncpg dialect passes URL
+    query params straight to ``asyncpg.connect()``, which spells the option ``ssl`` and
+    accepts the same libpq mode strings (``disable``/``prefer``/``require``/``verify-ca``/
+    ``verify-full``). A sync Postgres URI commonly carries ``sslmode`` (the official Helm
+    chart sets it unconditionally), so the derived async URI must translate it; otherwise
+    every async DB call raises ``TypeError: connect() got an unexpected keyword argument
+    'sslmode'``.
+    """
+    url = make_url(async_uri)
+    if "sslmode" not in url.query:
+        return async_uri
+    query = dict(url.query)
+    query["ssl"] = query.pop("sslmode")
+    return url.set(query=query).render_as_string(hide_password=False)
+
+
 def _get_async_conn_uri_from_sync(sync_uri):
-    # Mapping of backend to async driver:
+    # Derive the async URI scheme from the sync one by swapping in the async driver:
     AIO_LIBS_MAPPING = {
         "sqlite": "aiosqlite",
         "postgresql": "psycopg_async" if _USE_PSYCOPG3 else "asyncpg",
@@ -251,9 +271,15 @@ def _get_async_conn_uri_from_sync(sync_uri):
     scheme, rest = sync_uri.split(":", maxsplit=1)
     scheme = scheme.split("+", maxsplit=1)[0]
     aiolib = AIO_LIBS_MAPPING.get(scheme)
-    if aiolib:
-        return f"{scheme}+{aiolib}:{rest}"
-    return sync_uri
+    if not aiolib:
+        return sync_uri
+    async_uri = f"{scheme}+{aiolib}:{rest}"
+    if aiolib == "asyncpg":
+        # asyncpg-only: it is not libpq-based and has no ``sslmode`` connect arg. A
+        # libpq-based async driver (e.g. psycopg3) understands ``sslmode`` natively and
+        # would in turn reject ``ssl``, so this translation must stay gated on asyncpg.
+        async_uri = _translate_asyncpg_sslmode(async_uri)
+    return async_uri
 
 
 def configure_vars():
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index 8a152bebe0d3f..b2f15459234ea 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -34,6 +34,7 @@
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from sqlalchemy import select, update
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Session
 
 from airflow._shared.observability.traces import OverrideableRandomIdGenerator
@@ -2586,6 +2587,19 @@ def setup_method(self):
     def teardown_method(self):
         clear_db_runs()
 
+    # ti_heartbeat runs on the async engine. The async engine binds its pool to
+    # the event loop that created it (once per process), but the test harness
+    # builds a fresh FastAPI app and event loop per test, so a pooled connection
+    # from a prior test's closed loop gets reused and fails ("attached to a
+    # different loop"). Re-configuring the async session before each test rebuilds
+    # the engine on the current loop. Same workaround as TestWaitDagRun in
+    # tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py.
+    @pytest.fixture(autouse=True)
+    def reconfigure_async_db_engine(self):
+        from airflow.settings import _configure_async_session
+
+        _configure_async_session()
+
     @pytest.mark.parametrize(
         ("hostname", "pid", "expected_status_code", "expected_detail"),
         [
@@ -2844,10 +2858,10 @@ def test_ti_heartbeat_fallback_updates_on_fast_path_miss(
         new_time = time_now.add(minutes=10)
         time_machine.move_to(new_time, tick=False)
 
-        original_execute = Session.execute
+        original_execute = AsyncSession.execute
         fast_path_intercepted = False
 
-        def execute_with_fast_path_miss(session_obj, statement, *args, **kwargs):
+        async def execute_with_fast_path_miss(session_obj, statement, *args, **kwargs):
             nonlocal fast_path_intercepted
             if (
                 not fast_path_intercepted
@@ -2856,9 +2870,9 @@ def execute_with_fast_path_miss(session_obj, statement, *args, **kwargs):
             ):
                 fast_path_intercepted = True
                 return mock.MagicMock(rowcount=0)
-            return original_execute(session_obj, statement, *args, **kwargs)
+            return await original_execute(session_obj, statement, *args, **kwargs)
 
-        monkeypatch.setattr(Session, "execute", execute_with_fast_path_miss)
+        monkeypatch.setattr(AsyncSession, "execute", execute_with_fast_path_miss)
 
         response = client.put(
             f"/execution/task-instances/{ti.id}/heartbeat",
@@ -2890,10 +2904,10 @@ def test_ti_heartbeat_fallback_updates_on_unknown_fast_path_rowcount(
         new_time = time_now.add(minutes=10)
         time_machine.move_to(new_time, tick=False)
 
-        original_execute = Session.execute
+        original_execute = AsyncSession.execute
         fast_path_intercepted = False
 
-        def execute_with_unknown_fast_path_rowcount(session_obj, statement, *args, **kwargs):
+        async def execute_with_unknown_fast_path_rowcount(session_obj, statement, *args, **kwargs):
             nonlocal fast_path_intercepted
             if (
                 not fast_path_intercepted
@@ -2902,9 +2916,9 @@ def execute_with_unknown_fast_path_rowcount(session_obj, statement, *args, **kwa
             ):
                 fast_path_intercepted = True
                 return mock.MagicMock(rowcount=-1)
-            return original_execute(session_obj, statement, *args, **kwargs)
+            return await original_execute(session_obj, statement, *args, **kwargs)
 
-        monkeypatch.setattr(Session, "execute", execute_with_unknown_fast_path_rowcount)
+        monkeypatch.setattr(AsyncSession, "execute", execute_with_unknown_fast_path_rowcount)
 
         response = client.put(
             f"/execution/task-instances/{ti.id}/heartbeat",
@@ -2916,6 +2930,47 @@ def execute_with_unknown_fast_path_rowcount(session_obj, statement, *args, **kwa
         session.refresh(ti)
         assert ti.last_heartbeat_at == new_time
 
+    def test_ti_heartbeat_commit_failure_surfaces_error(
+        self, client, session, create_task_instance, monkeypatch
+    ):
+        """A commit failure must reach the worker as an error, never a silent 204.
+
+        ``AsyncSessionDep`` is function-scoped, so the yield-dependency commit runs
+        *before* the response is sent -- mirroring the sync ``SessionDep``. Were it
+        request-scoped (the FastAPI default for ``yield`` dependencies), the 204 would
+        be sent before the commit, so a commit failure (e.g. an asyncpg /
+        transaction-mode PgBouncer drop) would roll back *after* the worker already
+        saw success. Regression guard for the parity goal of this route conversion.
+        """
+        ti = create_task_instance(
+            task_id="test_ti_heartbeat_commit_failure",
+            state=State.RUNNING,
+            hostname="random-hostname",
+            pid=1789,
+            session=session,
+        )
+        session.commit()
+
+        async def failing_commit(self):
+            raise SQLAlchemyError("simulated commit failure (connection dropped)")
+
+        monkeypatch.setattr(AsyncSession, "commit", failing_commit)
+        # The default TestClient re-raises server exceptions, and it does so for *both*
+        # dependency scopes; the worker-visible status (500 vs a silent 204) is what
+        # tells them apart, so observe the response the worker would actually receive.
+        monkeypatch.setattr(client._transport, "raise_server_exceptions", False)
+
+        response = client.put(
+            f"/execution/task-instances/{ti.id}/heartbeat",
+            json={"hostname": "random-hostname", "pid": 1789},
+        )
+
+        # Function scope -> commit fails before the response -> 500. Request scope -> 204.
+        assert response.status_code == 500
+        # The transaction rolled back, so the heartbeat was not persisted.
+        session.refresh(ti)
+        assert ti.last_heartbeat_at is None
+
 
 class TestTIPutRTIF:
     def setup_method(self):
diff --git a/airflow-core/tests/unit/core/test_settings.py b/airflow-core/tests/unit/core/test_settings.py
index c703ee05ab941..be9e34f611abf 100644
--- a/airflow-core/tests/unit/core/test_settings.py
+++ b/airflow-core/tests/unit/core/test_settings.py
@@ -25,7 +25,7 @@
 from unittest.mock import MagicMock, call, patch
 
 import pytest
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import Engine, make_url
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.pool import NullPool
 
@@ -400,6 +400,61 @@ def test_override_via_local_settings(self):
             assert engine is not None
 
 
+@pytest.mark.parametrize(
+    ("sync_uri", "expected_drivername", "expected_query"),
+    [
+        # libpq's ``sslmode`` is not an asyncpg connect arg -- it must be translated to
+        # asyncpg's ``ssl`` (which accepts the same mode strings). Without this the
+        # derived async URI raises "connect() got an unexpected keyword argument 'sslmode'"
+        # on every async DB call. The official Helm chart sets sslmode=disable by default.
+        (
+            "postgresql://airflow:pw@postgres:5432/airflow?sslmode=disable",
+            "postgresql+asyncpg",
+            {"ssl": "disable"},
+        ),
+        (
+            "postgresql://airflow:pw@postgres/airflow?sslmode=verify-full",
+            "postgresql+asyncpg",
+            {"ssl": "verify-full"},
+        ),
+        # An explicit sync driver in the scheme is dropped before swapping to asyncpg.
+        (
+            "postgresql+psycopg2://u:p@h/db?sslmode=require",
+            "postgresql+asyncpg",
+            {"ssl": "require"},
+        ),
+        # No sslmode -> the query is left untouched.
+        ("postgresql://u:p@h/db", "postgresql+asyncpg", {}),
+    ],
+)
+def test_get_async_conn_uri_translates_pg_sslmode(sync_uri, expected_drivername, expected_query):
+    from airflow import settings
+    from airflow.settings import _get_async_conn_uri_from_sync
+
+    with patch.object(settings, "_USE_PSYCOPG3", False):
+        url = make_url(_get_async_conn_uri_from_sync(sync_uri))
+    assert url.drivername == expected_drivername
+    assert "sslmode" not in url.query
+    assert dict(url.query) == expected_query
+
+
+@pytest.mark.parametrize(
+    "sync_uri",
+    [
+        # Non-postgres backends use a different driver and SSL params; the sslmode
+        # translation must not touch their query strings.
+        "mysql://u:p@h/db?charset=utf8mb4",
+        "sqlite:////tmp/airflow.db",
+    ],
+)
+def test_get_async_conn_uri_leaves_non_postgres_query_untouched(sync_uri):
+    from airflow.settings import _get_async_conn_uri_from_sync
+
+    src_query = dict(make_url(sync_uri).query)
+    out_query = dict(make_url(_get_async_conn_uri_from_sync(sync_uri)).query)
+    assert out_query == src_query
+
+
 _local_db_path_error = pytest.raises(AirflowConfigException, match=r"Cannot use relative path:")
 
 
