From 450cb99131a55bdcdb367bd055951749694dfd13 Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Tue, 14 Jul 2026 20:13:22 +0530 Subject: [PATCH] Execution API: Resolve team lookup from TI history --- .../api_fastapi/execution_api/security.py | 16 +++++++-- .../execution_api/test_security.py | 35 +++++++++++++++++++ 2 files changed, 49 insertions(+), 2 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 9de3493061f32..d380fe855ec11 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -75,7 +75,7 @@ from fastapi.routing import APIRoute from fastapi.security import HTTPBearer, SecurityScopes from pydantic import ValidationError -from sqlalchemy import select +from sqlalchemy import select, union_all from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken, TokenScope @@ -264,9 +264,10 @@ def _team_name_for_ti_stmt(ti_id): """Build the select statement resolving ``TaskInstance.id -> Team.name``.""" from airflow.models import DagModel, TaskInstance from airflow.models.dagbundle import DagBundleModel + from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.models.team import Team - return ( + task_instance_stmt = ( select(Team.name) .select_from(TaskInstance) .join(DagModel, DagModel.dag_id == TaskInstance.dag_id) @@ -275,6 +276,17 @@ def _team_name_for_ti_stmt(ti_id): .where(TaskInstance.id == ti_id) ) + task_instance_history_stmt = ( + select(Team.name) + .select_from(TaskInstanceHistory) + .join(DagModel, DagModel.dag_id == TaskInstanceHistory.dag_id) + .join(DagBundleModel, DagBundleModel.name == DagModel.bundle_name) + .join(DagBundleModel.teams) + .where(TaskInstanceHistory.task_instance_id == ti_id) + ) + + return union_all(task_instance_stmt, task_instance_history_stmt) + def _team_name_for_dag_stmt(dag_id): """Build the select statement resolving ``DagModel.dag_id -> Team.name``.""" 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 e98d918385933..bf62f4508c9c4 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 @@ -30,7 +30,9 @@ from airflow.api_fastapi.execution_api.security import ( ExecutionAPIRoute, _jwt_bearer, + _team_name_for_ti_stmt, get_team_name_dep, + get_team_name_for_ti, require_auth, ) @@ -273,3 +275,36 @@ async def test_returns_none_without_session_when_multi_team_disabled(self): assert result is None mock_create_session.assert_not_called() + + +class TestGetTeamNameForTI: + @pytest.mark.parametrize( + ("multi_team", "expected_result"), + [ + (True, "team_a"), + (False, None), + ], + ) + def test_get_team_name_for_ti_respects_multi_team(self, multi_team, expected_result): + session = MagicMock() + session.scalar.return_value = "team_a" + + with patch("airflow.configuration.conf.getboolean", return_value=multi_team): + result = get_team_name_for_ti(UUID("d9edb890-fa95-4049-b66f-b6469d4fbc32"), session) + + assert result == expected_result + if multi_team: + session.scalar.assert_called_once() + else: + session.scalar.assert_not_called() + + +class TestTeamNameStatement: + def test_team_name_statement_checks_task_instance_and_history(self): + stmt = _team_name_for_ti_stmt(UUID("da17bd0e-aa95-4995-8c10-26bf5606a9fb")) + rendered = str(stmt) + + assert "UNION ALL" in rendered + assert "FROM task_instance" in rendered + assert "FROM task_instance_history" in rendered + assert "task_instance_history.task_instance_id" in rendered