Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/routes/public/xcom.py: 92%

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. 

17from __future__ import annotations 

18 

19import copy 

20from typing import Annotated 

21 

22from fastapi import Depends, HTTPException, Query, status 

23from sqlalchemy import and_, delete, select 

24from sqlalchemy.orm import joinedload 

25 

26from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity 

27from airflow.api_fastapi.common.dagbag import DagBagDep, get_dag_for_run_or_latest_version 

28from airflow.api_fastapi.common.db.common import SessionDep, paginated_select 

29from airflow.api_fastapi.common.parameters import ( 

30 FilterParam, 

31 QueryLimit, 

32 QueryOffset, 

33 QueryXComDagDisplayNamePatternSearch, 

34 QueryXComDagDisplayNamePrefixPatternSearch, 

35 QueryXComKeyPatternSearch, 

36 QueryXComKeyPrefixPatternSearch, 

37 QueryXComRunIdPatternSearch, 

38 QueryXComRunIdPrefixPatternSearch, 

39 QueryXComTaskIdPatternSearch, 

40 QueryXComTaskIdPrefixPatternSearch, 

41 RangeFilter, 

42 SortParam, 

43 datetime_range_filter_factory, 

44 filter_param_factory, 

45) 

46from airflow.api_fastapi.common.router import AirflowRouter 

47from airflow.api_fastapi.core_api.datamodels.xcom import ( 

48 XComCollectionResponse, 

49 XComCreateBody, 

50 XComResponseNative, 

51 XComResponseString, 

52 XComUpdateBody, 

53) 

54from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc 

55from airflow.api_fastapi.core_api.security import ReadableXComFilterDep, requires_access_dag 

56from airflow.api_fastapi.logging.decorators import action_logging 

57from airflow.exceptions import TaskNotFound 

58from airflow.models import DagRun as DR 

59from airflow.models.dag import DagModel 

60from airflow.models.xcom import XComModel 

61 

62xcom_router = AirflowRouter( 

63 tags=["XCom"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries" 

64) 

65 

66 

67@xcom_router.get( 

68 "/{xcom_key:path}", 

69 responses=create_openapi_http_exception_doc( 

70 [ 

71 status.HTTP_400_BAD_REQUEST, 

72 status.HTTP_404_NOT_FOUND, 

73 ] 

74 ), 

75 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], 

76) 

77def get_xcom_entry( 

78 dag_id: str, 

79 task_id: str, 

80 dag_run_id: str, 

81 xcom_key: str, 

82 session: SessionDep, 

83 map_index: Annotated[int, Query(ge=-1)] = -1, 

84 deserialize: Annotated[bool, Query()] = False, 

85 stringify: Annotated[bool, Query()] = False, 

86) -> XComResponseNative | XComResponseString: 

87 """Get an XCom entry.""" 

88 xcom_query = XComModel.get_many( 

89 run_id=dag_run_id, 

90 key=xcom_key, 

91 task_ids=task_id, 

92 dag_ids=dag_id, 

93 map_indexes=map_index, 

94 limit=1, 

95 ).options(joinedload(XComModel.task), joinedload(XComModel.dag_run).joinedload(DR.dag_model)) 

96 

97 # We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend. 

98 # This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead 

99 # retrieves the raw serialized value from the database. 

100 raw_result: tuple[XComModel] | None = session.scalars(xcom_query).first() 

101 

102 if raw_result is None: 

103 raise HTTPException(status.HTTP_404_NOT_FOUND, f"XCom entry with key: `{xcom_key}` not found") 

104 result = raw_result[0] if isinstance(raw_result, tuple) else raw_result 

105 

106 item = copy.copy(result) 

107 

108 if deserialize: 

109 # Custom XCom backends may store references (eg: object storage paths) in the database. 

110 # The custom XCom backend's deserialize_value() resolves these to actual values, but that is only 

111 # used on workers during task execution. The API reads directly from the database and uses 

112 # stringify() to convert DB values (references or serialized data) to human readable 

113 # format for UI display or for API users. 

114 import json 

115 

