This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 45bad9f54dd Use async DB session for Execution API task-instance
heartbeat (#67800)
45bad9f54dd is described below
commit 45bad9f54ddb186ba75f7062a76c0770cc125f29
Author: Dev-iL <[email protected]>
AuthorDate: Sat Sep 19 23:29:47 2026 +0300
Use async DB session for Execution API task-instance heartbeat (#67800)
---
.../src/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 8abcaf59fa8..5807cabcec3 100644
--- a/airflow-core/src/airflow/api_fastapi/common/db/common.py
+++ b/airflow-core/src/airflow/api_fastapi/common/db/common.py
@@ -94,7 +94,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 a49d8fe30e5..14910a4d866 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 import trace
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import
TraceContextTextMapPropagator
from pydantic import JsonValue, ValidationError
-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 contains_eager, joinedload
@@ -44,7 +44,7 @@ from airflow._shared.observability.traces import override_ids
from airflow._shared.state import TaskScope
from airflow._shared.timezones import timezone
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.db.dags import eager_load_teams
from airflow.api_fastapi.common.types import UtcDateTime
from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
@@ -891,11 +891,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,
@@ -932,10 +930,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))
@@ -945,7 +943,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,
@@ -966,7 +964,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
)
@@ -974,7 +972,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(
@@ -1006,7 +1007,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)
@@ -1044,7 +1047,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 b6774cc8657..3cf451d55fc 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -736,9 +736,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 ec2724abb5b..08d4fbbfe72 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 ccaecce822d..80278c28d07 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
@@ -36,6 +36,7 @@ from opentelemetry.trace.propagation.tracecontext import
TraceContextTextMapProp
from pydantic import ValidationError
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
@@ -2862,6 +2863,19 @@ class TestTIHealthEndpoint:
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"),
[
@@ -3120,10 +3134,10 @@ class TestTIHealthEndpoint:
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
@@ -3132,9 +3146,9 @@ class TestTIHealthEndpoint:
):
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",
@@ -3166,10 +3180,10 @@ class TestTIHealthEndpoint:
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
@@ -3178,9 +3192,9 @@ class TestTIHealthEndpoint:
):
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",
@@ -3192,6 +3206,47 @@ class TestTIHealthEndpoint:
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 c703ee05ab9..be9e34f611a 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 import mock
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 @@ class TestMetadataEngineHooks:
assert engine is not None
[email protected](
+ ("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
+
+
[email protected](
+ "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:")
