From f2c3ae6ff9c498db428e302ca337aa0e179bb248 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Wed, 15 Apr 2026 06:11:52 +0100 Subject: [PATCH] [v3-2-test] Fix get_team_name_dep creating wasted async sessions when multi_team=False (#65275) get_team_name_dep() declared AsyncSessionDep as a FastAPI dependency parameter, which meant FastAPI checked out an async database connection on every Execution API request -- even when multi_team=False (the default), where the function immediately returned None without using the session. Remove the AsyncSessionDep parameter and create the session inside the function only when multi_team=True. This avoids a connection checkout, no-op commit round-trip, and connection return on every connection and variable lookup from workers. (cherry picked from commit 1e64f418ad55c2fa8a9a56f8c544ac4a247ed1c1) Co-authored-by: Kaxil Naik --- .../api_fastapi/execution_api/security.py | 17 +++++++----- .../execution_api/test_security.py | 26 ++++++++++++++++++- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index 215997d28d9c7..19a29f7405da7 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -74,10 +74,8 @@ from fastapi.params import Security as SecurityParam from fastapi.routing import APIRoute from fastapi.security import HTTPBearer, SecurityScopes -from sqlalchemy import select from airflow.api_fastapi.auth.tokens import JWTValidator -from airflow.api_fastapi.common.db.common import AsyncSessionDep from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.deps import DepContainer @@ -222,16 +220,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.allowed_token_types = frozenset(token_scopes) if token_scopes else frozenset({"execution"}) -async def get_team_name_dep(session: AsyncSessionDep, token=CurrentTIToken) -> str | None: +async def get_team_name_dep(token=CurrentTIToken) -> str | None: """Return the team name associated to the task (if any).""" from airflow.configuration import conf - from airflow.models import DagModel, TaskInstance - from airflow.models.dagbundle import DagBundleModel - from airflow.models.team import Team if not conf.getboolean("core", "multi_team"): return None + from sqlalchemy import select + + from airflow.models import DagModel, TaskInstance + from airflow.models.dagbundle import DagBundleModel + from airflow.models.team import Team + from airflow.utils.session import create_session_async + stmt = ( select(Team.name) .select_from(TaskInstance) @@ -240,4 +242,5 @@ async def get_team_name_dep(session: AsyncSessionDep, token=CurrentTIToken) -> s .join(DagBundleModel.teams) .where(TaskInstance.id == token.id) ) - return await session.scalar(stmt) + async with create_session_async() as session: + return await session.scalar(stmt) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py index 8fff2c9f7322a..29ca8901efb53 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +from unittest.mock import MagicMock, patch from uuid import UUID import pytest @@ -23,7 +24,12 @@ from fastapi.testclient import TestClient from airflow.api_fastapi.execution_api.datamodels.token import TIToken -from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute, _jwt_bearer, require_auth +from airflow.api_fastapi.execution_api.security import ( + ExecutionAPIRoute, + _jwt_bearer, + get_team_name_dep, + require_auth, +) class TestExecutionAPIRoute: @@ -136,3 +142,21 @@ def test_execution_token_accepted_on_both_routes(self, token_type_app): run = client.get(f"/task-instances/{self.TI_ID}/run", headers={"Authorization": "Bearer fake"}) assert state.status_code == 200 assert run.status_code == 200 + + +class TestGetTeamNameDep: + """Tests for get_team_name_dep avoiding unnecessary async sessions.""" + + @pytest.mark.asyncio + async def test_returns_none_without_session_when_multi_team_disabled(self): + """When multi_team=False, no async session should be created.""" + token = MagicMock(spec=TIToken) + + with ( + patch("airflow.configuration.conf.getboolean", return_value=False), + patch("airflow.utils.session.create_session_async") as mock_create_session, + ): + result = await get_team_name_dep(token=token) + + assert result is None + mock_create_session.assert_not_called()