diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_versions.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_versions.py index e4f282f60880b..5f1c734148b2b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_versions.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_versions.py @@ -108,7 +108,7 @@ def get_dag_versions( This endpoint allows specifying `~` as the dag_id to retrieve DAG Versions for all DAGs. """ - query = select(DagVersion).options(joinedload(DagVersion.dag_model)) + query = select(DagVersion).options(joinedload(DagVersion.dag_model), joinedload(DagVersion.bundle)) if dag_id != "~": get_latest_version_of_dag(dag_bag, dag_id, session) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_versions.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_versions.py index 50307f8502da1..506875b2e351b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_versions.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_versions.py @@ -22,6 +22,7 @@ from airflow.providers.standard.operators.empty import EmptyOperator +from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_serialized_dags pytestmark = pytest.mark.db_test @@ -212,7 +213,7 @@ def test_should_respond_403(self, unauthorized_test_client): class TestGetDagVersions(TestDagVersionEndpoint): @pytest.mark.parametrize( - "dag_id, expected_response", + "dag_id, expected_response, expected_query_count", [ [ "~", @@ -261,6 +262,7 @@ class TestGetDagVersions(TestDagVersionEndpoint): ], "total_entries": 4, }, + 2, ], [ "dag_with_multiple_versions", @@ -299,17 +301,19 @@ class TestGetDagVersions(TestDagVersionEndpoint): ], "total_entries": 3, }, + 4, ], ], ) @pytest.mark.usefixtures("make_dag_with_multiple_versions") - def test_get_dag_versions(self, test_client, dag_id, expected_response): - response = test_client.get(f"/dags/{dag_id}/dagVersions") + def test_get_dag_versions(self, test_client, dag_id, expected_response, expected_query_count): + with assert_queries_count(expected_query_count): + response = test_client.get(f"/dags/{dag_id}/dagVersions") assert response.status_code == 200 assert response.json() == expected_response @pytest.mark.parametrize( - "dag_id, expected_response", + "dag_id, expected_response, expected_query_count", [ [ "~", @@ -358,6 +362,7 @@ def test_get_dag_versions(self, test_client, dag_id, expected_response): ], "total_entries": 4, }, + 2, ], [ "dag_with_multiple_versions", @@ -396,12 +401,16 @@ def test_get_dag_versions(self, test_client, dag_id, expected_response): ], "total_entries": 3, }, + 4, ], ], ) @pytest.mark.usefixtures("make_dag_with_multiple_versions") - def test_get_dag_versions_with_url_template(self, test_client, dag_id, expected_response): - response = test_client.get(f"/dags/{dag_id}/dagVersions") + def test_get_dag_versions_with_url_template( + self, test_client, dag_id, expected_response, expected_query_count + ): + with assert_queries_count(expected_query_count): + response = test_client.get(f"/dags/{dag_id}/dagVersions") assert response.status_code == 200 assert response.json() == expected_response @@ -479,7 +488,8 @@ def test_get_dag_versions_with_url_template(self, test_client, dag_id, expected_ def test_get_dag_versions_parameters( self, test_client, params, expected_versions, expected_total_entries ): - response = test_client.get("/dags/~/dagVersions", params=params) + with assert_queries_count(2): + response = test_client.get("/dags/~/dagVersions", params=params) assert response.status_code == 200 response_payload = response.json() assert response_payload["total_entries"] == expected_total_entries