amoghrajesh commented on code in PR #51568:
URL: https://github.com/apache/airflow/pull/51568#discussion_r2139494634
##########
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:
Makes sense, let me do that
--
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]
