kaxil commented on code in PR #73005:
URL: https://github.com/apache/airflow/pull/73005#discussion_r3992952551
##########
airflow-core/src/airflow/models/xcom.py:
##########
@@ -245,6 +253,7 @@ def set(
dag_id=dag_id,
map_index=map_index,
dag_result=dag_result,
+ mapped_length=mapped_length,
Review Comment:
`XComModel.set` is delete-then-insert, so any caller that doesn't pass
`mapped_length` wipes it, and `update_xcom_entry` in
`core_api/routes/public/xcom.py` is one of those. A `PATCH
.../xcomEntries/return_value` re-inserts the row with `mapped_length=None`,
`get_task_map_length` then reads NULL and the next expansion raises
`NotFullyPopulated`. If the downstream hasn't expanded yet the unmapped TI gets
marked `UPSTREAM_FAILED`; if it already has, `verify_integrity` marks the
existing mapped TIs `REMOVED`.
Reachable from the UI too, since `XCom.tsx` puts `EditXComButton` on every
row including `return_value`, with nothing gating it beyond normal DAG edit
access.
Pre-PR this couldn't happen: the only thing writing `task_map` was the exec
API route, `public/xcom.py` isn't in this diff, and the old `XComModel.set` had
no `mapped_length` parameter at all, so editing a value left expansion alone.
The fix looks simple, you've already got the row loaded nine lines up, so
`mapped_length=xcom_entry.mapped_length` would carry it through (or recompute
it from the new value).
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -397,7 +396,7 @@ def set_xcom(
map_index: Annotated[int, Query()] = -1,
dag_result: Annotated[bool, Query(description="Whether this XCom is a dag
result")] = False,
mapped_length: Annotated[
- int | None, Query(description="Number of mapped tasks this value
expands into")
+ int | None, Query(ge=0, description="Number of mapped tasks this value
expands into")
Review Comment:
Not blocking, more a question. The read side is pinned to `key ==
XCOM_RETURN_KEY` now, but this still accepts `mapped_length` on any key, so a
write under any other key gets stored and then never looked at.
`test_xcom_set_mapped` in this PR does exactly that, posting `mapped_length=3`
under `xcom_1` and asserting it sticks. With `task_map` it didn't matter, since
the row had no key column at all.
Worth rejecting `mapped_length` unless `key == XCOM_RETURN_KEY`, so the
invariant is enforced rather than something the SDK just happens to respect?
##########
airflow-core/src/airflow/migrations/versions/0134_3_4_0_fold_task_map_into_xcom_mapped_length.py:
##########
@@ -0,0 +1,144 @@
+#
+# 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.
+
+"""
+Fold task_map into xcom.mapped_length.
+
+Every task_map row is written by the same execution API call that writes the
+pushing task's ``return_value`` XCom row, at the same coordinates, so the
length
+lives on that row instead. ``task_map.keys`` is dropped rather than migrated:
the
+only writer has always set it to NULL, so downgrade restores every map as the
+list variant.
+
+Revision ID: 3b7a91c5df20
+Revises: f8c2a1d94e03
+Create Date: 2026-09-10 10:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+from airflow.migrations.db_types import StringID
+from airflow.migrations.utils import disable_sqlite_fkeys
+from airflow.utils.sqlalchemy import ExtendedJSON
+
+# revision identifiers, used by Alembic.
+revision = "3b7a91c5df20"
+down_revision = "f8c2a1d94e03"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+XCOM_RETURN_KEY = "return_value"
+
+_xcom = sa.table(
+ "xcom",
+ sa.column("dag_id"),
+ sa.column("task_id"),
+ sa.column("run_id"),
+ sa.column("map_index"),
+ sa.column("key"),
+ sa.column("mapped_length"),
+)
+_task_map = sa.table(
+ "task_map",
+ sa.column("dag_id"),
+ sa.column("task_id"),
+ sa.column("run_id"),
+ sa.column("map_index"),
+ sa.column("length"),
+ sa.column("keys"),
+)
+
+_JOIN = sa.and_(
+ _task_map.c.dag_id == _xcom.c.dag_id,
+ _task_map.c.task_id == _xcom.c.task_id,
+ _task_map.c.run_id == _xcom.c.run_id,
+ _task_map.c.map_index == _xcom.c.map_index,
+)
+
+# A correlated subquery rather than UPDATE ... FROM: MySQL has no such form
and SQLite only
+# gained it in 3.33. The EXISTS keeps the statement a no-op for rows with no
task_map row.
+BACKFILL = (
+ _xcom.update()
+ .where(
+ _xcom.c.key == XCOM_RETURN_KEY,
+ sa.exists(sa.select(sa.literal(1)).where(_JOIN)),
+ )
+
.values(mapped_length=sa.select(_task_map.c.length).where(_JOIN).scalar_subquery())
+)
+
+RESTORE = _task_map.insert().from_select(
+ ["dag_id", "task_id", "run_id", "map_index", "length", "keys"],
+ sa.select(
+ _xcom.c.dag_id,
+ _xcom.c.task_id,
+ _xcom.c.run_id,
+ _xcom.c.map_index,
+ _xcom.c.mapped_length,
+ sa.null(),
+ ).where(_xcom.c.mapped_length.is_not(None)),
Review Comment:
`RESTORE` isn't filtered to `return_value`, but `task_map`'s PK is only
(dag_id, task_id, run_id, map_index), so two xcom rows at the same coordinates
under different keys, both carrying a `mapped_length`, collide. I ran the real
`RESTORE` against `task_map` as `downgrade()` creates it, with such a pair
present, and it fails on all three backends:
- sqlite: `UNIQUE constraint failed: task_map.dag_id, task_map.task_id,
task_map.run_id, task_map.map_index`
- postgres: `duplicate key value violates unique constraint "task_map_pkey"`
- mysql: `Duplicate entry 'dag-op1-test--1' for key 'task_map.PRIMARY'`
Adding `_xcom.c.key == XCOM_RETURN_KEY` to the where clause clears all three.
The shipped SDK only ever sends `mapped_length` on `return_value`, so a real
worker won't produce that pair today. But nothing server-side stops it either:
the set route takes `mapped_length` on any key, and `XComModel.set`'s delete is
scoped to the single key, so a second key at the same coordinates keeps its own
value. And `BACKFILL` above pins to `return_value`, as does
`get_task_map_length` with a comment saying why, which leaves `RESTORE` as the
odd one out of the three. Pinning it also means the downgrade reads back
exactly what `BACKFILL` wrote.
##########
airflow-core/src/airflow/migrations/versions/0134_3_4_0_fold_task_map_into_xcom_mapped_length.py:
##########
@@ -0,0 +1,144 @@
+#
+# 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.
+
+"""
+Fold task_map into xcom.mapped_length.
+
+Every task_map row is written by the same execution API call that writes the
+pushing task's ``return_value`` XCom row, at the same coordinates, so the
length
+lives on that row instead. ``task_map.keys`` is dropped rather than migrated:
the
Review Comment:
The conclusion is right but the premise isn't, and the PR body puts it more
strongly still ("The `keys` column has never been written to"), so worth a
tweak.
`keys` was written non-NULL for the whole of 2.x, on the normal worker path
rather than in a corner. `_record_task_map_for_downstreams` was called from
`_execute_task` (`airflow/models/taskinstance.py:801` at 2.11.0, `:1667` at
2.3.0) via `_run_raw_task`, and went through `TaskMap.from_task_instance_xcom`,
which does `keys=(list(value) if isinstance(value, collections.abc.Mapping)
else None)` (`airflow/models/taskmap.py:109` at 2.11.0). Dicts reach it,
`_is_mappable_value` allows them (`taskinstance.py:675`). So any 2.x deployment
where a task returned a dict feeding a mapped downstream has non-NULL `keys` in
its DB, and nothing in the 2->3 chain clears those rows
(`0042_3_0_0_add_uuid_primary_key_to_task_instance_.py` just recreates the FK).
That writer survived into 3.0.0 and 3.0.1 on the `airflow dags test`, `airflow
tasks test` and `BaseOperator.run()` paths (`models/taskinstance.py:346`, call
site `:2220` at 3.0.0 and `:2219` at 3.0.1), and went in 3.0.2 with #50980.
Dropping the column is still fine, just for a different reason:
`TaskMapVariant` has never appeared outside `models/taskmap.py` at any tag from
2.3.0 to 3.3.0, and `TaskMap.keys` is never selected anywhere. "Nothing has
ever read it" is the part that holds.
The first sentence is true at write time, but the rows come apart
afterwards, so "downgrade restores every map" isn't quite right over a row's
lifetime. The per-attempt purge only deletes XCom rows:
`TIRunContext.xcom_keys_to_clear`
(`execution_api/datamodels/taskinstance.py:428`, filled at
`routes/task_instances.py:283-309`) goes to the SDK, which calls `XCom.delete`
per key (`task_runner.py:1571`), and the exec API DELETE route is a bare
`delete(XComModel)` (`routes/xcoms.py:473`), same for the public one
(`core_api/routes/public/xcom.py:435`). A cleared or retried TI that purges and
then fails before re-pushing leaves a task_map row with no `return_value` row,
and `BACKFILL`'s `EXISTS` drops that length.
`0049_3_0_0_remove_pickled_data_from_xcom_table.py:290` creates the same orphan
outright, deleting pickled xcom rows during the 3.0.0 upgrade without touching
`task_map`. That's what `test_mapped_length_dies_with_the_pushed_value` is
getting at anyway, so no objection to the be
haviour, just worth saying here.
##########
devel-common/src/tests_common/test_utils/mapping.py:
##########
@@ -18,31 +18,68 @@
from typing import TYPE_CHECKING
-from airflow.models.taskmap import TaskMap
+from sqlalchemy import select
+
+from airflow.models.taskinstance import TaskInstance
+from airflow.models.xcom import XCOM_RETURN_KEY, XComModel
if TYPE_CHECKING:
+ from collections.abc import Collection, Sequence
+
from sqlalchemy.orm import Session
from airflow.serialization.definitions.mappedoperator import Operator
+def push_mapped_length(ti: TaskInstance, value: Collection, *, session:
Session) -> None:
+ """Record ``value`` as ``ti``'s return value, usable as an expansion
input."""
+ XComModel.set(
+ key=XCOM_RETURN_KEY,
+ value=list(value),
Review Comment:
Minor: `list(value)` turns a dict return into its key list, so this records
something the real push path never would (`task_runner.py` pushes `result`
unchanged alongside `mapped_length`). `test_map_product_expansion`'s
`emit_letters` returns `{"a": "x", "b": "y", "c": "z"}` and this stores `["a",
"b", "c"]`.
Nothing fails today since expansion only reads `mapped_length`, so it's a
trap for the next test rather than a bug: the two shapes aren't interchangeable
once anyone asserts on a received value, because `_expand_mapped_field` gives
`value[i]` for a sequence and a `(key, value)` tuple for a dict. The coercion
looks like it's only there for the two callers passing a `range`,
`expand_mapped_task` below and `range(count)` in
`core_api/routes/public/test_task_instances.py`, so could those pass
`list(range(...))` and leave `value` alone here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
