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:")