From 047e0ad5b3754ef7bf8df9ab933b41da0bb1d6d1 Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Mon, 3 Nov 2025 12:53:51 +0100 Subject: [PATCH] [v3-1-test] Add number of queries guard for public hitls (#57549) * Add number of queries guard in public HITL list endpoints * Fix CI (cherry picked from commit bc6fe0ff93c6f5e97b651033f5bd5ac6ba0e61c0) Co-authored-by: Pierre Jeambrun --- .../src/airflow/api_fastapi/core_api/routes/public/hitl.py | 7 +++++-- .../unit/api_fastapi/core_api/routes/public/test_hitl.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py index d8635c7034351..361d5faf933f2 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py @@ -53,6 +53,7 @@ from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag from airflow.api_fastapi.logging.decorators import action_logging +from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun from airflow.models.hitl import HITLDetail as HITLDetailModel, HITLUser from airflow.models.taskinstance import TaskInstance as TI @@ -252,8 +253,10 @@ def get_hitl_details( .join(TI.dag_run) .options( joinedload(HITLDetailModel.task_instance).options( - joinedload(TI.dag_run), - ) + joinedload(TI.dag_run).joinedload(DagRun.dag_model), + joinedload(TI.task_instance_note), + joinedload(TI.dag_version).joinedload(DagVersion.bundle), + ), ) ) if dag_id != "~": diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py index 161e81607d414..2429f2ea09313 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py @@ -34,6 +34,7 @@ from airflow.sdk.execution_time.hitl import HITLUser from airflow.utils.state import TaskInstanceState +from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.format_datetime import from_datetime_to_zulu_without_ms if TYPE_CHECKING: @@ -526,7 +527,8 @@ def test_should_respond_200_with_existing_response( test_client: TestClient, expected_sample_hitl_detail_dict: dict[str, Any], ) -> None: - response = test_client.get("/dags/~/dagRuns/~/hitlDetails") + with assert_queries_count(3): + response = test_client.get("/dags/~/dagRuns/~/hitlDetails") assert response.status_code == 200 assert response.json() == { "hitl_details": [expected_sample_hitl_detail_dict], @@ -595,7 +597,8 @@ def test_should_respond_200_with_existing_response_and_query( params: dict[str, Any], expected_ti_count: int, ) -> None: - response = test_client.get("/dags/~/dagRuns/~/hitlDetails", params=params) + with assert_queries_count(3): + response = test_client.get("/dags/~/dagRuns/~/hitlDetails", params=params) assert response.status_code == 200 assert response.json()["total_entries"] == expected_ti_count assert len(response.json()["hitl_details"]) == expected_ti_count