potiuk commented on code in PR #67998:
URL: https://github.com/apache/airflow/pull/67998#discussion_r3680764936
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_security.py:
##########
@@ -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):
Review Comment:
Asserting `"UNION ALL" in rendered` and `"FROM task_instance_history" in
rendered` tests the SQL text rather than the behaviour — it would pass on a
statement that compiles correctly but resolves the wrong team, and it'll break
on any cosmetic change to how SQLAlchemy renders the query.
What would actually protect this path is a DB-backed test: create a Dag on a
bundle owned by a team, archive a try into `task_instance_history`, then assert
`get_team_name_for_ti(...)` returns that team — and returns `None` for a TI
whose bundle has no team. Given this decides authorization scope, behavioural
coverage seems worth the extra setup.
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_security.py:
##########
@@ -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):
Review Comment:
`patch("airflow.configuration.conf.getboolean", return_value=multi_team)`
replaces *every* `getboolean` call for the duration, not just `core/multi_team`
— so anything else consulting a boolean config inside this block silently gets
the parametrized value too.
Airflow's convention here is `conf_vars` from
`tests_common.test_utils.config`, which scopes the override to the specific
key. Since the value varies with `@pytest.mark.parametrize`, use it as a
context manager:
```python
with conf_vars({("core", "multi_team"): str(multi_team)}):
result = get_team_name_for_ti(UUID(...), session)
```
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
##########
airflow-core/src/airflow/api_fastapi/execution_api/security.py:
##########
@@ -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)
Review Comment:
Worth a thought on `union_all` vs `union`. `task_instance_id` is the sole
primary key of `task_instance_history`, so this is at most two rows — one live,
one archived — rather than one per try, which is much better than I first
assumed.
But for a retried TI both branches do match, and the caller uses
`session.scalar()`, which takes the first row with no `ORDER BY`. Normally both
resolve to the same team, so it's benign; if a Dag ever moved between bundles
owned by different teams, the answer becomes whichever row the backend happened
to return first. `union` would collapse identical values and make the duplicate
case explicit if it ever isn't.
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
