This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9ef5f3d463f Add XCom read access to deadline callbacks run by
executors (#74779)
9ef5f3d463f is described below
commit 9ef5f3d463f030884aeb9824dacd22a4d1784789
Author: Sean Ghaeli <[email protected]>
AuthorDate: Fri Oct 23 11:53:35 2026 -0700
Add XCom read access to deadline callbacks run by executors (#74779)
Executor-run deadline callbacks could read Connections and Variables
through the callback supervisor but not XComs, so an alert could not
include an upstream task's output. Serve GetXCom through the shared
handle_get_xcom handler, and add context["xcom_pull"] bound to the
callback's Dag run so a callback only names the task to pull from.
Only reads are exposed. The request uses the execution token from the
callback token exchange, so no API routes or token scopes change.
---
airflow-core/docs/howto/deadline-alerts.rst | 3 +
.../sdk/execution_time/callback_supervisor.py | 50 ++++++++++++++--
.../execution_time/test_callback_supervisor.py | 66 ++++++++++++++++++++++
3 files changed, 113 insertions(+), 6 deletions(-)
diff --git a/airflow-core/docs/howto/deadline-alerts.rst
b/airflow-core/docs/howto/deadline-alerts.rst
index 538173c3505..714397dafc0 100644
--- a/airflow-core/docs/howto/deadline-alerts.rst
+++ b/airflow-core/docs/howto/deadline-alerts.rst
@@ -387,6 +387,9 @@ as long as the variables being templated are included in
the simplified context.
ID and the calculated deadline time of the Deadline Alert as well as the data
included in the ``GET`` REST API
response for Dag Run. Support for more comprehensive context and templating
will be added in future versions.
+Synchronous callbacks can also read XComs from tasks in the same Dag run
through ``context["xcom_pull"]``,
+for example ``context["xcom_pull"]("extract")`` or
``context["xcom_pull"]("extract", key="row_count")``.
+
Deadline Calculation
^^^^^^^^^^^^^^^^^^^^
diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
index bc38e2abb20..a2861de413c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
@@ -39,12 +39,14 @@ from airflow.sdk.execution_time.comms import (
GetConnection,
GetVariable,
GetVariableKeys,
+ GetXCom,
MaskSecret,
)
from airflow.sdk.execution_time.request_handlers import (
handle_get_connection,
handle_get_variable,
handle_get_variable_keys,
+ handle_get_xcom,
handle_mask_secret,
)
from airflow.sdk.execution_time.supervisor import (
@@ -78,14 +80,46 @@ log: FilteringBoundLogger =
structlog.get_logger(logger_name="callback_superviso
# 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"):
+ log.warning("Callback context has no Dag run, so xcom_pull is not
available")
+ return
+
+ def xcom_pull(
+ task_id: str, *, key: str = "return_value", map_index: int | None =
None, default: Any = None
+ ) -> Any:
+ from airflow.sdk.execution_time.xcom import XCom
+
+ value = XCom.get_one(
+ key=key,
+ dag_id=dag_run["dag_id"],
+ run_id=dag_run["dag_run_id"],
+ task_id=task_id,
+ map_index=map_index,
+ )
+ return default if value is None else value
+
+ context["xcom_pull"] = xcom_pull
+
+
SERDE_KWARGS_ENCODING = "serde"
@@ -209,9 +243,10 @@ class CallbackSubprocess(WatchedSubprocess):
Uses the WatchedSubprocess infrastructure for fork/monitor/signal handling
while keeping a simple lifecycle: start, run callback, exit.
- Provides a limited set of comms channels (Connections and Variables) so
- that callback code can access runtime services like
- ``Connection.get()`` and ``Variable.get()`` via the supervisor's API
client.
+ Provides a limited set of comms channels (Connections, Variables, and XCom)
+ so that callback code can access runtime services like
+ ``Connection.get()``, ``Variable.get()``, and ``context["xcom_pull"]()``
via
+ the supervisor's API client.
"""
client: Client # The HTTP client to use for communication with the API
server.
@@ -270,6 +305,7 @@ class CallbackSubprocess(WatchedSubprocess):
exc_info=True,
)
+ _add_xcom_pull_to_context(callback_kwargs)
success, error_msg = execute_callback(
callback_path=callback_path,
callback_kwargs=callback_kwargs,
@@ -391,6 +427,8 @@ class CallbackSubprocess(WatchedSubprocess):
resp, dump_opts = handle_get_variable(self.client, msg)
elif isinstance(msg, GetVariableKeys):
resp, dump_opts = handle_get_variable_keys(self.client, msg)
+ elif isinstance(msg, GetXCom):
+ resp, dump_opts = handle_get_xcom(self.client, msg)
elif isinstance(msg, MaskSecret):
handle_mask_secret(msg)
else:
diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
index 7eb5bc9c949..36c0cf52ee8 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
@@ -31,6 +31,7 @@ from unittest.mock import ANY, Mock, patch
import pytest
import structlog
+from airflow.sdk.api.datamodels._generated import XComResponse
from airflow.sdk.execution_time.callback_supervisor import (
CallbackSubprocess,
Path,
@@ -43,9 +44,11 @@ from airflow.sdk.execution_time.comms import (
GetConnection,
GetVariable,
GetVariableKeys,
+ GetXCom,
MaskSecret,
VariableKeysResult,
VariableResult,
+ XComResult,
_RequestFrame,
)
@@ -320,6 +323,38 @@ class TestCallbackHandleRequest:
response=VariableKeysResult(keys=["test_key"],
total_entries=1),
),
),
+ RequestCase(
+ message=GetXCom(
+ key="return_value",
+ dag_id="test_dag",
+ run_id="test_run_1",
+ task_id="upstream_task",
+ ),
+ test_id="get_xcom",
+ client_mock=ClientMock(
+ method_path="xcoms.get",
+ args=("test_dag", "test_run_1", "upstream_task",
"return_value", None, False),
+ kwargs={"previous_iteration": False},
+ response=XComResponse(key="return_value",
value="xcom_payload"),
+ ),
+ ),
+ RequestCase(
+ message=GetXCom(
+ key="custom_key",
+ dag_id="dag_a",
+ run_id="run_42",
+ task_id="task_b",
+ map_index=3,
+ include_prior_dates=True,
+ ),
+ test_id="get_xcom_with_map_index",
+ client_mock=ClientMock(
+ method_path="xcoms.get",
+ args=("dag_a", "run_42", "task_b", "custom_key", 3, True),
+ kwargs={"previous_iteration": False},
+ response=XComResponse(key="custom_key", value={"nested":
"data"}),
+ ),
+ ),
RequestCase(
message=MaskSecret(value="super_secret", name="api_key"),
test_id="mask_secret",
@@ -639,6 +674,37 @@ class TestCallbackSubprocessStart:
assert exc_info.value.code == 1
+ def _run_target_with_context(self, base_start_kwargs, context):
+ CallbackSubprocess.start(**{**base_start_kwargs, "callback_kwargs":
{"context": context}})
+ self.mock_super_start.call_args.kwargs["target"]()
+ return
self.mock_execute_callback.call_args.kwargs["callback_kwargs"]["context"]
+
+ @pytest.mark.parametrize(
+ ("xcom_value", "expected"),
+ [
+ pytest.param("payload", "payload", id="value"),
+ pytest.param(None, "fallback", id="missing"),
+ ],
+ )
+ def test_context_xcom_pull_uses_dag_run_ids(self, base_start_kwargs,
xcom_value, expected):
+ comms = self.mock_comms_decoder.__getitem__.return_value.return_value
+ comms.send.return_value = XComResult(key="out", value=xcom_value)
+
+ context = self._run_target_with_context(
+ base_start_kwargs, {"dag_run": {"dag_id": "dag_a", "dag_run_id":
"run_1"}}
+ )
+
+ assert context["xcom_pull"]("upstream", key="out", map_index=2,
default="fallback") == expected
+ comms.send.assert_called_once_with(
+ GetXCom(key="out", dag_id="dag_a", run_id="run_1",
task_id="upstream", map_index=2)
+ )
+
+ def test_context_without_dag_run_has_no_xcom_pull(self, base_start_kwargs,
caplog):
+ context = self._run_target_with_context(base_start_kwargs,
{"deadline": {"id": "x"}})
+
+ assert "xcom_pull" not in context
+ assert "Callback context has no Dag run, so xcom_pull is not
available" in caplog
+
@pytest.mark.parametrize("option_value", ["True", "False"])
def test_start_keeps_bare_fork_regardless_of_exec_option(self,
base_start_kwargs, option_value):
"""
