Skip to content

Fix xcom_pull for task_ids=None #47407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amoghrajesh merged 4 commits into from
Mar 8, 2025
Merged

Fix xcom_pull for task_ids=None #47407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
31f91aa
Fix xcom_pull for task_ids=None
insomnes Mar 5, 2025
d1d78bf
Remove print
insomnes Mar 5, 2025
687eaa4
Use task_id provided to operator, instead of one from runtime_ti 'mock'
insomnes Mar 7, 2025
ff1bbec
Merge branch 'main' into xcom-task-ids-none
amoghrajesh Mar 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion task-sdk/src/airflow/sdk/execution_time/task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ def xcom_pull(

if task_ids is None:
# default to the current task if not provided
task_ids = self.task_id
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
if isinstance(map_indexes, ArgNotSet):
Expand Down
17 changes: 14 additions & 3 deletions task-sdk/tests/task_sdk/execution_time/test_task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import json
import os
import uuid
from collections.abc import Iterable
from datetime import datetime, timedelta
from pathlib import Path
from socket import socketpair
Expand Down Expand Up @@ -93,6 +94,7 @@
)
from airflow.utils import timezone
from airflow.utils.state import TaskInstanceState
from airflow.utils.types import NOTSET, ArgNotSet

from tests_common.test_utils.mock_operators import AirflowLink

Expand Down Expand Up @@ -1134,28 +1136,37 @@ def test_get_variable_from_context(
"push_task",
["push_task1", "push_task2"],
{"push_task1", "push_task2"},
None,
NOTSET,
],
)
def test_xcom_pull(self, create_runtime_ti, mock_supervisor_comms, spy_agency, task_ids):
"""Test that a task pulls the expected XCom value if it exists."""

class CustomOperator(BaseOperator):
def execute(self, context):
value = context["ti"].xcom_pull(task_ids=task_ids, key="key")
if isinstance(task_ids, ArgNotSet):
value = context["ti"].xcom_pull(key="key")
else:
value = context["ti"].xcom_pull(task_ids=task_ids, key="key")
print(f"Pulled XCom Value: {value}")

task = CustomOperator(task_id="pull_task")
test_task_id = "pull_task"
task = CustomOperator(task_id=test_task_id)

runtime_ti = create_runtime_ti(task=task)

mock_supervisor_comms.get_message.return_value = XComResult(key="key", value='"value"')

run(runtime_ti, log=mock.MagicMock())

if isinstance(task_ids, str):
if not isinstance(task_ids, Iterable) or isinstance(task_ids, str):
task_ids = [task_ids]

for task_id in task_ids:
# Without task_ids (or None) expected behavior is to pull with calling task_id
if task_id is None or isinstance(task_id, ArgNotSet):
task_id = test_task_id
mock_supervisor_comms.send_request.assert_any_call(
log=mock.ANY,
msg=GetXCom(
Expand Down