amoghrajesh commented on code in PR #49692:
URL: https://github.com/apache/airflow/pull/49692#discussion_r2060278920
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -320,10 +320,29 @@ def xcom_pull(
a non-str iterable), a list of matching XComs is returned. Elements in
the list is ordered by item ordering in ``task_id`` and ``map_index``.
"""
- if dag_id is None:
- dag_id = self.dag_id
- if run_id is None:
- run_id = self.run_id
+ dag_id = dag_id or self.dag_id
+ run_id = run_id or self.run_id
+
+ is_single_task = task_ids is None or isinstance(task_ids, str)
+ is_single_map_index = not isinstance(map_indexes, Iterable)
+
+ # We are pulling one single task or handling case where single map
index is passed.
+ if is_single_task and is_single_map_index:
+ t_id = self.task_id if task_ids is None else task_ids
+ m_id = (
+ self.map_index if isinstance(map_indexes, ArgNotSet) or
map_indexes is None else map_indexes
+ )
+ value = XCom.get_one(
+ run_id=run_id,
+ key=key,
+ task_id=t_id,
+ dag_id=dag_id,
+ map_index=m_id,
+ include_prior_dates=include_prior_dates,
+ )
+ if value is None:
+ return default
+ return value
Review Comment:
Didnt want to perform a heavy refactoring effort, so tried to keep it to a
minimum and somethign that gets things working.
This can surely be reworked if needed, but i think its not needed.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -363,8 +363,6 @@ def xcom_pull(
else:
xcoms.append(value)
- if len(xcoms) == 1:
- return xcoms[0]
Review Comment:
Thanks, i have handled it
--
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]
