This is an automated email from the ASF dual-hosted git repository.
rahulvats pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-2-test by this push:
new 84f7eaf3852 [v3-2-test] Fix get_team_name_dep creating wasted async
sessions when multi_team=False (#65275) (#65282)
84f7eaf3852 is described below
commit 84f7eaf38520ff47b26f16211ff07dc0f6eb057c
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Apr 15 11:56:38 2026 +0530
[v3-2-test] Fix get_team_name_dep creating wasted async sessions when
multi_team=False (#65275) (#65282)
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 <[email protected]>
---
.../airflow/api_fastapi/execution_api/security.py | 17 ++++++++------
.../api_fastapi/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 215997d28d9..19a29f7405d 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 import Depends, HTTPException, Request, status
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 @@ class ExecutionAPIRoute(APIRoute):
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 8fff2c9f732..29ca8901efb 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 import APIRouter, FastAPI, Request, Security
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 @@ class TestTokenTypeScopeEnforcement:
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()