116 from airflow.serialization.stringify import ( 

117 StringifyNotSupportedError, 

118 stringify as stringify_xcom, 

119 ) 

120 

121 try: 

122 parsed_value = json.loads(result.value) 

123 except (ValueError, TypeError): 

124 # Already deserialized (e.g., set via Task Execution API) 

125 parsed_value = result.value 

126 

127 try: 

128 item.value = stringify_xcom(parsed_value) 

129 except StringifyNotSupportedError: 

130 item.value = XComModel.deserialize_value(result) 

131 else: 

132 # For native format, return the raw serialized value from the database 

133 # This preserves the JSON string format that the API expects 

134 item.value = result.value 

135 

136 if stringify: 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true

137 return XComResponseString.model_validate(item) 

138 return XComResponseNative.model_validate(item) 

139 

140 

141@xcom_router.get( 

142 "", 

143 responses=create_openapi_http_exception_doc( 

144 [ 

145 status.HTTP_400_BAD_REQUEST, 

146 status.HTTP_404_NOT_FOUND, 

147 ] 

148 ), 

149 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.XCOM))], 

150) 

151def get_xcom_entries( 

152 dag_id: str, 

153 dag_run_id: str, 

154 task_id: str, 

155 limit: QueryLimit, 

156 offset: QueryOffset, 

157 readable_xcom_filter: ReadableXComFilterDep, 

158 session: SessionDep, 

159 xcom_key_pattern: QueryXComKeyPatternSearch, 

160 xcom_key_prefix_pattern: QueryXComKeyPrefixPatternSearch, 

161 dag_display_name_pattern: QueryXComDagDisplayNamePatternSearch, 

162 dag_display_name_prefix_pattern: QueryXComDagDisplayNamePrefixPatternSearch, 

163 run_id_pattern: QueryXComRunIdPatternSearch, 

164 run_id_prefix_pattern: QueryXComRunIdPrefixPatternSearch, 

165 task_id_pattern: QueryXComTaskIdPatternSearch, 

166 task_id_prefix_pattern: QueryXComTaskIdPrefixPatternSearch, 

167 map_index_filter: Annotated[ 

168 FilterParam[int | None], 

169 Depends(filter_param_factory(XComModel.map_index, int | None, filter_name="map_index_filter")), 

170 ], 

171 logical_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("logical_date", DR))], 

172 run_after_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("run_after", DR))], 

173 order_by: Annotated[ 

174 SortParam, 

175 Depends( 

176 SortParam( 

177 ["key", "dag_id", "run_id", "task_id", "map_index", "timestamp"], 

178 XComModel, 

179 to_replace={"run_after": DR.run_after}, 

180 ).dynamic_depends(default=("dag_id", "task_id", "run_id", "map_index", "key")) 

181 ), 

182 ], 

183 xcom_key: Annotated[str | None, Query()] = None, 

184 map_index: Annotated[int | None, Query(ge=-1)] = None, 

185) -> XComCollectionResponse: 

186 """ 

187 Get all XCom entries. 

188 

189 This endpoint allows specifying `~` as the dag_id, dag_run_id, task_id to retrieve XCom entries for all Dags. 

190 """ 

191 query = select(XComModel) 

192 if dag_id != "~": 

193 query = query.where(XComModel.dag_id == dag_id) 

194 query = ( 

195 query.join(DR, and_(XComModel.dag_id == DR.dag_id, XComModel.run_id == DR.run_id)) 

196 .join(DagModel, DR.dag_id == DagModel.dag_id) 

197 .options(joinedload(XComModel.task), joinedload(XComModel.dag_run).joinedload(DR.dag_model)) 

198 ) 

199 

200 if task_id != "~": 200 ↛ 202line 200 didn't jump to line 202 because the condition on line 200 was always true

201 query = query.where(XComModel.task_id == task_id) 

202 if dag_run_id != "~": 202 ↛ 204line 202 didn't jump to line 204 because the condition on line 202 was always true

203 query = query.where(DR.run_id == dag_run_id) 

204 if map_index is not None: 

205 query = query.where(XComModel.map_index == map_index) 

