This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 04d9aaa62a1 Pass team name information to `is_authorized_dag` method
(#55040)
04d9aaa62a1 is described below
commit 04d9aaa62a111809d213756750cb943a8ed83eb3
Author: Vincent <[email protected]>
AuthorDate: Fri Aug 29 09:37:32 2025 -0400
Pass team name information to `is_authorized_dag` method (#55040)
---
.../auth/managers/models/resource_details.py | 1 +
.../src/airflow/api_fastapi/core_api/security.py | 6 +++-
airflow-core/src/airflow/models/dag.py | 14 +++++++++
.../unit/api_fastapi/core_api/test_security.py | 6 ++--
airflow-core/tests/unit/models/test_dag.py | 36 ++++++++++++++++++++++
devel-common/src/tests_common/test_utils/db.py | 7 +++++
6 files changed, 67 insertions(+), 3 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/managers/models/resource_details.py
b/airflow-core/src/airflow/api_fastapi/auth/managers/models/resource_details.py
index 038f82b3067..aed68f27aa8 100644
---
a/airflow-core/src/airflow/api_fastapi/auth/managers/models/resource_details.py
+++
b/airflow-core/src/airflow/api_fastapi/auth/managers/models/resource_details.py
@@ -42,6 +42,7 @@ class DagDetails:
"""Represents the details of a DAG."""
id: str | None = None
+ team_name: str | None = None
@dataclass
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/security.py
b/airflow-core/src/airflow/api_fastapi/core_api/security.py
index d4f6e11e7e9..258ca2dc13d 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/security.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/security.py
@@ -103,10 +103,14 @@ def requires_access_dag(
user: GetUserDep,
) -> None:
dag_id: str | None = request.path_params.get("dag_id")
+ team_name = DagModel.get_team_name(dag_id) if dag_id else None
_requires_access(
is_authorized_callback=lambda:
get_auth_manager().is_authorized_dag(
- method=method, access_entity=access_entity,
details=DagDetails(id=dag_id), user=user
+ method=method,
+ access_entity=access_entity,
+ details=DagDetails(id=dag_id, team_name=team_name),
+ user=user,
)
)
diff --git a/airflow-core/src/airflow/models/dag.py
b/airflow-core/src/airflow/models/dag.py
index 99f881c8055..fad01299377 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -76,6 +76,7 @@ from airflow.models.asset import (
)
from airflow.models.base import Base, StringID
from airflow.models.dag_version import DagVersion
+from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import RUN_ID_REGEX, DagRun
from airflow.models.taskinstance import (
TaskInstance,
@@ -83,6 +84,7 @@ from airflow.models.taskinstance import (
clear_task_instances,
)
from airflow.models.tasklog import LogTemplate
+from airflow.models.team import Team
from airflow.sdk import TaskGroup
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey,
BaseAsset
from airflow.sdk.definitions.dag import DAG as TaskSDKDag, dag as
task_sdk_dag_decorator
@@ -2269,6 +2271,18 @@ class DagModel(Base):
# an empty dict as there's no asset info to get. This method should
thus return None.
return get_asset_triggered_next_run_info([self.dag_id],
session=session).get(self.dag_id, None)
+ @staticmethod
+ @provide_session
+ def get_team_name(dag_id: str, session=NEW_SESSION) -> str | None:
+ """Return the team name associated to a Dag or None if it is not owned
by a specific team."""
+ stmt = (
+ select(Team.name)
+ .join(DagBundleModel.teams)
+ .join(DagModel, DagModel.bundle_name == DagBundleModel.name)
+ .where(DagModel.dag_id == dag_id)
+ )
+ return session.scalar(stmt)
+
STATICA_HACK = True
globals()["kcah_acitats"[::-1].upper()] = False
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
b/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
index 92b4dbf2e7b..f407fbc7f8f 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
@@ -84,25 +84,27 @@ class TestFastApiSecurity:
auth_manager.get_user_from_token.assert_called_once_with(token_str)
+ @pytest.mark.db_test
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
async def test_requires_access_dag_authorized(self, mock_get_auth_manager):
auth_manager = Mock()
auth_manager.is_authorized_dag.return_value = True
mock_get_auth_manager.return_value = auth_manager
fastapi_request = Mock()
- fastapi_request.path_params.return_value = {}
+ fastapi_request.path_params = {}
requires_access_dag("GET", DagAccessEntity.CODE)(fastapi_request,
Mock())
auth_manager.is_authorized_dag.assert_called_once()
+ @pytest.mark.db_test
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
async def test_requires_access_dag_unauthorized(self,
mock_get_auth_manager):
auth_manager = Mock()
auth_manager.is_authorized_dag.return_value = False
mock_get_auth_manager.return_value = auth_manager
fastapi_request = Mock()
- fastapi_request.path_params.return_value = {}
+ fastapi_request.path_params = {}
mock_request = Mock()
mock_request.path_params.return_value = {}
diff --git a/airflow-core/tests/unit/models/test_dag.py
b/airflow-core/tests/unit/models/test_dag.py
index b9af8f15f2c..867dc5aaaa5 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -21,6 +21,7 @@ import datetime
import logging
import os
import pickle
+import uuid
from datetime import timedelta
from pathlib import Path
from typing import TYPE_CHECKING
@@ -64,6 +65,7 @@ from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.taskinstance import TaskInstance as TI
+from airflow.models.team import Team
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
@@ -92,6 +94,7 @@ from tests_common.test_utils.db import (
clear_db_dags,
clear_db_runs,
clear_db_serialized_dags,
+ clear_db_teams,
)
from tests_common.test_utils.mapping import expand_mapped_task
from tests_common.test_utils.mock_plugins import mock_plugin_manager
@@ -145,6 +148,18 @@ def test_dags_bundle(configure_testing_dag_bundle):
yield
[email protected]
+def testing_team():
+ from airflow.utils.session import create_session
+
+ with create_session() as session:
+ team = session.query(Team).filter_by(name="testing").one_or_none()
+ if not team:
+ team = Team(id=uuid.uuid4(), name="testing")
+ session.add(team)
+ yield team
+
+
def _create_dagrun(
dag: DAG,
*,
@@ -2086,6 +2101,8 @@ class TestDagModel:
clear_db_dags()
clear_db_assets()
clear_db_runs()
+ clear_db_dag_bundles()
+ clear_db_teams()
def setup_method(self):
self._clean()
@@ -2522,6 +2539,25 @@ class TestDagModel:
]
}
+ def test_get_team_name(self, testing_team):
+ session = settings.Session()
+ dag_bundle = DagBundleModel(name="testing-team")
+ dag_bundle.teams.append(testing_team)
+ session.add(dag_bundle)
+ session.flush()
+
+ dag_id = "test_get_team_name"
+ dag = DAG(dag_id, schedule=None)
+ orm_dag = DagModel(
+ dag_id=dag.dag_id,
+ bundle_name="testing-team",
+ is_stale=False,
+ )
+ session.add(orm_dag)
+ session.flush()
+ assert DagModel.get_dagmodel(dag_id) is not None
+ assert DagModel.get_team_name(dag_id, session=session) == "testing"
+
class TestQueries:
def setup_method(self) -> None:
diff --git a/devel-common/src/tests_common/test_utils/db.py
b/devel-common/src/tests_common/test_utils/db.py
index 5849ec5b1d1..28764aeb90d 100644
--- a/devel-common/src/tests_common/test_utils/db.py
+++ b/devel-common/src/tests_common/test_utils/db.py
@@ -355,6 +355,13 @@ def clear_db_dag_bundles():
session.query(DagBundleModel).delete()
+def clear_db_teams():
+ with create_session() as session:
+ from airflow.models.team import Team
+
+ session.query(Team).delete()
+
+
def clear_dag_specific_permissions():
if "FabAuthManager" not in conf.get("core", "auth_manager"):
return
