You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/08/26 21:57:51 UTC
(airflow) 03/03: Fix XCom deletion failure for mapped task instances (#54954)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit fc2239b75f22bd7376ab56c5b4ed7f3a391b8237
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Tue Aug 26 22:55:47 2025 +0100
Fix XCom deletion failure for mapped task instances (#54954)
XCom deletion was not working properly for mapped task instances due to
missing map_index parameter in deletion requests. This caused XCom data
to persist incorrectly when attempting to clear XComs from specific
mapped task executions.
Fixed in both `BaseXCom.delete()` and task_runner XCom clearing logic.
(cherry picked from commit c44bff4d45da96989ada6f340c2370aa14040fc7)
---
task-sdk/src/airflow/sdk/bases/xcom.py | 1 +
.../src/airflow/sdk/execution_time/task_runner.py | 1 +
task-sdk/tests/task_sdk/bases/test_xcom.py | 72 ++++++++++++++++++++++
.../task_sdk/execution_time/test_task_runner.py | 65 +++++++++++++++++++
4 files changed, 139 insertions(+)
diff --git a/task-sdk/src/airflow/sdk/bases/xcom.py b/task-sdk/src/airflow/sdk/bases/xcom.py
index 7c982a050dd..e35912d39f1 100644
--- a/task-sdk/src/airflow/sdk/bases/xcom.py
+++ b/task-sdk/src/airflow/sdk/bases/xcom.py
@@ -374,5 +374,6 @@ class BaseXCom:
dag_id=dag_id,
task_id=task_id,
run_id=run_id,
+ map_index=map_index,
),
)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index e366136f699..968d83e794e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -906,6 +906,7 @@ def run(
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=ti.run_id,
+ map_index=ti.map_index,
)
with set_current_context(context):
diff --git a/task-sdk/tests/task_sdk/bases/test_xcom.py b/task-sdk/tests/task_sdk/bases/test_xcom.py
new file mode 100644
index 00000000000..5bc2fb85b25
--- /dev/null
+++ b/task-sdk/tests/task_sdk/bases/test_xcom.py
@@ -0,0 +1,72 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.execution_time.comms import DeleteXCom, XComResult
+
+
+class TestBaseXCom:
+ @pytest.mark.parametrize(
+ "map_index",
+ [
+ pytest.param(None, id="map_index_none"),
+ pytest.param(-1, id="map_index_negative_one"),
+ pytest.param(0, id="map_index_zero"),
+ pytest.param(5, id="map_index_positive"),
+ ],
+ )
+ def test_delete_includes_map_index_in_delete_xcom_message(self, map_index, mock_supervisor_comms):
+ """Test that BaseXCom.delete properly passes map_index to the DeleteXCom message."""
+ with mock.patch.object(
+ BaseXCom, "_get_xcom_db_ref", return_value=XComResult(key="test_key", value="test_value")
+ ) as mock_get_ref:
+ with mock.patch.object(BaseXCom, "purge") as mock_purge:
+ BaseXCom.delete(
+ key="test_key",
+ task_id="test_task",
+ dag_id="test_dag",
+ run_id="test_run",
+ map_index=map_index,
+ )
+
+ mock_get_ref.assert_called_once_with(
+ key="test_key",
+ dag_id="test_dag",
+ task_id="test_task",
+ run_id="test_run",
+ map_index=map_index,
+ )
+
+ # Verify purge was called
+ mock_purge.assert_called_once()
+
+ # Verify DeleteXCom message was sent with map_index
+ mock_supervisor_comms.send.assert_called_once()
+ sent_message = mock_supervisor_comms.send.call_args[0][0]
+
+ assert isinstance(sent_message, DeleteXCom)
+ assert sent_message.key == "test_key"
+ assert sent_message.dag_id == "test_dag"
+ assert sent_message.task_id == "test_task"
+ assert sent_message.run_id == "test_run"
+ assert sent_message.map_index == map_index
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 48550de9fc6..0b7f157ee52 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -54,6 +54,7 @@ from airflow.sdk.api.datamodels._generated import (
DagRunState,
TaskInstance,
TaskInstanceState,
+ TIRunContext,
)
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
@@ -1908,6 +1909,70 @@ class TestRuntimeTaskInstance:
assert dr.run_id == "prev_success_run"
assert dr.state == "success"
+ @pytest.mark.parametrize(
+ "map_index",
+ [
+ pytest.param(-1, id="map_index_negative_one"),
+ pytest.param(0, id="map_index_zero"),
+ pytest.param(5, id="map_index_positive"),
+ ],
+ )
+ def test_xcom_clearing_includes_map_index(self, create_runtime_ti, mock_supervisor_comms, map_index):
+ """Test that XCom clearing during task execution properly passes map_index."""
+
+ task = BaseOperator(task_id="test_task")
+ runtime_ti = create_runtime_ti(task=task, map_index=map_index)
+
+ # Mock the ti_context_from_server to include xcom_keys_to_clear
+ runtime_ti._ti_context_from_server = TIRunContext(
+ dag_run=runtime_ti._ti_context_from_server.dag_run,
+ task_reschedule_count=0,
+ max_tries=1,
+ should_retry=False,
+ xcom_keys_to_clear=["key1", "key2", "key3"],
+ )
+
+ with mock.patch.object(XCom, "delete") as mock_delete:
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ # Verify that XCom.delete was called for each key with the correct map_index
+ expected_calls = [
+ mock.call(
+ key="key1",
+ dag_id=runtime_ti.dag_id,
+ task_id=runtime_ti.task_id,
+ run_id=runtime_ti.run_id,
+ map_index=map_index,
+ ),
+ mock.call(
+ key="key2",
+ dag_id=runtime_ti.dag_id,
+ task_id=runtime_ti.task_id,
+ run_id=runtime_ti.run_id,
+ map_index=map_index,
+ ),
+ mock.call(
+ key="key3",
+ dag_id=runtime_ti.dag_id,
+ task_id=runtime_ti.task_id,
+ run_id=runtime_ti.run_id,
+ map_index=map_index,
+ ),
+ ]
+ mock_delete.assert_has_calls(expected_calls)
+
+ def test_xcom_clearing_without_keys_to_clear(self, create_runtime_ti, mock_supervisor_comms):
+ """Test that no XCom clearing occurs when xcom_keys_to_clear is empty."""
+ task = BaseOperator(task_id="test_task")
+ runtime_ti = create_runtime_ti(task=task)
+
+ assert runtime_ti._ti_context_from_server.xcom_keys_to_clear is None
+
+ with mock.patch.object(XCom, "delete") as mock_delete:
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ mock_delete.assert_not_called()
+
class TestXComAfterTaskExecution:
@pytest.mark.parametrize(