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 e17f92776da12..5e9fae52e6574 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/security.py @@ -133,7 +133,7 @@ def inner( request: Request, user: GetUserDep, ) -> None: - dag_id: str | None = request.path_params.get("dag_id") + dag_id = request.path_params.get("dag_id") or request.query_params.get("dag_id") dag_id = dag_id if dag_id != "~" else None team_name = DagModel.get_team_name(dag_id) if dag_id else None diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py index 45f08f95d1e28..f1b4c16259251 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py @@ -89,7 +89,7 @@ class TestGetDagWarnings: ], ) def test_get_dag_warnings(self, test_client, query_params, expected_total_entries, expected_messages): - with assert_queries_count(3): + with assert_queries_count(3 if query_params.get("dag_id") is None else 4): response = test_client.get("/dagWarnings", params=query_params) assert response.status_code == 200 response_json = response.json() diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py index 8b7bdde96d577..f7574936428f0 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py @@ -153,7 +153,7 @@ def test_should_response_200( expected_response = [] for backfill in response_params: expected_response.append(backfill_responses[backfill]) - with assert_queries_count(2): + with assert_queries_count(2 if test_params.get("dag_id") is None else 3): response = test_client.get("/backfills", params=test_params) assert response.status_code == 200 assert response.json() == { diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index f8dd0116c6158..155b01094190e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -265,7 +265,7 @@ class TestStructureDataEndpoint: }, ], }, - 3, + 6, ), ( { @@ -273,7 +273,7 @@ class TestStructureDataEndpoint: "root": "unknown_task", }, {"edges": [], "nodes": []}, - 3, + 6, ), ( { @@ -298,7 +298,7 @@ class TestStructureDataEndpoint: }, ], }, - 3, + 6, ), ( {"dag_id": DAG_ID_EXTERNAL_TRIGGER, "external_dependencies": True}, @@ -337,7 +337,7 @@ class TestStructureDataEndpoint: }, ], }, - 10, + 13, ), ], ) @@ -534,7 +534,7 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a ], } - with assert_queries_count(10): + with assert_queries_count(13): response = test_client.get("/structure/structure_data", params=params) assert response.status_code == 200 assert response.json() == expected 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 96d7451c8f26f..f86a511aed01b 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 @@ -154,6 +154,7 @@ async def test_requires_access_dag_authorized(self, mock_get_auth_manager): mock_get_auth_manager.return_value = auth_manager fastapi_request = Mock() fastapi_request.path_params = {} + fastapi_request.query_params = {} requires_access_dag("GET", DagAccessEntity.CODE)(fastapi_request, Mock()) @@ -167,9 +168,11 @@ async def test_requires_access_dag_unauthorized(self, mock_get_auth_manager): mock_get_auth_manager.return_value = auth_manager fastapi_request = Mock() fastapi_request.path_params = {} + fastapi_request.query_params = {} mock_request = Mock() mock_request.path_params.return_value = {} + mock_request.query_params.return_value = {} with pytest.raises(HTTPException, match="Forbidden"): requires_access_dag("GET", DagAccessEntity.CODE)(fastapi_request, Mock())