ashb commented on code in PR #73005:
URL: https://github.com/apache/airflow/pull/73005#discussion_r4005170311
##########
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:
API updated to reject mapped_length for other keys.
##########
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:
Fixed.
--
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]
