This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch backport-ff839573e9-v3-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 484125ba6c6b653d4a6eff7903d2b4d08925d1a5 Author: Pierre Jeambrun <[email protected]> AuthorDate: Wed Sep 9 18:21:30 2026 +0200 Bound remaining single-row lookups with .limit(1) (#72699) * Bound remaining single-row lookups with .limit(1) Two more callers of ``session.scalars(...).first()`` / ``session.execute(select(...)).first()`` were still missing an explicit ``.limit(1)``. Because SQLAlchemy 2.0's ``Result.first()`` does not add one (unlike the legacy ``Query.first``), each call sends every matching row to the Python process and discards all but the first. The favorite existence check is bounded to a single row per (dag_id, user_id) by the schema, but the previous-TI lookup scans every earlier matching TI. Introduced independently of #52325's migration: get_previous_task_instance came in #59712, unfavorite_dag's existence probe in #51264. Same anti-pattern as the four sites #72554 fixes; the tests reuse the ``capture_orm_selects`` helper that PR introduces. * Also assert the returned row on the previous-TI regression test (cherry picked from commit ff839573e94f0c564680b15d1ff1bc87d0200728) --- .../api_fastapi/core_api/routes/public/dags.py | 4 +++- .../execution_api/routes/task_instances.py | 2 +- .../core_api/routes/public/test_dags.py | 18 +++++++++++++++- .../versions/head/test_task_instances.py | 25 ++++++++++++++++++++++ 4 files changed, 46 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py index 240c93220b1..eeb034cf490 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py @@ -426,10 +426,12 @@ def unfavorite_dag(dag_id: str, session: SessionDep, user: GetUserDep): user_id = str(user.get_id()) favorite_exists = session.execute( - select(DagFavorite).where( + select(DagFavorite) + .where( DagFavorite.dag_id == dag_id, DagFavorite.user_id == user_id, ) + .limit(1) ).first() if not favorite_exists: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c5415e45e77..83d8dc7c51e 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -1200,7 +1200,7 @@ def get_previous_task_instance( if state: query = query.where(TI.state == state) - ti = session.scalars(query).first() + ti = session.scalars(query.limit(1)).first() if not ti: return None diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py index eb5d4a9f918..93795d1befe 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import re from datetime import datetime, timedelta, timezone from unittest import mock @@ -31,7 +32,7 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType -from tests_common.test_utils.asserts import assert_queries_count, count_queries +from tests_common.test_utils.asserts import assert_queries_count, capture_orm_selects, count_queries from tests_common.test_utils.db import ( clear_db_assets, clear_db_connections, @@ -952,6 +953,21 @@ class TestUnfavoriteDag(TestDagEndpoint): response = test_client.post(f"/dags/{DAG1_ID}/unfavorite") assert response.status_code == 409 + def test_unfavorite_dag_existence_check_is_bounded(self, test_client, session): + """The existing-favorite existence probe must ask the DB for one row.""" + session.execute(insert(DagFavorite).values(dag_id=DAG1_ID, user_id="test")) + session.commit() + + with capture_orm_selects("dag_favorite") as statements: + response = test_client.post(f"/dags/{DAG1_ID}/unfavorite") + + assert response.status_code == 204 + assert statements, "expected the endpoint to query the dag_favorite table" + for sql in statements: + assert re.search(r"\bLIMIT 1\b", sql), ( + f"favorite existence check is not bounded to one row: {sql}" + ) + class TestDagDetails(TestDagEndpoint): """Unit tests for DAG Details.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index bb3c0f7e5a7..65bbe0ff77c 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -17,6 +17,7 @@ from __future__ import annotations +import re from datetime import datetime from types import SimpleNamespace from typing import TYPE_CHECKING @@ -57,6 +58,7 @@ from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group from airflow.state.metastore import MetastoreBackend from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState +from tests_common.test_utils.asserts import capture_orm_selects from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import ( clear_db_assets, @@ -3712,6 +3714,29 @@ class TestGetPreviousTI: assert data["run_id"] == "target_run_1" assert data["state"] == State.SUCCESS + def test_get_previous_ti_query_is_bounded(self, client, session, create_task_instance): + """The single-row previous-TI lookup must ask the DB for one row.""" + for i in range(5): + create_task_instance( + task_id="test_task", + state=State.SUCCESS, + logical_date=timezone.datetime(2025, 1, i + 1), + run_id=f"run{i + 1}", + ) + session.commit() + + with capture_orm_selects("task_instance") as statements: + response = client.get( + "/execution/task-instances/previous/dag/test_task", + params={"logical_date": "2025-01-05T00:00:00Z"}, + ) + + assert response.status_code == 200 + assert response.json()["run_id"] == "run4" + assert statements, "expected the endpoint to query the task_instance table" + for sql in statements: + assert re.search(r"\bLIMIT 1\b", sql), f"previous-TI lookup is not bounded to one row: {sql}" + class TestGetTaskStates: def setup_method(self):
