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 adca02f2e7796..5dc543e3dd1f2 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/security.py @@ -150,7 +150,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 0d9eb862e29f6..68fb9473d4d6a 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 63a15042dbb88..b14ef9cf1a069 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 ded50d54a30af..f4a56db93e4a4 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 @@ -303,7 +303,7 @@ class TestStructureDataEndpoint: }, ], }, - 3, + 6, ), ( { @@ -311,7 +311,7 @@ class TestStructureDataEndpoint: "root": "unknown_task", }, {"edges": [], "nodes": []}, - 3, + 6, ), ( { @@ -336,7 +336,7 @@ class TestStructureDataEndpoint: }, ], }, - 3, + 6, ), ( {"dag_id": DAG_ID_EXTERNAL_TRIGGER, "external_dependencies": True}, @@ -375,7 +375,7 @@ class TestStructureDataEndpoint: }, ], }, - 10, + 13, ), ], ) @@ -572,7 +572,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 b00001474386a..5aaf509b66bde 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 @@ 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 @@ 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())