This is an automated email from the ASF dual-hosted git repository.
jason810496 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 91b18259e93 Require explicit limits for single-row lookups in core
(#72856)
91b18259e93 is described below
commit 91b18259e938bf0a00a4a6aa4e3dddf719add301
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Sat Sep 12 20:27:01 2026 +0800
Require explicit limits for single-row lookups in core (#72856)
* Require explicit limits for single-row lookups in core
Guard core first() lookups against unbounded SQLAlchemy results with a
core-specific prek hook. Follow local query assignments, require a
statement-level limit, and bound the remaining operator-link XCom lookup.
* Name first-limit diagnostic fields
---
airflow-core/.pre-commit-config.yaml | 6 +
.../api_fastapi/core_api/routes/public/xcom.py | 3 +-
.../serialization/definitions/operatorlink.py | 4 +-
.../tests/unit/serialization/test_operatorlink.py | 43 +++++
scripts/ci/prek/check_first_limit.py | 215 +++++++++++++++++++++
scripts/tests/ci/prek/test_check_first_limit.py | 141 ++++++++++++++
6 files changed, 409 insertions(+), 3 deletions(-)
diff --git a/airflow-core/.pre-commit-config.yaml
b/airflow-core/.pre-commit-config.yaml
index c59444e234d..d7b96c960d7 100644
--- a/airflow-core/.pre-commit-config.yaml
+++ b/airflow-core/.pre-commit-config.yaml
@@ -24,6 +24,12 @@ default_language_version:
repos:
- repo: local
hooks:
+ - id: check-first-limit
+ name: Check single-row lookups have an explicit limit
+ entry: ../scripts/ci/prek/check_first_limit.py
+ language: python
+ files: ^src/airflow/.*\.py$
+ pass_filenames: true
- id: check-taskinstance-tis-attrs
name: Check that TI and TIS have the same attributes
entry: ../scripts/ci/prek/check_ti_vs_tis_attributes.py
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
index e2c644ea16a..48ed3d11bf1 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
@@ -94,7 +94,6 @@ def get_xcom_entry(
task_ids=task_id,
dag_ids=dag_id,
map_indexes=map_index,
- limit=1,
).options(
joinedload(XComModel.task),
joinedload(XComModel.dag_run).joinedload(DR.dag_model),
@@ -104,7 +103,7 @@ def get_xcom_entry(
# We use `BaseXCom.get_many` to fetch XComs directly from the database,
bypassing the XCom Backend.
# This avoids deserialization via the backend (e.g., from a remote storage
like S3) and instead
# retrieves the raw serialized value from the database.
- raw_result: tuple[XComModel] | None = session.scalars(xcom_query).first()
+ raw_result: tuple[XComModel] | None =
session.scalars(xcom_query.limit(1)).first()
if raw_result is None:
raise HTTPException(status.HTTP_404_NOT_FOUND, f"XCom entry with key:
`{xcom_key}` not found")
diff --git a/airflow-core/src/airflow/serialization/definitions/operatorlink.py
b/airflow-core/src/airflow/serialization/definitions/operatorlink.py
index b7a6eef11b8..86e245b4c63 100644
--- a/airflow-core/src/airflow/serialization/definitions/operatorlink.py
+++ b/airflow-core/src/airflow/serialization/definitions/operatorlink.py
@@ -62,7 +62,9 @@ class XComOperatorLink(LoggingMixin):
dag_ids=ti_key.dag_id,
task_ids=ti_key.task_id,
map_indexes=ti_key.map_index,
- ).with_only_columns(XComModel.value)
+ )
+ .with_only_columns(XComModel.value)
+ .limit(1)
).first()
if not result:
self.log.debug(
diff --git a/airflow-core/tests/unit/serialization/test_operatorlink.py
b/airflow-core/tests/unit/serialization/test_operatorlink.py
new file mode 100644
index 00000000000..ef3a399279c
--- /dev/null
+++ b/airflow-core/tests/unit/serialization/test_operatorlink.py
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest import mock
+
+from sqlalchemy.engine import Result
+from sqlalchemy.orm import Session
+
+from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.serialization.definitions.operatorlink import XComOperatorLink
+
+
[email protected]("airflow.serialization.definitions.operatorlink.create_session",
autospec=True)
+def test_operator_link_lookup_is_bounded(create_session):
+ session = mock.create_autospec(Session, instance=True)
+ create_session.return_value.__enter__.return_value = session
+ result = mock.create_autospec(Result, instance=True)
+ result.first.return_value =
SimpleNamespace(value='"https://example.com/task";')
+ session.execute.return_value = result
+ link = XComOperatorLink(name="Example", xcom_key="_link_example")
+ ti_key = TaskInstanceKey(dag_id="example", task_id="task", run_id="run",
try_number=1, map_index=-1)
+
+ assert link.get_link(None, ti_key=ti_key) == "https://example.com/task";
+
+ session.execute.assert_called_once()
+ statement = session.execute.call_args.args[0]
+ assert "LIMIT 1" in str(statement.compile(compile_kwargs={"literal_binds":
True}))
diff --git a/scripts/ci/prek/check_first_limit.py
b/scripts/ci/prek/check_first_limit.py
new file mode 100755
index 00000000000..08b78180473
--- /dev/null
+++ b/scripts/ci/prek/check_first_limit.py
@@ -0,0 +1,215 @@
+#!/usr/bin/env python
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["rich>=13.6.0"]
+# ///
+"""Require an explicit statement-level limit for core's single-row lookups.
+
+Track straight-line local assignments, without inferring limits across helper
+calls or control-flow joins. Put ``.limit(1)`` at the execution site when a
+query's bound cannot be established locally.
+"""
+
+from __future__ import annotations
+
+import ast
+import sys
+from enum import Enum, auto
+from pathlib import Path
+from typing import NamedTuple
+
+from common_prek_utils import console
+
+
+class QueryState(Enum):
+ UNKNOWN = auto()
+ LIMITED = auto()
+ RESULT = auto()
+ LIMITED_RESULT = auto()
+
+
+class ErrorLocation(NamedTuple):
+ line: int
+ column: int
+
+
+class FirstLimitVisitor(ast.NodeVisitor):
+ def __init__(self) -> None:
+ self.bindings: dict[str, QueryState] = {}
+ self.errors: list[ErrorLocation] = []
+
+ def get_state(self, node: ast.AST | None) -> QueryState:
+ if isinstance(node, ast.Name):
+ return self.bindings.get(node.id, QueryState.UNKNOWN)
+ if isinstance(node, (ast.Await, ast.NamedExpr)):
+ return self.get_state(node.value)
+ if not isinstance(node, ast.Call) or not isinstance(node.func,
ast.Attribute):
+ return QueryState.UNKNOWN
+ method = node.func.attr
+ receiver = self.get_state(node.func.value)
+ if receiver in {QueryState.RESULT, QueryState.LIMITED_RESULT}:
+ if method in {"scalars", "unique", "mappings", "tuples",
"columns", "yield_per"}:
+ return receiver
+ return QueryState.RESULT
+ if method in {"execute", "scalars", "stream", "stream_scalars"} and (
+ node.args or any(keyword.arg == "statement" for keyword in
node.keywords)
+ ):
+ statement = (
+ node.args[0]
+ if node.args
+ else next(keyword.value for keyword in node.keywords if
keyword.arg == "statement")
+ )
+ return (
+ QueryState.LIMITED_RESULT
+ if self.get_state(statement) == QueryState.LIMITED
+ else QueryState.RESULT
+ )
+ if method == "limit":
+ if (
+ len(node.args) == 1
+ and not node.keywords
+ and isinstance(node.args[0], ast.Constant)
+ and type(node.args[0].value) is int
+ and node.args[0].value == 1
+ ):
+ return QueryState.LIMITED
+ return QueryState.UNKNOWN
+ if method in {
+ "where",
+ "filter",
+ "filter_by",
+ "order_by",
+ "options",
+ "with_only_columns",
+ "offset",
+ "join",
+ "outerjoin",
+ "select_from",
+ "with_for_update",
+ "execution_options",
+ "distinct",
+ }:
+ return receiver
+ return QueryState.UNKNOWN
+
+ def visit_Call(self, node: ast.Call) -> None:
+ self.generic_visit(node)
+ if (
+ isinstance(node.func, ast.Attribute)
+ and node.func.attr == "first"
+ and self.get_state(node.func.value) not in {QueryState.LIMITED,
QueryState.LIMITED_RESULT}
+ ):
+ self.errors.append(ErrorLocation(node.lineno, node.col_offset + 1))
+
+ def visit_Assign(self, node: ast.Assign) -> None:
+ self.visit(node.value)
+ state = self.get_state(node.value)
+ for target in node.targets:
+ self.set_binding(target, state)
+
+ def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+ if node.value is not None:
+ self.visit(node.value)
+ self.set_binding(node.target, self.get_state(node.value))
+
+ def visit_NamedExpr(self, node: ast.NamedExpr) -> None:
+ self.visit(node.value)
+ self.set_binding(node.target, self.get_state(node.value))
+
+ def set_binding(self, target: ast.AST, state: QueryState) -> None:
+ if isinstance(target, ast.Name):
+ self.bindings[target.id] = state
+ else:
+ self.invalidate_bindings(target)
+
+ def invalidate_bindings(self, node: ast.AST) -> None:
+ for child in ast.walk(node):
+ if isinstance(child, ast.Name) and isinstance(child.ctx,
(ast.Store, ast.Del)):
+ self.bindings.pop(child.id, None)
+
+ def visit_Name(self, node: ast.Name) -> None:
+ if isinstance(node.ctx, (ast.Store, ast.Del)):
+ self.bindings.pop(node.id, None)
+
+ def generic_visit(self, node: ast.AST) -> None:
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef,
ast.ClassDef, ast.Lambda)):
+ previous = self.bindings
+ self.bindings = {}
+ super().generic_visit(node)
+ self.bindings = previous
+ elif isinstance(
+ node,
+ (
+ ast.If,
+ ast.For,
+ ast.AsyncFor,
+ ast.While,
+ ast.Try,
+ ast.Match,
+ ast.ListComp,
+ ast.SetComp,
+ ast.DictComp,
+ ast.GeneratorExp,
+ ),
+ ):
+ self.invalidate_bindings(node)
+ previous = self.bindings.copy()
+ for field, value in ast.iter_fields(node):
+ self.bindings = previous.copy()
+ if isinstance(value, list):
+ for child in value:
+ if field in {"handlers", "cases"}:
+ self.bindings = previous.copy()
+ self.visit(child)
+ elif isinstance(value, ast.AST):
+ self.visit(value)
+ self.bindings = previous
+ else:
+ if isinstance(node, (ast.AugAssign, ast.Delete)):
+ self.invalidate_bindings(node)
+ super().generic_visit(node)
+
+
+def check_source(source: str, filename: str = "<unknown>") ->
list[ErrorLocation]:
+ visitor = FirstLimitVisitor()
+ visitor.visit(ast.parse(source, filename=filename))
+ return sorted(visitor.errors)
+
+
+def main(filenames: list[str]) -> int:
+ failed = False
+ for filename in filenames:
+ try:
+ errors = check_source(Path(filename).read_text(encoding="utf-8"),
filename)
+ except SyntaxError as error:
+ console.print(f"{filename}:{error.lineno}: {error.msg}",
markup=False, highlight=False)
+ failed = True
+ continue
+ for location in errors:
+ console.print(
+ f"{filename}:{location.line}:{location.column}: .first()
requires .limit(1) on the statement before execution.",
+ markup=False,
+ highlight=False,
+ )
+ failed = True
+ return int(failed)
+
+
+if __name__ == "__main__":
+ sys.exit(main(sys.argv[1:]))
diff --git a/scripts/tests/ci/prek/test_check_first_limit.py
b/scripts/tests/ci/prek/test_check_first_limit.py
new file mode 100644
index 00000000000..4d09fdc5fb0
--- /dev/null
+++ b/scripts/tests/ci/prek/test_check_first_limit.py
@@ -0,0 +1,141 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from textwrap import dedent
+
+import pytest
+from check_first_limit import check_source, main
+
+
[email protected](
+ "source",
+ [
+ "session.execute(select(Model).limit(1)).first()",
+ "session.scalars(query.limit(1)).first()",
+ "session.execute(statement=query.limit(1)).first()",
+ "session.execute(query.limit(1)).scalars().unique().first()",
+ "session.execute(query.limit(1)).scalars(0).first()",
+ "session.execute(query.limit(1)).mappings().first()",
+ "session.query(Model).limit(1).first()",
+ "query = select(Model).limit(1)\nsession.scalars(query).first()",
+ "query = select(Model)\nquery =
query.limit(1)\nsession.execute(query).first()",
+ "query: Select =
select(Model).limit(1)\nsession.execute(query).first()",
+ "query = select(Model).limit(1)\nother = query\nquery =
select(Other)\nsession.execute(other).first()",
+ "result = session.execute(query.limit(1))\nresult.first()",
+ "(result := session.execute(query.limit(1))).first()",
+ "session.execute(query.limit(1).where(Model.id ==
1).order_by(Model.id)).first()",
+
"session.execute(query.limit(1).options(joinedload(Model.other))).first()",
+ "(await session.execute(query.limit(1))).first()",
+ "(await session.stream_scalars(query.limit(1))).first()",
+ "if condition:\n query = select(Model).limit(1)\n
session.execute(query).first()",
+ "query = select(Model).limit(1)\nif condition:\n
pass\nsession.execute(query).first()",
+ "# session.execute(query).first()\ntext = 'result.first()'",
+ "obj.first",
+ "query: Select\nquery =
select(Model).limit(1)\nsession.execute(query).first()",
+ ],
+)
+def test_accepts_bounded_first(source: str):
+ assert check_source(source) == []
+
+
[email protected](
+ "source",
+ [
+ "session.execute(select(Model)).first()",
+ "session.scalars(query).first()",
+ "session.execute(statement=query).first()",
+ "session.execute(query).scalars().unique().first()",
+ "session.execute(query.limit(2)).first()",
+ "session.execute(query.limit(True)).first()",
+ "session.execute(query.limit(1.0)).first()",
+ "session.execute(query.limit(count)).first()",
+ "session.execute(query.limit(1).limit(None)).first()",
+ "session.execute(query.limit(1).fetch(10)).first()",
+
"session.execute(select(Model).where(Model.id.in_(select(Other.id).limit(1)))).first()",
+ "session.execute(select(query.limit(1).subquery())).first()",
+ "session.execute(query).limit(1).first()",
+ "session.execute(query).limit(1).limit(1).first()",
+ "query = select(Model).limit(1)\nquery =
select(Other)\nsession.execute(query).first()",
+ "query =
select(Model)\nquery.limit(1)\nsession.execute(query).first()",
+ "result = session.execute(query)\nquery =
query.limit(1)\nresult.first()",
+ "result = session.execute(query.limit(1))\nresult =
session.execute(query)\nresult.first()",
+ "session.execute(query).first()\nquery = query.limit(1)",
+ "query = select(Model).limit(1)\ndel
query\nsession.execute(query).first()",
+ "query = select(Model).limit(1)\nquery +=
other\nsession.execute(query).first()",
+ "query = select(Model).limit(1)\nquery, other =
get_queries()\nsession.execute(query).first()",
+ "if condition:\n query =
select(Model).limit(1)\nsession.execute(query).first()",
+ "if condition:\n query = select(Model).limit(1)\nelse:\n
session.execute(query).first()",
+ "query = select(Model).limit(1)\nif condition:\n query =
select(Other)\nsession.execute(query).first()",
+ "query = select(Model).limit(1)\nfor query in queries:\n
session.execute(query).first()",
+ "query = select(Model).limit(1)\nwith context() as query:\n
session.execute(query).first()",
+ "query = select(Model).limit(1)\ndef lookup(query):\n return
session.execute(query).first()",
+ "def a():\n query = select(Model).limit(1)\ndef b():\n return
session.execute(query).first()",
+ "query = select(Model).limit(1)\nlookup = lambda query:
session.execute(query).first()",
+ "try:\n query = select(Model).limit(1)\nexcept Error:\n
session.execute(query).first()",
+ "(await session.execute(query)).first()",
+ "session.execute(build_query(limit=1)).first()",
+ ],
+)
+def test_rejects_unbounded_first(source: str):
+ assert len(check_source(source)) == 1
+
+
+def test_reports_all_locations():
+ assert check_source("session.execute(query).first()\nresult.first()") == [
+ (1, 1),
+ (2, 1),
+ ]
+
+
[email protected]("invalid", [False, True])
+def test_main(tmp_path, capsys, invalid):
+ bounded = tmp_path / "bounded.py"
+ bounded.write_text("session.execute(query.limit(1)).first()",
encoding="utf-8")
+ other = tmp_path / "other.py"
+ other.write_text("session.execute(query).first()" if invalid else "pass",
encoding="utf-8")
+
+ assert main([str(bounded), str(other)]) == int(invalid)
+ output = capsys.readouterr().out
+ if invalid:
+ assert f"{other}:1:1:" in output
+ assert ".limit(1) on the statement before execution" in output
+ assert str(bounded) not in output
+ else:
+ assert output == ""
+
+
+def test_syntax_error_does_not_hide_other_failures(tmp_path, capsys):
+ broken = tmp_path / "broken.py"
+ broken.write_text("def broken(", encoding="utf-8")
+ unbounded = tmp_path / "unbounded.py"
+ unbounded.write_text("result.first()", encoding="utf-8")
+ assert main([str(broken), str(unbounded)]) == 1
+ output = capsys.readouterr().out
+ assert f"{broken}:1:" in output
+ assert f"{unbounded}:1:1:" in output
+
+
+def test_multiline_first():
+ source = dedent("""\
+ def lookup():
+ return session.execute(
+ select(Model)
+ .where(Model.id == 1)
+ ).first()
+ """)
+ assert check_source(source) == [(2, 12)]
