Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/execution_api/routes/xcoms.py: 41%
144 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 14:22 +0000
1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
18from __future__ import annotations
20import logging
21from typing import Annotated
23from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status
24from pydantic import JsonValue
25from sqlalchemy import delete
26from sqlalchemy.sql.selectable import Select
28from airflow.api_fastapi.common.db.common import SessionDep
29from airflow.api_fastapi.core_api.base import BaseModel
30from airflow.api_fastapi.execution_api.datamodels.xcom import (
31 XComResponse,
32 XComSequenceIndexResponse,
33 XComSequenceSliceResponse,
34)
35from airflow.api_fastapi.execution_api.security import CurrentTIToken
36from airflow.models.taskmap import TaskMap
37from airflow.models.xcom import XComModel
38from airflow.utils.db import get_query_count
41def has_xcom_access(
42 dag_id: str,
43 run_id: str,
44 task_id: str,
45 xcom_key: Annotated[str, Path(alias="key", min_length=1)],
46 request: Request,
47 session: SessionDep,
48 token=CurrentTIToken,
49) -> bool:
50 """
51 Check whether the requesting task may access the XCom for ``dag_id``.
53 In multi-team mode, XCom access is scoped by team ownership (resolved via the
54 ``dag -> bundle -> team`` chain). There is no cross-team XCom sharing:
56 * reads (``GET``/``HEAD``) are allowed for the requester's own team or for
57 global (teamless) dags;
58 * writes and deletes are allowed only for the requester's own team; a team
59 task may not mutate a global dag's XCom, mirroring how team-scoped
60 Variables and Connections behave.
62 When multi-team mode is disabled this is a no-op and all access is allowed,
63 consistent with Airflow's single-team security model where workers within a
64 deployment trust each other. Note this enforces the boundary at the Execution
65 API only; it does not constrain code paths with direct database access (e.g.
66 the Dag File Processor or Triggerer).
67 """
68 from airflow.configuration import conf
70 write = request.method not in {"GET", "HEAD", "OPTIONS"}
72 log.debug(
73 "Checking %s XCom access for task instance '%s' to XCom '%s' on dag '%s'",
74 "write" if write else "read",
75 token.id,
76 xcom_key,
77 dag_id,
78 )
80 if not conf.getboolean("core", "multi_team"): 80 ↛ 83line 80 didn't jump to line 83 because the condition on line 80 was always true
81 return True
83 from airflow.api_fastapi.execution_api.security import (
84 _team_name_for_dag_stmt,
85 _team_name_for_ti_stmt,
86 )
88 requester_team = session.scalar(_team_name_for_ti_stmt(token.id))
89 target_team = session.scalar(_team_name_for_dag_stmt(dag_id))
91 # Same team (including a teamless task accessing a global, teamless dag) is always allowed.
92 if target_team == requester_team:
93 return True
94 # Reads may additionally reach global (teamless) dags; writes and deletes may not.
95 if not write and target_team is None:
96 return True
98 raise HTTPException(
99 status_code=status.HTTP_403_FORBIDDEN,
100 detail={
101 "reason": "access_denied",
102 "message": "Task does not have access to this XCom in multi-team mode",
103 },
104 )
107router = APIRouter(
108 responses={
109 status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
110 status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the XCom"},
111 status.HTTP_404_NOT_FOUND: {"description": "XCom not found"},
112 },
113 dependencies=[Depends(has_xcom_access)],
114)
116log = logging.getLogger(__name__)
119async def xcom_query(
120 dag_id: str,
121 run_id: str,
122 task_id: str,
123 key: str,
124 map_index: Annotated[int | None, Query()] = None,
125) -> Select:
126 query = XComModel.get_many(
127 run_id=run_id,
128 key=key,
129 task_ids=task_id,
130 dag_ids=dag_id,
131 map_indexes=map_index,
132 )
133 return query
136@router.get(
137 "/{dag_id}/{run_id}/{task_id}/{key:path}/item/{offset}",
138 description="Get a single XCom value from a mapped task by sequence index",
139)
140def get_mapped_xcom_by_index(
141 dag_id: str,
142 run_id: str,
143 task_id: str,
144 key: Annotated[str, Path(min_length=1)],
145 offset: int,
146 session: SessionDep,
147) -> XComSequenceIndexResponse:
148 xcom_query = XComModel.get_many(
149 run_id=run_id,
150 key=key,
151 task_ids=task_id,
152 dag_ids=dag_id,
153 )
154 xcom_query = xcom_query.order_by(None)
155 if offset >= 0:
156 xcom_query = xcom_query.order_by(XComModel.map_index.asc()).offset(offset)
157 else:
158 xcom_query = xcom_query.order_by(XComModel.map_index.desc()).offset(-1 - offset)
160 result: tuple[XComModel] | None
161 if (result := session.scalars(xcom_query.limit(1)).first()) is None:
162 message = (
163 f"XCom with {key=} {offset=} not found for task {task_id!r} in DAG run {run_id!r} of {dag_id!r}"
164 )
165 raise HTTPException(
166 status_code=status.HTTP_404_NOT_FOUND,
167 detail={"reason": "not_found", "message": message},
168 )
169 return XComSequenceIndexResponse((result[0] if isinstance(result, tuple) else result).value)
172class GetXComSliceFilterParams(BaseModel):
173 """Class to house slice params."""
175 start: int | None = None
176 stop: int | None = None
177 step: int | None = None
178 include_prior_dates: bool = False
181@router.get(
182 "/{dag_id}/{run_id}/{task_id}/{key:path}/slice",
183 description="Get XCom values from a mapped task by sequence slice",
184)
185def get_mapped_xcom_by_slice(
186 dag_id: str,
187 run_id: str,
188 task_id: str,
189 key: Annotated[str, Path(min_length=1)],
190 params: Annotated[GetXComSliceFilterParams, Query()],
191 session: SessionDep,
192) -> XComSequenceSliceResponse:
193 query = XComModel.get_many(
194 run_id=run_id,
195 key=key,
196 task_ids=task_id,
197 dag_ids=dag_id,
198 include_prior_dates=params.include_prior_dates,
199 )
200 query = query.order_by(None)
202 step = params.step or 1
204 # We want to optimize negative slicing (e.g. seq[-10:]) by not doing an
205 # additional COUNT query if possible. This is possible unless both start and
206 # stop are explicitly given and have different signs.
207 if (start := params.start) is None: 207 ↛ 227line 207 didn't jump to line 227 because the condition on line 207 was always true
208 if (stop := params.stop) is None: 208 ↛ 214line 208 didn't jump to line 214 because the condition on line 208 was always true
209 if step >= 0: 209 ↛ 212line 209 didn't jump to line 212 because the condition on line 209 was always true
210 query = query.order_by(XComModel.map_index.asc())
211 else:
212 query = query.order_by(XComModel.map_index.desc())
213 step = -step
214 elif stop >= 0:
215 query = query.order_by(XComModel.map_index.asc())
216 if step >= 0:
217 query = query.limit(stop)
218 else:
219 query = query.offset(stop + 1)
220 else:
221 query = query.order_by(XComModel.map_index.desc())
222 step = -step
223 if step > 0:
224 query = query.limit(-stop - 1)
225 else:
226 query = query.offset(-stop)
227 elif start >= 0:
228 query = query.order_by(XComModel.map_index.asc())
229 if (stop := params.stop) is None:
230 if step >= 0:
231 query = query.offset(start)
232 else:
233 query = query.limit(start + 1)
234 else:
235 if stop < 0:
236 stop += get_query_count(query, session=session)
237 if step >= 0:
238 query = query.slice(start, stop)
239 else:
240 query = query.slice(stop + 1, start + 1)
241 else:
242 query = query.order_by(XComModel.map_index.desc())
243 step = -step
244 if (stop := params.stop) is None:
245 if step > 0:
246 query = query.offset(-start - 1)
247 else:
248 query = query.limit(-start)
249 else:
250 if stop >= 0:
251 stop -= get_query_count(query, session=session)
252 if step > 0:
253 query = query.slice(-1 - start, -1 - stop)
254 else:
255 query = query.slice(-stop, -start)
257 values = [row.value for row in session.execute(query.with_only_columns(XComModel.value)).all()]
258 if step != 1: 258 ↛ 259line 258 didn't jump to line 259 because the condition on line 258 was never true
259 values = values[::step]
260 return XComSequenceSliceResponse(values)
263@router.head(
264 "/{dag_id}/{run_id}/{task_id}/{key:path}",
265 responses={
266 status.HTTP_200_OK: {
267 "description": "Metadata about the number of matching XCom values",
268 "headers": {
269 "Content-Range": {
270 "schema": {"pattern": r"^map_indexes \d+$"},
271 "description": "The number of (mapped) XCom values found for this task.",
272 },
273 },
274 },
275 },
276 description="Returns the count of mapped XCom values found in the `Content-Range` response header",
277)
278def head_xcom(
279 response: Response,
280 session: SessionDep,
281 xcom_query: Annotated[Select, Depends(xcom_query)],
282 map_index: Annotated[int | None, Query()] = None,
283) -> None:
284 """Get the count of XComs from database - not other XCom Backends."""
285 if map_index is not None:
286 raise HTTPException(
287 status_code=status.HTTP_400_BAD_REQUEST,
288 detail={"reason": "invalid_request", "message": "Cannot specify map_index in a HEAD request"},
289 )
291 count = get_query_count(xcom_query, session=session)
292 # Tell the caller how many items in this query. We define a custom range unit (HTTP spec only defines
293 # "bytes" but we can add our own)
294 response.headers["Content-Range"] = f"map_indexes {count}"
297class GetXcomFilterParams(BaseModel):
298 """Class to house the params that can optionally be set for Get XCom."""
300 map_index: int = -1
301 include_prior_dates: bool = False
302 offset: int | None = None
305@router.get(
306 "/{dag_id}/{run_id}/{task_id}/{key:path}",
307 description="Get a single XCom Value",
308)
309def get_xcom(
310 dag_id: str,
311 run_id: str,
312 task_id: str,
313 key: Annotated[str, Path(min_length=1)],
314 session: SessionDep,
315 params: Annotated[GetXcomFilterParams, Query()],
316) -> XComResponse:
317 """Get an Airflow XCom from database - not other XCom Backends."""
318 xcom_query = XComModel.get_many(
319 run_id=run_id,
320 key=key,
321 task_ids=task_id,
322 dag_ids=dag_id,
323 include_prior_dates=params.include_prior_dates,
324 )
325 if params.offset is not None: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 xcom_query = xcom_query.where(XComModel.value.is_not(None)).order_by(None)
327 if params.offset >= 0:
328 xcom_query = xcom_query.order_by(XComModel.map_index.asc()).offset(params.offset)
329 else:
330 xcom_query = xcom_query.order_by(XComModel.map_index.desc()).offset(-1 - params.offset)
331 else:
332 xcom_query = xcom_query.where(XComModel.map_index == params.map_index)
334 # We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
335 # This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
336 # retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
337 # (which automatically deserializes using the backend), we avoid potential
338 # performance hits from retrieving large data files into the API server.
339 result: tuple[XComModel] | None
340 if (result := session.scalars(xcom_query.limit(1)).first()) is None:
341 if params.offset is None: 341 ↛ 347line 341 didn't jump to line 347 because the condition on line 341 was always true
342 message = (
343 f"XCom with {key=} map_index={params.map_index} not found for "
344 f"task {task_id!r} in DAG run {run_id!r} of {dag_id!r}"
345 )
346 else:
347 message = (
348 f"XCom with {key=} offset={params.offset} not found for "
349 f"task {task_id!r} in DAG run {run_id!r} of {dag_id!r}"
350 )
351 raise HTTPException(
352 status_code=status.HTTP_404_NOT_FOUND,
353 detail={"reason": "not_found", "message": message},
354 )
356 return XComResponse(key=key, value=(result[0] if isinstance(result, tuple) else result).value)
359# TODO: once we have JWT tokens, then remove dag_id/run_id/task_id from the URL and just use the info in
360# the token
361@router.post(
362 "/{dag_id}/{run_id}/{task_id}/{key:path}",
363 status_code=status.HTTP_201_CREATED,
364)
365def set_xcom(
366 dag_id: str,
367 run_id: str,
368 task_id: str,
369 key: Annotated[str, Path(min_length=1)],
370 session: SessionDep,
371 value: Annotated[
372 JsonValue,
373 Body(
374 description="A JSON-formatted string representing the value to set for the XCom.",
375 openapi_examples={
376 "simple_value": {
377 "summary": "Simple value",
378 "value": '"value1"',
379 },
380 "dict_value": {
381 "summary": "Dictionary value",
382 "value": '{"key2": "value2"}',
383 },
384 "list_value": {
385 "summary": "List value",
386 "value": '["value1"]',
387 },
388 },
389 ),
390 ] = None,
391 map_index: Annotated[int, Query()] = -1,
392 dag_result: Annotated[bool, Query(description="Whether this XCom is a dag result")] = False,
393 mapped_length: Annotated[
394 int | None, Query(description="Number of mapped tasks this value expands into")
395 ] = None,
396):
397 """Set an Airflow XCom."""
398 from airflow.configuration import conf
400 # Validate that the provided key is not empty
401 # XCom keys must be non-empty strings to ensure proper data retrieval and avoid ambiguity.
402 if not key: 402 ↛ 403line 402 didn't jump to line 403 because the condition on line 402 was never true
403 raise HTTPException(
404 status_code=status.HTTP_400_BAD_REQUEST,
405 detail={
406 "reason": "invalid_key",
407 "message": "XCom key must be a non-empty string.",
408 },
409 )
411 if mapped_length is not None: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true
412 task_map = TaskMap(
413 dag_id=dag_id,
414 task_id=task_id,
415 run_id=run_id,
416 map_index=map_index,
417 length=mapped_length,
418 keys=None,
419 )
420 max_map_length = conf.getint("core", "max_map_length", fallback=1024)
421 if task_map.length > max_map_length:
422 raise HTTPException(
423 status_code=status.HTTP_400_BAD_REQUEST,
424 detail={
425 "reason": "unmappable_return_value_length",
426 "message": "pushed value is too large to map as a downstream's dependency",
427 },
428 )
429 session.merge(task_map)
431 # else:
432 # TODO: Can/should we check if a client _hasn't_ provided this for an upstream of a mapped task? That
433 # means loading the serialized dag and that seems like a relatively costly operation for minimal benefit
434 # (the mapped task would fail in a moment as it can't be expanded anyway.)
435 try:
436 # We expect serialised value from the caller - sdk, do not serialise in here
437 XComModel.set(
438 key=key,
439 value=value,
440 run_id=run_id,
441 task_id=task_id,
442 dag_id=dag_id,
443 map_index=map_index,
444 serialize=False,
445 dag_result=dag_result,
446 session=session,
447 )
448 except ValueError as e:
449 raise HTTPException(status.HTTP_404_NOT_FOUND, str(e))
450 except TypeError as e:
451 raise HTTPException(
452 status_code=status.HTTP_400_BAD_REQUEST,
453 detail={
454 "reason": "invalid_format",
455 "message": f"XCom value is not a valid JSON: {e}",
456 },
457 )
459 return {"message": "XCom successfully set"}
462@router.delete(
463 "/{dag_id}/{run_id}/{task_id}/{key:path}",
464 responses={status.HTTP_404_NOT_FOUND: {"description": "XCom not found"}},
465 description="Delete a single XCom Value",
466)
467def delete_xcom(
468 session: SessionDep,
469 dag_id: str,
470 run_id: str,
471 task_id: str,
472 key: Annotated[str, Path(min_length=1)],
473 map_index: Annotated[int, Query()] = -1,
474):
475 """Delete a single XCom Value."""
476 query = delete(XComModel).where(
477 XComModel.key == key,
478 XComModel.run_id == run_id,
479 XComModel.task_id == task_id,
480 XComModel.dag_id == dag_id,
481 XComModel.map_index == map_index,
482 )
483 session.execute(query)
484 return {"message": f"XCom with key: {key} successfully deleted."}