206 if xcom_key is not None: 

207 query = query.where(XComModel.key == xcom_key) 

208 

209 query, total_entries = paginated_select( 

210 statement=query, 

211 filters=[ 

212 readable_xcom_filter, 

213 xcom_key_pattern, 

214 xcom_key_prefix_pattern, 

215 dag_display_name_pattern, 

216 dag_display_name_prefix_pattern, 

217 run_id_pattern, 

218 run_id_prefix_pattern, 

219 task_id_pattern, 

220 task_id_prefix_pattern, 

221 map_index_filter, 

222 logical_date_range, 

223 run_after_range, 

224 ], 

225 order_by=order_by, 

226 offset=offset, 

227 limit=limit, 

228 session=session, 

229 ) 

230 return XComCollectionResponse(xcom_entries=session.scalars(query), total_entries=total_entries) 

231 

232 

233@xcom_router.post( 

234 "", 

235 status_code=status.HTTP_201_CREATED, 

236 responses=create_openapi_http_exception_doc( 

237 [ 

238 status.HTTP_400_BAD_REQUEST, 

239 status.HTTP_404_NOT_FOUND, 

240 status.HTTP_409_CONFLICT, 

241 ] 

242 ), 

243 dependencies=[ 

244 Depends(action_logging()), 

245 Depends(requires_access_dag(method="POST", access_entity=DagAccessEntity.XCOM)), 

246 ], 

247) 

248def create_xcom_entry( 

249 dag_id: str, 

250 task_id: str, 

251 dag_run_id: str, 

252 request_body: XComCreateBody, 

253 session: SessionDep, 

254 dag_bag: DagBagDep, 

255) -> XComResponseNative: 

256 """Create an XCom entry.""" 

257 from airflow.models.dagrun import DagRun 

258 

259 dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id)) 

260 # Validate Dag ID 

261 dag = get_dag_for_run_or_latest_version(dag_bag, dag_run, dag_id, session) 

262 

263 # Validate Task ID 

264 try: 

265 dag.get_task(task_id) 

266 except TaskNotFound: 

267 raise HTTPException( 

268 status.HTTP_404_NOT_FOUND, f"Task with ID: `{task_id}` not found in dag: `{dag_id}`" 

269 ) 

270 

271 # Validate Dag Run ID 

272 if not dag_run: 

273 raise HTTPException( 

274 status.HTTP_404_NOT_FOUND, f"Dag Run with ID: `{dag_run_id}` not found for dag: `{dag_id}`" 

275 ) 

276 

277 # Check existing XCom 

278 already_existing_query = XComModel.get_many( 

279 key=request_body.key, 

280 task_ids=task_id, 

281 dag_ids=dag_id, 

282 run_id=dag_run_id, 

283 map_indexes=request_body.map_index, 

284 ) 

285 result = session.execute(already_existing_query.with_only_columns(XComModel.value).limit(1)).first() 

286 if result: 

287 raise HTTPException( 

288 status_code=status.HTTP_409_CONFLICT, 

289 detail=f"The XCom with key: `{request_body.key}` with mentioned task instance already exists.", 

290 ) 

291 

292 try: 

293 XComModel.set( 

294 key=request_body.key, 

295 value=request_body.value, 

296 dag_id=dag_id, 

297 task_id=task_id, 

298 run_id=dag_run_id, 

299 map_index=request_body.map_index, 

300 serialize=False, 

301 session=session, 

302 ) 

303 except (ValueError, TypeError) as e: 

304 raise HTTPException( 

305 status.HTTP_400_BAD_REQUEST, f"Couldn't serialise the XCom with key: `{request_body.key}`" 

306 ) from e 

307 

308 xcom = session.scalar( 

309 select(XComModel) 

310 .where( 

311 XComModel.dag_id == dag_id, 

312 XComModel.task_id == task_id, 

313 XComModel.run_id == dag_run_id, 

314 XComModel.key == request_body.key, 

315 XComModel.map_index == request_body.map_index, 

316 ) 

317 .limit(1) 

318 .options(joinedload(XComModel.task), joinedload(XComModel.dag_run).joinedload(DR.dag_model)) 

319 ) 

