ferruzzi commented on code in PR #74779:
URL: https://github.com/apache/airflow/pull/74779#discussion_r4352712135
##########
task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py:
##########
@@ -78,14 +80,45 @@ class _BundleInfoLike(Protocol):
# The set of messages that a callback subprocess can send to the supervisor.
-# This is a minimal subset of ToSupervisor: read-only access to Connections
-# and Variables, plus MaskSecret for the secrets masker.
+# This is a minimal subset of ToSupervisor: read-only access to Connections,
+# Variables, and XCom values, plus MaskSecret for the secrets masker.
CallbackToSupervisor = Annotated[
- GetConnection | GetVariable | GetVariableKeys | MaskSecret,
+ GetConnection | GetVariable | GetVariableKeys | GetXCom | MaskSecret,
Field(discriminator="type"),
]
+def _add_xcom_pull_to_context(callback_kwargs: dict) -> None:
+ """
+ Add an ``xcom_pull`` helper bound to the callback's Dag run to its context.
+
+ Callbacks have no task instance, so the helper takes ``dag_id`` and
``dag_run_id`` from the
+ ``dag_run`` in the context and the callback only names the task to pull
from.
+ """
+ context = callback_kwargs.get("context")
+ if not isinstance(context, dict):
+ return
+ dag_run = context.get("dag_run")
+ if not isinstance(dag_run, dict) or not dag_run.get("dag_id") or not
dag_run.get("dag_run_id"):
+ return
+
+ def xcom_pull(
+ task_ids: str, *, key: str = "return_value", map_indexes: int | None =
None, default: Any = None
Review Comment:
Ah, yeah, that makes sense. I can see both sides here. Judgement call,
I'll leave it up to you.
--
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]
