This is an automated email from the ASF dual-hosted git repository.
ashb 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 489316fbc72 Migrate execution API reads to async sessions (#73403)
489316fbc72 is described below
commit 489316fbc72cc0c3ca99f7656cf4d1c9eca3b558
Author: Subhramit Basu <[email protected]>
AuthorDate: Thu Oct 1 02:42:31 2026 +0530
Migrate execution API reads to async sessions (#73403)
Migrate a small group of Execution API read endpoints from sync DB access
to AsyncSessionDep.
This updates direct ORM read handlers that can safely switch to async
without pulling in sync-only helpers.
---
.../api_fastapi/execution_api/routes/assets.py | 14 +++++------
.../api_fastapi/execution_api/routes/dag_runs.py | 18 +++++++-------
.../api_fastapi/execution_api/routes/hitl.py | 12 ++++++----
.../execution_api/routes/task_instances.py | 28 ++++++++++++----------
.../execution_api/routes/task_reschedules.py | 6 ++---
5 files changed, 43 insertions(+), 35 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/assets.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/assets.py
index 385ec509f9f..d6f03691eda 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/assets.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/assets.py
@@ -22,7 +22,7 @@ from typing import Annotated
from fastapi import APIRouter, HTTPException, Query, status
from sqlalchemy import select
-from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.common.db.common import AsyncSessionDep, SessionDep
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.models.asset import AssetModel, expand_alias_to_assets
@@ -35,24 +35,24 @@ router = APIRouter(
@router.get("/by-name")
-def get_asset_by_name(
+async def get_asset_by_name(
name: Annotated[str, Query(description="The name of the Asset")],
- session: SessionDep,
+ session: AsyncSessionDep,
) -> AssetResponse:
"""Get an Airflow Asset by `name`."""
- asset = session.scalar(select(AssetModel).where(AssetModel.name == name,
AssetModel.active.has()))
+ asset = await session.scalar(select(AssetModel).where(AssetModel.name ==
name, AssetModel.active.has()))
_raise_if_not_found(asset, f"Asset with name {name} not found")
return AssetResponse.model_validate(asset)
@router.get("/by-uri")
-def get_asset_by_uri(
+async def get_asset_by_uri(
uri: Annotated[str, Query(description="The URI of the Asset")],
- session: SessionDep,
+ session: AsyncSessionDep,
) -> AssetResponse:
"""Get an Airflow Asset by `uri`."""
- asset = session.scalar(select(AssetModel).where(AssetModel.uri == uri,
AssetModel.active.has()))
+ asset = await session.scalar(select(AssetModel).where(AssetModel.uri ==
uri, AssetModel.active.has()))
_raise_if_not_found(asset, f"Asset with URI {uri} not found")
return AssetResponse.model_validate(asset)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
index 00a9ab37f25..14bebddbd1d 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
@@ -27,7 +27,7 @@ from sqlalchemy.exc import NoResultFound
from airflow.api.common.trigger_dag import trigger_dag
from airflow.api_fastapi.common.dagbag import DagBagDep, get_dag_for_run,
resolve_run_on_latest_version
-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.execution_api.datamodels.dagrun import
DagRunStateResponse, TriggerDAGRunPayload
@@ -223,15 +223,17 @@ def clear_dag_run(
"/{dag_id}/{run_id}/state",
responses={status.HTTP_404_NOT_FOUND: {"description": "Dag run not
found"}},
)
-def get_dagrun_state(
+async def get_dagrun_state(
dag_id: str,
run_id: str,
- session: SessionDep,
+ session: AsyncSessionDep,
) -> DagRunStateResponse:
"""Get a Dag run State."""
try:
- state: DagRunState = session.scalars(
- select(DagRunModel.state).where(DagRunModel.dag_id == dag_id,
DagRunModel.run_id == run_id)
+ state: DagRunState = (
+ await session.scalars(
+ select(DagRunModel.state).where(DagRunModel.dag_id == dag_id,
DagRunModel.run_id == run_id)
+ )
).one()
except NoResultFound:
raise HTTPException(
@@ -245,9 +247,9 @@ def get_dagrun_state(
@router.get("/count", status_code=status.HTTP_200_OK)
-def get_dr_count(
+async def get_dr_count(
dag_id: str,
- session: SessionDep,
+ session: AsyncSessionDep,
logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
run_ids: Annotated[list[str] | None, Query()] = None,
states: Annotated[list[str] | None, Query()] = None,
@@ -260,7 +262,7 @@ def get_dr_count(
stmt = stmt.where(DagRunModel.run_id.in_(run_ids))
if states:
stmt = stmt.where(DagRunModel.state.in_(states))
- return session.scalar(stmt) or 0
+ return (await session.scalar(stmt)) or 0
@router.get("/previous", status_code=status.HTTP_200_OK)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
index eb70fbd60ad..395a3ea30e8 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py
@@ -24,7 +24,7 @@ from fastapi import HTTPException, Security, status
from sqlalchemy import select
from airflow._shared.timezones import timezone
-from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.common.db.common import AsyncSessionDep, SessionDep
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.hitl import (
HITLDetailRequest,
@@ -154,13 +154,15 @@ def update_hitl_detail(
"/{task_instance_id}",
status_code=status.HTTP_200_OK,
)
-def get_hitl_detail(
+async def get_hitl_detail(
task_instance_id: UUID,
- session: SessionDep,
+ session: AsyncSessionDep,
) -> HITLDetailResponse:
"""Get Human-in-the-loop detail for a specific Task Instance."""
- hitl_detail_model_result = session.execute(
- select(HITLDetail).where(HITLDetail.ti_id == task_instance_id),
+ hitl_detail_model_result = (
+ await session.execute(
+ select(HITLDetail).where(HITLDetail.ti_id == task_instance_id),
+ )
).scalar()
hitl_detail_model = _check_hitl_detail_exists(hitl_detail_model_result)
return HITLDetailResponse.from_hitl_detail_orm(hitl_detail_model)
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 aa2c3ac5b11..4b5402d2342 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
@@ -1222,8 +1222,8 @@ def ti_patch_rendered_map_index(
]
),
)
-def get_previous_successful_dagrun(
- task_instance_id: UUID, session: SessionDep
+async def get_previous_successful_dagrun(
+ task_instance_id: UUID, session: AsyncSessionDep
) -> PrevSuccessfulDagRunResponse:
"""
Get the previous successful DagRun for a TaskInstance.
@@ -1233,12 +1233,12 @@ def get_previous_successful_dagrun(
bind_contextvars(ti_id=str(task_instance_id))
log.debug("Retrieving previous successful DAG run")
- task_instance = session.scalar(select(TI).where(TI.id == task_instance_id))
+ task_instance = await session.scalar(select(TI).where(TI.id ==
task_instance_id))
if not task_instance or not task_instance.logical_date:
log.debug("No task instance or logical date found")
return PrevSuccessfulDagRunResponse()
- dag_run = session.scalar(
+ dag_run = await session.scalar(
select(DR)
.where(
DR.dag_id == task_instance.dag_id,
@@ -1319,10 +1319,10 @@ def get_task_instance_count(
@router.get("/previous/{dag_id}/{task_id}", status_code=status.HTTP_200_OK)
-def get_previous_task_instance(
+async def get_previous_task_instance(
dag_id: str,
task_id: str,
- session: SessionDep,
+ session: AsyncSessionDep,
logical_date: Annotated[UtcDateTime | None, Query()] = None,
map_index: Annotated[int, Query()] = -1,
state: Annotated[TaskInstanceState | None, Query()] = None,
@@ -1351,7 +1351,7 @@ def get_previous_task_instance(
if state:
query = query.where(TI.state == state)
- ti = session.scalars(query.limit(1)).first()
+ ti = (await session.scalars(query.limit(1))).first()
if not ti:
return None
@@ -1420,11 +1420,15 @@ def get_task_instance_states(
@router.get("/breadcrumbs", status_code=status.HTTP_200_OK)
-def get_task_instance_breadcrumbs(dag_id: str, run_id: str, session:
SessionDep) -> TaskBreadcrumbsResponse:
- result = session.execute(
- select(TI.task_id, TI.map_index, TI.state, TI.operator, TI.duration)
- .where(TI.dag_id == dag_id, TI.run_id == run_id,
TI.state.in_(TerminalTIState))
- .order_by(TI.task_id, TI.map_index)
+async def get_task_instance_breadcrumbs(
+ dag_id: str, run_id: str, session: AsyncSessionDep
+) -> TaskBreadcrumbsResponse:
+ result = (
+ await session.execute(
+ select(TI.task_id, TI.map_index, TI.state, TI.operator,
TI.duration)
+ .where(TI.dag_id == dag_id, TI.run_id == run_id,
TI.state.in_(TerminalTIState))
+ .order_by(TI.task_id, TI.map_index)
+ )
).mappings()
def _iter_breadcrumbs() -> Iterator[dict[str, Any]]:
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py
index 3930fd945ea..c98d9417f50 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py
@@ -22,7 +22,7 @@ from uuid import UUID
from fastapi import APIRouter, Security, status
from sqlalchemy import select
-from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.common.db.common import AsyncSessionDep
from airflow.api_fastapi.common.types import UtcDateTime
from airflow.api_fastapi.execution_api.security import require_auth
from airflow.models.taskreschedule import TaskReschedule
@@ -40,9 +40,9 @@ router = APIRouter(
@router.get("/{task_instance_id}/start_date")
-def get_start_date(task_instance_id: UUID, session: SessionDep) -> UtcDateTime
| None:
+async def get_start_date(task_instance_id: UUID, session: AsyncSessionDep) ->
UtcDateTime | None:
"""Get the first reschedule date if found, None if no records exist."""
- start_date = session.scalar(
+ start_date = await session.scalar(
select(TaskReschedule.start_date)
.where(TaskReschedule.ti_id == task_instance_id)
.order_by(TaskReschedule.id.asc())
