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 1fba54ce54385..709cd7b8a65eb 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 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 @@ 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 ea58e610813bf..f8dd0116c6158 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.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 @@ def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, a ], } - 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