320 

321 return XComResponseNative.model_validate(xcom) 

322 

323 

324@xcom_router.patch( 

325 "/{xcom_key:path}", 

326 status_code=status.HTTP_200_OK, 

327 responses=create_openapi_http_exception_doc( 

328 [ 

329 status.HTTP_400_BAD_REQUEST, 

330 status.HTTP_404_NOT_FOUND, 

331 ] 

332 ), 

333 dependencies=[ 

334 Depends(action_logging()), 

335 Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.XCOM)), 

336 ], 

337) 

338def update_xcom_entry( 

339 dag_id: str, 

340 task_id: str, 

341 dag_run_id: str, 

342 xcom_key: str, 

343 patch_body: XComUpdateBody, 

344 *, 

345 session: SessionDep, 

346) -> XComResponseNative: 

347 """Update an existing XCom entry.""" 

348 # Check if XCom entry exists 

349 xcom_query = ( 

350 select(XComModel) 

351 .where( 

352 XComModel.dag_id == dag_id, 

353 XComModel.task_id == task_id, 

354 XComModel.run_id == dag_run_id, 

355 XComModel.key == xcom_key, 

356 XComModel.map_index == patch_body.map_index, 

357 ) 

358 .limit(1) 

359 .options(joinedload(XComModel.task), joinedload(XComModel.dag_run).joinedload(DR.dag_model)) 

360 ) 

361 xcom_entry = session.scalar(xcom_query) 

362 

363 if not xcom_entry: 

364 raise HTTPException( 

365 status.HTTP_404_NOT_FOUND, 

366 f"The XCom with key: `{xcom_key}` with mentioned task instance doesn't exist.", 

367 ) 

368 

369 try: 

370 XComModel.set( 

371 key=xcom_key, 

372 value=patch_body.value, 

373 dag_id=dag_id, 

374 task_id=task_id, 

375 run_id=dag_run_id, 

376 map_index=patch_body.map_index, 

377 serialize=False, 

378 session=session, 

379 ) 

380 except (ValueError, TypeError) as e: 

381 raise HTTPException( 

382 status.HTTP_400_BAD_REQUEST, f"Couldn't serialise the XCom with key: `{xcom_key}`" 

383 ) from e 

384 

385 # Fetch after setting, to get fresh object for response 

386 xcom_entry = session.scalar(xcom_query) 

387 return XComResponseNative.model_validate(xcom_entry) 

388 

389 

390@xcom_router.delete( 

391 "/{xcom_key:path}", 

392 status_code=status.HTTP_204_NO_CONTENT, 

393 responses=create_openapi_http_exception_doc( 

394 [ 

395 status.HTTP_400_BAD_REQUEST, 

396 status.HTTP_404_NOT_FOUND, 

397 ] 

398 ), 

399 dependencies=[ 

400 Depends(action_logging()), 

401 Depends(requires_access_dag(method="DELETE", access_entity=DagAccessEntity.XCOM)), 

402 ], 

403) 

404def delete_xcom_entry( 

405 dag_id: str, 

406 task_id: str, 

407 dag_run_id: str, 

408 xcom_key: str, 

409 session: SessionDep, 

410 map_index: Annotated[int, Query(ge=-1)] = -1, 

411): 

412 """Delete an XCom entry.""" 

413 # Delete XCom entry 

414 result = session.execute( 

415 delete(XComModel).where( 

416 XComModel.dag_id == dag_id, 

417 XComModel.task_id == task_id, 

418 XComModel.run_id == dag_run_id, 

419 XComModel.key == xcom_key, 

420 XComModel.map_index == map_index, 

421 ) 

422 ) 

423 

424 if getattr(result, "rowcount", 0) == 0: 424 ↛ exitline 424 didn't return from function 'delete_xcom_entry' because the condition on line 424 was always true

425 raise HTTPException( 

426 status.HTTP_404_NOT_FOUND, 

427 f"The XCom with key: `{xcom_key}` with mentioned task instance doesn't exist.", 

428 )