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 <[email protected]> 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(
