This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ff839573e94 Bound remaining single-row lookups with .limit(1) (#72699)
ff839573e94 is described below
commit ff839573e94f0c564680b15d1ff1bc87d0200728
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
---
.../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 fd1516b99ef..2140707ab2d 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
@@ -436,10 +436,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 625e2ee9116..bc00ccf07c9 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
@@ -1225,7 +1225,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 c376d43f759..b37ccf1a130 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
@@ -33,7 +34,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.config import conf_vars
from tests_common.test_utils.db import (
clear_db_assets,
@@ -1185,6 +1186,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 571057c12df..44a1fe5e088 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
@@ -58,6 +59,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,
@@ -3884,6 +3886,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):
