This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch v3-1-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-1-test by this push:
new abcac64f624 Add number of queries guard for ui structure (#58051)
(#58632)
abcac64f624 is described below
commit abcac64f624db48d352268e995b7302f6aa876f2
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Tue Nov 25 11:46:40 2025 +0100
Add number of queries guard for ui structure (#58051) (#58632)
(cherry picked from commit 213294e573e67a63d15d0254d5226817e4c94520)
---
.../airflow/api_fastapi/core_api/routes/ui/structure.py | 3 +++
.../unit/api_fastapi/core_api/routes/ui/test_structure.py | 15 +++++++++++----
2 files changed, 14 insertions(+), 4 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
index 1fba54ce543..709cd7b8a65 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
@@ -18,6 +18,7 @@ from __future__ import annotations
from fastapi import Depends, HTTPException, status
from sqlalchemy import select
+from sqlalchemy.orm import joinedload
from airflow.api_fastapi.auth.managers.models.resource_details import
DagAccessEntity
from airflow.api_fastapi.common.db.common import SessionDep
@@ -31,6 +32,7 @@ from airflow.api_fastapi.core_api.services.ui.structure
import (
get_upstream_assets,
)
from airflow.api_fastapi.core_api.services.ui.task_group import
task_group_to_dict
+from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.models.serialized_dag import SerializedDagModel
from airflow.utils.dag_edges import dag_edges
@@ -70,6 +72,7 @@ def structure_data(
select(SerializedDagModel)
.join(DagVersion)
.where(SerializedDagModel.dag_id == dag_id, DagVersion.version_number
== version_number)
+
.options(joinedload(SerializedDagModel.dag_model).joinedload(DagModel.task_outlet_asset_references)),
)
if serialized_dag is None:
raise HTTPException(
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 ea58e610813..f8dd0116c61 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
@@ -33,6 +33,7 @@ from airflow.providers.standard.sensors.external_task import
ExternalTaskSensor
from airflow.sdk import Metadata, task
from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset
+from tests_common.test_utils.asserts import assert_queries_count
from tests_common.test_utils.db import clear_db_assets, clear_db_runs
pytestmark = pytest.mark.db_test
@@ -207,7 +208,7 @@ def asset3_id(make_dags, asset3, session) -> str:
class TestStructureDataEndpoint:
@pytest.mark.parametrize(
- "params, expected",
+ ("params", "expected", "expected_queries_count"),
[
(
{"dag_id": DAG_ID},
@@ -264,6 +265,7 @@ class TestStructureDataEndpoint:
},
],
},
+ 3,
),
(
{
@@ -271,6 +273,7 @@ class TestStructureDataEndpoint:
"root": "unknown_task",
},
{"edges": [], "nodes": []},
+ 3,
),
(
{
@@ -295,6 +298,7 @@ class TestStructureDataEndpoint:
},
],
},
+ 3,
),
(
{"dag_id": DAG_ID_EXTERNAL_TRIGGER, "external_dependencies":
True},
@@ -333,12 +337,14 @@ class TestStructureDataEndpoint:
},
],
},
+ 10,
),
],
)
@pytest.mark.usefixtures("make_dags")
- def test_should_return_200(self, test_client, params, expected):
- response = test_client.get("/structure/structure_data", params=params)
+ def test_should_return_200(self, test_client, params, expected,
expected_queries_count):
+ with assert_queries_count(expected_queries_count):
+ response = test_client.get("/structure/structure_data",
params=params)
assert response.status_code == 200
assert response.json() == expected
@@ -528,7 +534,8 @@ class TestStructureDataEndpoint:
],
}
- response = test_client.get("/structure/structure_data", params=params)
+ with assert_queries_count(10):
+ response = test_client.get("/structure/structure_data",
params=params)
assert response.status_code == 200
assert response.json() == expected
