seanghaeli commented on code in PR #75377:
URL: https://github.com/apache/airflow/pull/75377#discussion_r4383520081
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -125,10 +125,10 @@ def has_xcom_access(
from airflow.api_fastapi.execution_api.security import (
_team_name_for_dag_stmt,
- _team_name_for_ti_stmt,
+ get_team_name_for_execution_identity,
)
- requester_team = session.scalar(_team_name_for_ti_stmt(token.id))
+ requester_team = get_team_name_for_execution_identity(token.id, session)
Review Comment:
And the matching call site:
```suggestion
requester_team = get_team_name_for_execution_identity(token.id,
session=session)
```
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_security.py:
##########
@@ -604,6 +604,122 @@ async def
test_dag_processor_token_resolves_the_team_of_its_bundle(self, session
assert result == "team_a"
+ @pytest.mark.db_test
+ @pytest.mark.asyncio
+ @pytest.mark.usefixtures("async_db_engine")
+ async def
test_callback_execution_token_resolves_team_from_its_bundle(self, session):
+ team = Team(name=f"team_{uuid4().hex}")
+ bundle = DagBundleModel(name=f"bundle_{uuid4().hex}")
+ bundle.teams.append(team)
+ callback = Callback()
+ callback.bundle_name = bundle.name
+ callback.fetch_method = CallbackFetchMethod.IMPORT_PATH
+ session.add_all([bundle, callback])
+ session.commit()
+ token = TIToken(id=callback.id, claims=TIClaims(scope="execution"))
+
+ try:
+ with conf_vars({("core", "multi_team"): "True"}):
+ result = await get_team_name_dep(token=token)
+ finally:
+ session.delete(callback)
+ session.delete(bundle)
+ session.delete(team)
+ session.commit()
+
+ assert result == team.name
+
+ @pytest.mark.db_test
+ @pytest.mark.asyncio
+ @pytest.mark.usefixtures("async_db_engine")
+ async def test_task_execution_token_still_resolves_team_from_its_dag(self,
session, create_task_instance):
Review Comment:
nit: this one and `test_teamless_callback_execution_token_remains_teamless`
also pass on main, so they don't cover anything new. I'd drop them, or fold the
teamless case into a parametrized version of the callback test above.
##########
airflow-core/src/airflow/api_fastapi/execution_api/security.py:
##########
@@ -526,6 +527,41 @@ def get_team_name_for_ti(ti_id, session) -> str | None:
return session.scalar(_team_name_for_ti_stmt(ti_id))
+def get_team_name_for_execution_identity(identity_id, session) -> str | None:
Review Comment:
nit: could `session` be keyword-only here? That's the convention for new
airflow-core code.
```suggestion
def get_team_name_for_execution_identity(identity_id, *, session) -> str |
None:
```
--
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]
