This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new d9391a701e7 Restore LIMIT 1 on single-row XCom lookups (#72554)
(#72702)
d9391a701e7 is described below
commit d9391a701e74189c2a74340b9aaf91ee43f47435
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Tue Sep 8 15:46:57 2026 +0200
Restore LIMIT 1 on single-row XCom lookups (#72554) (#72702)
Signed-off-by: PoAn Yang <[email protected]>
(cherry picked from commit 39a5b937bdc669816775d31f083ed9ab6e381fe7)
Co-authored-by: PoAn Yang <[email protected]>
---
.../api_fastapi/core_api/routes/public/xcom.py | 2 +-
.../api_fastapi/execution_api/routes/xcoms.py | 4 +-
airflow-core/src/airflow/models/taskinstance.py | 2 +-
.../core_api/routes/public/test_xcom.py | 16 ++++-
.../execution_api/versions/head/test_xcoms.py | 77 ++++++++++++++++++++++
.../tests/unit/models/test_taskinstance.py | 24 ++++++-
.../src/tests_common/test_utils/asserts.py | 40 ++++++++++-
7 files changed, 158 insertions(+), 7 deletions(-)
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 40bb97cfc9f..99c5f5c1741 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
@@ -282,7 +282,7 @@ def create_xcom_entry(
run_id=dag_run_id,
map_indexes=request_body.map_index,
)
- result =
session.execute(already_existing_query.with_only_columns(XComModel.value)).first()
+ result =
session.execute(already_existing_query.with_only_columns(XComModel.value).limit(1)).first()
if result:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
index 0cb6ccb23ce..2e1736a5290 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
@@ -158,7 +158,7 @@ def get_mapped_xcom_by_index(
xcom_query = xcom_query.order_by(XComModel.map_index.desc()).offset(-1
- offset)
result: tuple[XComModel] | None
- if (result := session.scalars(xcom_query).first()) is None:
+ if (result := session.scalars(xcom_query.limit(1)).first()) is None:
message = (
f"XCom with {key=} {offset=} not found for task {task_id!r} in DAG
run {run_id!r} of {dag_id!r}"
)
@@ -337,7 +337,7 @@ def get_xcom(
# (which automatically deserializes using the backend), we avoid potential
# performance hits from retrieving large data files into the API server.
result: tuple[XComModel] | None
- if (result := session.scalars(xcom_query).first()) is None:
+ if (result := session.scalars(xcom_query.limit(1)).first()) is None:
if params.offset is None:
message = (
f"XCom with {key=} map_index={params.map_index} not found for "
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index aca45da6391..cd5ab584587 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -2063,7 +2063,7 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
XComModel.dag_id,
XComModel.map_index,
XComModel.value,
- )
+ ).limit(1)
).first()
if first is None: # No matching XCom at all.
return default
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
index 64ce784775c..eec95b3b9c2 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import json
+import re
from typing import TYPE_CHECKING
from unittest import mock
@@ -36,7 +37,7 @@ from airflow.sdk.execution_time.xcom import
resolve_xcom_backend
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.types import DagRunType
-from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.asserts import assert_queries_count,
capture_orm_selects
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags,
clear_db_runs, clear_db_xcom
@@ -696,6 +697,19 @@ class TestCreateXComEntry(TestXComEndpoint):
assert current_data["map_index"] == request_body.map_index
check_last_log(session, dag_id=TEST_DAG_ID, event="create_xcom_entry",
logical_date=None)
+ def test_create_xcom_entry_duplicate_check_is_bounded(self, test_client):
+ """Checking for an existing XCom before inserting must ask the
database for one row."""
+ with capture_orm_selects("xcom") as statements:
+ response = test_client.post(
+
f"/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries",
+ json=XComCreateBody(key=TEST_XCOM_KEY,
value=TEST_XCOM_VALUE).model_dump(),
+ )
+
+ assert response.status_code == 201
+ assert statements, "expected the endpoint to query the xcom table"
+ for sql in statements:
+ assert re.search(r"\bLIMIT 1\b", sql), f"XCom lookup is not
bounded to one row: {sql}"
+
def test_should_respond_401(self, unauthenticated_test_client):
response = unauthenticated_test_client.post(
"/dags/dag_id/dagRuns/dag_run_id/taskInstances/task_id/xcomEntries",
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
index 277e2a78fd8..cce53e69f35 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import logging
+import re
import urllib.parse
from uuid import uuid4
@@ -37,6 +38,7 @@ from airflow.serialization.serde import deserialize, serialize
from airflow.utils.session import create_session
from airflow.utils.state import DagRunState
+from tests_common.test_utils.asserts import capture_orm_selects
from tests_common.test_utils.config import conf_vars
pytestmark = pytest.mark.db_test
@@ -328,6 +330,81 @@ class TestXComsGetEndpoint:
assert set(response.json()) == set(expected_xcoms)
+ @pytest.mark.parametrize("offset", [0, 2, -1])
+ def test_xcom_get_by_index_query_is_bounded(self, client, dag_maker,
session, offset):
+ """Reading one item of a mapped XCom must ask the database for one
row, not every row after ``offset``."""
+ xcom_values = ["f", "o", "o", "b"]
+
+ class MyOperator(EmptyOperator):
+ def __init__(self, *, x, **kwargs):
+ super().__init__(**kwargs)
+ self.x = x
+
+ with dag_maker(dag_id="dag"):
+ MyOperator.partial(task_id="task").expand(x=xcom_values)
+ dag_run = dag_maker.create_dagrun(run_id="runid")
+ for ti in dag_run.task_instances:
+ session.add(
+ XComModel(
+ key="xcom_1",
+ value=xcom_values[ti.map_index],
+ dag_run_id=ti.dag_run.id,
+ run_id=ti.run_id,
+ task_id=ti.task_id,
+ dag_id=ti.dag_id,
+ map_index=ti.map_index,
+ )
+ )
+ session.commit()
+
+ with capture_orm_selects("xcom") as statements:
+ response =
client.get(f"/execution/xcoms/dag/runid/task/xcom_1/item/{offset}")
+
+ assert response.status_code == 200
+ assert response.json() == xcom_values[offset]
+ assert statements, "expected the endpoint to query the xcom table"
+ for sql in statements:
+ assert re.search(r"\bLIMIT 1\b", sql), f"XCom lookup is not
bounded to one row: {sql}"
+
+ @pytest.mark.parametrize(
+ "query_string",
+ [
+ pytest.param("", id="map_index"),
+ pytest.param("?include_prior_dates=true",
id="include_prior_dates"),
+ pytest.param("?offset=0", id="offset"),
+ ],
+ )
+ def test_xcom_get_query_is_bounded(self, client, dag_maker, session,
query_string):
+ """Reading one XCom value must ask the database for one row, however
many rows match the filters."""
+ with dag_maker(dag_id="dag"):
+ EmptyOperator(task_id="task")
+
+ for run_id, logical_date, value in [
+ ("earlier_run", "2024-01-01T00:00:00Z", "earlier_value"),
+ ("later_run", "2024-01-02T00:00:00Z", "later_value"),
+ ]:
+ dag_run = dag_maker.create_dagrun(run_id=run_id,
logical_date=timezone.parse(logical_date))
+ session.add(
+ XComModel(
+ key="xcom_1",
+ value=value,
+ dag_run_id=dag_run.id,
+ run_id=run_id,
+ task_id="task",
+ dag_id="dag",
+ )
+ )
+ session.commit()
+
+ with capture_orm_selects("xcom") as statements:
+ response =
client.get(f"/execution/xcoms/dag/later_run/task/xcom_1{query_string}")
+
+ assert response.status_code == 200
+ assert response.json() == {"key": "xcom_1", "value": "later_value"}
+ assert statements, "expected the endpoint to query the xcom table"
+ for sql in statements:
+ assert re.search(r"\bLIMIT 1\b", sql), f"XCom lookup is not
bounded to one row: {sql}"
+
class TestXComsSetEndpoint:
@pytest.mark.parametrize(
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index d02fcc593b1..ef89099b79e 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -23,6 +23,7 @@ import json
import operator
import os
import pathlib
+import re
from typing import TYPE_CHECKING, cast
from unittest import mock
from unittest.mock import patch
@@ -113,7 +114,7 @@ from airflow.utils.state import DagRunState, State,
TaskInstanceState
from airflow.utils.types import DagRunTriggeredByType, DagRunType
from tests_common.test_utils import db
-from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.asserts import assert_queries_count,
capture_orm_selects
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import clear_db_runs
from tests_common.test_utils.mock_operators import MockOperator
@@ -3399,6 +3400,27 @@ class TestMappedTaskInstanceReceiveValue:
assert isinstance(result, dict), f"Expected dict for unmapped task,
got {type(result)}"
assert result == {"key": "value"}
+ def test_xcom_pull_single_value_query_is_bounded(self, dag_maker, session):
+ """Pulling one value from one task must ask the database for one
row."""
+ with dag_maker(dag_id="test_xcom_pull_bounded", session=session):
+ upstream = PythonOperator(task_id="unmapped_task",
python_callable=lambda: {"key": "value"})
+ downstream = PythonOperator(task_id="downstream",
python_callable=lambda: None)
+ upstream >> downstream
+
+ dag_run = dag_maker.create_dagrun(logical_date=timezone.utcnow())
+ dag_maker.run_ti("unmapped_task", dag_run=dag_run, session=session)
+
+ ti_downstream = dag_run.get_task_instance("downstream",
session=session)
+ ti_downstream.task = dag_maker.dag.task_dict["downstream"]
+
+ with capture_orm_selects("xcom") as statements:
+ result = ti_downstream.xcom_pull(task_ids="unmapped_task",
session=session)
+
+ assert result == {"key": "value"}
+ assert statements, "expected xcom_pull to query the xcom table"
+ for sql in statements:
+ assert re.search(r"\bLIMIT 1\b", sql), f"xcom_pull is not bounded
to one row: {sql}"
+
def test_xcom_pull_returns_lazy_sequence_for_mapped_xcom(self, dag_maker,
session):
"""
Test that xcom_pull returns LazyXComSelectSequence when XComs are
mapped (map_index >= 0)
diff --git a/devel-common/src/tests_common/test_utils/asserts.py
b/devel-common/src/tests_common/test_utils/asserts.py
index 59ca9bad968..4ad1a2637f0 100644
--- a/devel-common/src/tests_common/test_utils/asserts.py
+++ b/devel-common/src/tests_common/test_utils/asserts.py
@@ -25,12 +25,16 @@ from contextlib import contextmanager
from typing import TYPE_CHECKING, NamedTuple
from sqlalchemy import event
+from sqlalchemy.orm import Session
+from sqlalchemy.sql import Select
# Long import to not create a copy of the reference, but to refer to one place.
import airflow.settings
if TYPE_CHECKING:
- from sqlalchemy.orm.session import Session
+ from collections.abc import Generator
+
+ from sqlalchemy.orm import ORMExecuteState
log = logging.getLogger(__name__)
@@ -177,3 +181,37 @@ def assert_queries_count(
message += f"\n\t{location}:\t{count}"
raise AssertionError(message)
+
+
+@contextmanager
+def capture_orm_selects(table: str) -> Generator[list[str], None, None]:
+ """
+ Collect the ORM ``SELECT`` statements issued against ``table`` while the
context is active.
+
+ Each statement is rendered with SQLAlchemy's default dialect and with its
bound values inlined,
+ so assertions about the shape of a query (``LIMIT 1``, ``OFFSET`` ...)
read the same on every
+ backend. The raw text seen by ``before_cursor_execute`` is not enough for
``LIMIT``: SQLite,
+ Postgres and MySQL all emit a ``LIMIT`` of some sort for an offset-only
query.
+
+ The listener is attached to the ``Session`` class, so statements executed
by sessions the code
+ under test opens itself (for example inside an API request handler) are
captured too.
+
+ :param table: Name of the table the captured statements must select from.
+ """
+ statements: list[str] = []
+ selects_from_table = re.compile(rf"\bFROM {re.escape(table)}\b")
+
+ def capture(orm_execute_state: ORMExecuteState) -> None:
+ statement = orm_execute_state.statement
+ if not isinstance(statement, Select):
+ return
+ if not selects_from_table.search(" ".join(str(statement).split())):
+ return
+ rendered = str(statement.compile(compile_kwargs={"literal_binds":
True}))
+ statements.append(" ".join(rendered.split()))
+
+ event.listen(Session, "do_orm_execute", capture)
+ try:
+ yield statements
+ finally:
+ event.remove(Session, "do_orm_execute", capture)
