diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py index 9ab8eb200712a..6c8b4ad6d7468 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py @@ -20,7 +20,7 @@ import logging from typing import Annotated -from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status +from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, Security, status from pydantic import JsonValue from sqlalchemy import delete from sqlalchemy.sql.selectable import Select @@ -32,7 +32,7 @@ XComSequenceIndexResponse, XComSequenceSliceResponse, ) -from airflow.api_fastapi.execution_api.security import CurrentTIToken +from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute, require_auth from airflow.models.taskmap import TaskMap from airflow.models.xcom import XComModel from airflow.utils.db import get_query_count @@ -105,6 +105,7 @@ def has_xcom_access( router = APIRouter( + route_class=ExecutionAPIRoute, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the XCom"}, @@ -135,6 +136,7 @@ async def xcom_query( @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}/item/{offset}", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Get a single XCom value from a mapped task by sequence index", ) def get_mapped_xcom_by_index( @@ -180,6 +182,7 @@ class GetXComSliceFilterParams(BaseModel): @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}/slice", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Get XCom values from a mapped task by sequence slice", ) def get_mapped_xcom_by_slice( @@ -273,6 +276,7 @@ def get_mapped_xcom_by_slice( }, }, }, + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Returns the count of mapped XCom values found in the `Content-Range` response header", ) def head_xcom( @@ -305,6 +309,7 @@ class GetXcomFilterParams(BaseModel): @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}", description="Get a single XCom Value", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], ) def get_xcom( dag_id: str, diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 6912d02d1be0f..ca597b17ebf09 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -4050,7 +4050,14 @@ def test_workload_scope_rejected_on_state_endpoint(self, client, session, create assert "Token type 'workload' not allowed" in resp.json()["detail"] def test_workload_scope_accepted_on_connections_endpoint(self, client, session, create_task_instance): - """Workload scoped tokens are accepted on GET /connections for deadline callback subprocesses.""" + """Workload scoped tokens are accepted on GET /connections (read access for callbacks). + + The connections router declares ``token:workload`` on its read route and sets + ``route_class=ExecutionAPIRoute`` so the scope is enforced; deadline callback + subprocesses (which carry workload tokens) must be able to read connections. + A missing connection therefore returns 404 (the request reached the route), + not a 403 token-type rejection. + """ ti = create_task_instance(task_id="test_workload_conn", state=State.RUNNING) session.commit() @@ -4059,6 +4066,7 @@ def test_workload_scope_accepted_on_connections_endpoint(self, client, session, resp = client.get("/execution/connections/test_conn") # Workload tokens are now accepted; 404 because the connection doesn't exist in the test DB. assert resp.status_code == 404 + assert "Token type 'workload' not allowed" not in str(resp.json()) def test_execution_scope_accepted_on_all_endpoints(self, client, session, create_task_instance): """Execution scoped tokens should be accepted on all endpoints.""" 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 3e62a90f23539..ad819a673049e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -47,12 +47,14 @@ GetDagRun, 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 ( @@ -102,11 +104,11 @@ class CallbackContextFetchError(RuntimeError): # 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, plus GetDagRun for -# building context from DagRun identifiers. +# This is a minimal subset of ToSupervisor: read-only access to Connections, +# Variables, and XCom values, plus MaskSecret for the secrets masker, plus +# GetDagRun for building context from DagRun identifiers. CallbackToSupervisor = Annotated[ - GetConnection | GetDagRun | GetVariable | GetVariableKeys | MaskSecret, + GetConnection | GetDagRun | GetVariable | GetVariableKeys | GetXCom | MaskSecret, Field(discriminator="type"), ] @@ -220,9 +222,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 ``XCom.get()`` via the + supervisor's API client. """ client: Client # The HTTP client to use for communication with the API server. @@ -435,6 +438,8 @@ def _handle_request(self, msg: CallbackToSupervisor, log: FilteringBoundLogger, 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 647bda8ea87f8..14ee6854d42f4 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 @@ -33,7 +33,7 @@ from airflow.sdk._shared.template_rendering import render_callback_kwargs from airflow.sdk._shared.timezones import timezone -from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, DagRunType +from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, DagRunType, XComResponse from airflow.sdk.execution_time.callback_supervisor import ( CALLBACK_CONTEXT_FETCH_EXIT_CODE, CallbackContextFetchError, @@ -43,8 +43,16 @@ supervise_callback, ) from airflow.sdk.execution_time.comms import ( + ConnectionResult, ErrorResponse, + GetConnection, GetDagRun, + GetVariable, + GetVariableKeys, + GetXCom, + MaskSecret, + VariableKeysResult, + VariableResult, _RequestFrame, ) @@ -172,6 +180,70 @@ class RequestCase: response=_MOCK_DAG_RUN, ), ), + RequestCase( + message=GetConnection(conn_id="test_conn"), + test_id="get_connection_with_password", + client_mock=ClientMock( + method_path="connections.get", + args=("test_conn",), + response=ConnectionResult(conn_id="test_conn", conn_type="mysql", password="secret"), + ), + mask_secret_args=("secret",), + ), + RequestCase( + message=GetVariable(key="test_key"), + test_id="get_variable", + client_mock=ClientMock( + method_path="variables.get", + args=("test_key",), + response=VariableResult(key="test_key", value="test_value"), + ), + ), + RequestCase( + message=GetVariableKeys(prefix="test_"), + test_id="get_variable_keys", + client_mock=ClientMock( + method_path="variables.keys", + kwargs={"prefix": "test_", "limit": 1000, "offset": 0}, + 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", + map_index=None, + ), + test_id="get_xcom", + client_mock=ClientMock( + method_path="xcoms.get", + args=("test_dag", "test_run_1", "upstream_task", "return_value", None, 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), + response=XComResponse(key="custom_key", value={"nested": "data"}), + ), + ), + RequestCase( + message=MaskSecret(value="super_secret", name="api_key"), + test_id="mask_secret", + mask_secret_args=("super_secret", "api_key"), + ), ] @pytest.fixture