Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@
AssetWatcherModel,
TaskOutletAssetReference,
)
from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.typing_compat import Unpack
from airflow.utils.state import DagRunState
Expand Down Expand Up @@ -459,7 +460,10 @@ def materialize_asset(
if not get_auth_manager().is_authorized_dag(
method="POST",
access_entity=DagAccessEntity.RUN,
details=DagDetails(id=dag_id),
# The Dag is resolved from the asset here rather than named by the caller, so its team has
# to be looked up too. A team-aware auth manager distinguishes a team-scoped Dag from a
# global one by this field, so leaving it None asks the wrong question.
details=DagDetails(id=dag_id, team_name=DagModel.get_team_name(dag_id, session=session)),
Comment thread
potiuk marked this conversation as resolved.
user=user,
):
raise HTTPException(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -869,7 +869,10 @@ def wait_dag_run_until_finished(
if not get_auth_manager().is_authorized_dag(
method="GET",
access_entity=DagAccessEntity.XCOM,
details=DagDetails(id=dag_id),
# The route dependency above already authorizes RUN access with the Dag's team resolved;
# this second, XCom-specific check has to resolve it the same way, or the two checks ask
# a team-aware auth manager about differently-scoped resources.
details=DagDetails(id=dag_id, team_name=DagModel.get_team_name(dag_id, session=session)),
user=user,
):
if result_task_ids:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from sqlalchemy import delete, func, select, update

from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
from airflow.models import DagModel
from airflow.models.asset import (
Expand Down Expand Up @@ -2151,6 +2152,33 @@ def test_should_respond_200_with_partition_date_for_partitioned_dag(
assert dag_run.partition_key == "2025-06-01T00:00:00"
assert dag_run.partition_date == timezone.datetime(2025, 6, 1)

@pytest.mark.parametrize("team_name", ["team_b", None])
def test_authorizes_against_the_dags_team(self, test_client, session, team_name):
"""The Dag is resolved from the asset, so its team must be resolved and passed too — see the
call site's comment for why an unresolved team asks about the wrong resource."""
recorded = []

auth_manager = mock.Mock(spec=BaseAuthManager)
auth_manager.is_authorized_dag.side_effect = lambda **kw: recorded.append(kw) or True

with (
mock.patch(
"airflow.api_fastapi.core_api.routes.public.assets.get_auth_manager",
return_value=auth_manager,
),
mock.patch.object(
DagModel, "get_team_name", return_value=team_name, autospec=True
) as mock_get_team_name,
):
test_client.post("/assets/1/materialize")

assert len(recorded) == 1, "expected exactly one authorization check"
details = recorded[0]["details"]
assert details.id == self.DAG_ASSET1_ID
assert details.team_name == team_name
Comment thread
potiuk marked this conversation as resolved.
# resolved for the Dag the asset led to, not for some other Dag
mock_get_team_name.assert_called_once_with(self.DAG_ASSET1_ID, session=mock.ANY)


class TestGetAssetQueuedEvents(TestQueuedEventEndpoint):
@pytest.mark.usefixtures("time_freezer")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4394,6 +4394,32 @@ def test_should_respond_403_when_user_lacks_xcom_permission(self, test_client):
user=mock.ANY,
)

@pytest.mark.parametrize("team_name", ["team_b", None])
def test_authorizes_xcom_against_the_dags_team(self, test_client, team_name):
"""The XCom check must carry the Dag's team, matching what the route dependency above already
resolves — see the call site's comment for why."""
with (
mock.patch(
"airflow.api_fastapi.core_api.routes.public.dag_run.get_auth_manager",
autospec=True,
) as mock_get_auth_manager,
mock.patch.object(DagModel, "get_team_name", return_value=team_name, autospec=True),
):
mock_get_auth_manager.return_value.is_authorized_dag.return_value = True

response = test_client.get(
f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}/wait",
params={"interval": "1", "result": "task_1"},
)

assert response.status_code == 200
mock_get_auth_manager.return_value.is_authorized_dag.assert_called_once_with(
method="GET",
access_entity=DagAccessEntity.XCOM,
details=DagDetails(id=DAG1_ID, team_name=team_name),
user=mock.ANY,
)

def test_should_respond_200_without_result_when_user_lacks_xcom_permission(self, test_client):
"""Waiting without result parameter should not require XCom permissions."""
with mock.patch(
Expand Down