uranusjr commented on code in PR #51568:
URL: https://github.com/apache/airflow/pull/51568#discussion_r2139467069
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -274,6 +281,61 @@ def get_one(
)
return None
+ @classmethod
+ def get_all(
+ cls,
+ *,
+ key: str,
+ dag_id: str,
+ task_id: str,
+ run_id: str,
+ ) -> Any | None:
+ """
+ Retrieve all XCom values for a task, typically from all map indexes.
+
+ This method returns "full" XCom values (i.e. uses ``deserialize_value``
+ from the XCom backend).
+
+ If there are no results, *None* is returned. If XCom entries exist,
+ a list containing all matching XCom values is returned.
+
+ This is particularly useful for getting all XCom values from all map
+ indexes of a mapped task at once.
+
+ :param key: A key for the XCom. Only XComs with this key will be
returned.
+ :param run_id: DAG run ID for the task.
+ :param dag_id: DAG ID to pull XComs from.
+ :param task_id: Task ID to pull XComs from.
+ :return: List of all XCom values if found, None if no XComs exist.
+ """
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ # Since Triggers can hit this code path via `sync_to_async` (which
uses threads internally)
+ # we need to make sure that we "atomically" send a request and get the
response to that
+ # back so that two triggers don't end up interleaving requests and
create a possible
+ # race condition where the wrong trigger reads the response.
+ with SUPERVISOR_COMMS.lock:
+ SUPERVISOR_COMMS.send_request(
+ log=log,
+ msg=GetXComSequenceSlice(
+ key=key,
+ dag_id=dag_id,
+ task_id=task_id,
+ run_id=run_id,
+ start=None,
+ stop=None,
+ step=None,
+ ),
+ )
+ msg = SUPERVISOR_COMMS.get_message()
+
+ if not isinstance(msg, XComSequenceSliceResult):
+ raise TypeError(f"Expected XComSequenceSliceResult, received:
{type(msg)} {msg}")
+
+ if msg.root is not None:
+ return msg.root
+ return None
Review Comment:
XComSequenceSliceResult should never contain None, `root` would always be a
list. It returns an empty list if there are no values (same behaviour as
slicing a Python list). This should probably simply `return msg.root`
unconditionally.
--
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]
