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

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. 

17 

18from __future__ import annotations 

19 

20from collections.abc import Sequence 

21from typing import Annotated, Literal, cast 

22 

23import structlog 

24from fastapi import Depends, HTTPException, Query, status 

25from sqlalchemy import or_, select 

26from sqlalchemy.orm import joinedload 

27from sqlalchemy.sql.selectable import Select 

28 

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

30from airflow.api_fastapi.common.cursors import ( 

31 apply_cursor_filter, 

32 encode_cursor, 

33 make_backward_cursor, 

34 parse_cursor, 

35) 

36from airflow.api_fastapi.common.dagbag import ( 

37 DagBagDep, 

38 get_dag_for_run, 

39 get_dag_for_run_or_latest_version, 

40 get_latest_version_of_dag, 

41 resolve_run_on_latest_version, 

42) 

43from airflow.api_fastapi.common.db.common import ( 

44 SessionDep, 

45 apply_filters_to_select, 

46 bounded_total_entries, 

47 paginated_select, 

48) 

49from airflow.api_fastapi.common.db.task_instances import eager_load_TI_and_TIH_for_validation 

50from airflow.api_fastapi.common.parameters import ( 

51 FilterOptionEnum, 

52 FilterParam, 

53 LimitFilter, 

54 OffsetFilter, 

55 QueryLimit, 

56 QueryOffset, 

57 QueryTIDagVersionFilter, 

58 QueryTIExecutorFilter, 

59 QueryTIMapIndexFilter, 

60 QueryTIOperatorFilter, 

61 QueryTIOperatorNamePatternSearch, 

62 QueryTIOperatorNamePrefixPatternSearch, 

63 QueryTIPoolFilter, 

64 QueryTIPoolNamePatternSearch, 

65 QueryTIPoolNamePrefixPatternSearch, 

66 QueryTIQueueFilter, 

67 QueryTIQueueNamePatternSearch, 

68 QueryTIQueueNamePrefixPatternSearch, 

69 QueryTIRenderedMapIndexPatternSearch, 

70 QueryTIRenderedMapIndexPrefixPatternSearch, 

71 QueryTIStateFilter, 

72 QueryTITaskDisplayNamePatternSearch, 

73 QueryTITaskDisplayNamePrefixPatternSearch, 

74 QueryTITaskGroupFilter, 

75 QueryTITryNumberFilter, 

76 Range, 

77 RangeFilter, 

78 SortParam, 

79 _PrefixSearchParam, 

80 _SearchParam, 

81 datetime_range_filter_factory, 

82 filter_param_factory, 

83 float_range_filter_factory, 

84 prefix_search_param_factory, 

85 search_param_factory, 

86) 

87from airflow.api_fastapi.common.router import AirflowRouter 

88from airflow.api_fastapi.core_api.base import OrmClause 

89from airflow.api_fastapi.core_api.datamodels.common import BulkBody, BulkResponse 

90from airflow.api_fastapi.core_api.datamodels.task_instance_history import ( 

91 TaskInstanceHistoryCollectionResponse, 

92 TaskInstanceHistoryResponse, 

93) 

94from airflow.api_fastapi.core_api.datamodels.task_instances import ( 

95 BulkTaskInstanceBody, 

96 ClearTaskInstancesBody, 

97 PatchTaskInstanceBody, 

98 TaskDependencyCollectionResponse, 

99 TaskInstanceCollectionResponse, 

100 TaskInstanceResponse, 

101 TaskInstancesBatchBody, 

102) 

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

104from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag 

105from airflow.api_fastapi.core_api.services.public.task_instances import ( 

106 BulkTaskInstanceService, 

107 _get_task_group_task_instances, 

108 _patch_task_group_state, 

109 _patch_task_instance_note, 

110 _patch_task_instance_state, 

111 _patch_ti_group_validate_request, 

112 _patch_ti_validate_request, 

113 _reload_tis_with_rendered_fields, 

114) 

115from airflow.api_fastapi.logging.decorators import action_logging 

116from airflow.exceptions import AirflowClearRunningTaskException, TaskNotFound 

117from airflow.models import Base, DagRun 

118from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances 

119from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH 

120from airflow.ti_deps.dep_context import DepContext 

121from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS 

122from airflow.utils.db import get_query_count 

123from airflow.utils.state import DagRunState, TaskInstanceState 

124 

125log = structlog.get_logger(__name__) 

126 

127task_instances_router = AirflowRouter(tags=["Task Instance"], prefix="/dags/{dag_id}") 

128task_instances_prefix = "/dagRuns/{dag_run_id}/taskInstances" 

129 

130 

131@task_instances_router.get( 

132 task_instances_prefix + "/{task_id}", 

133 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

134 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

135) 

136def get_task_instance( 

137 dag_id: str, 

138 dag_run_id: str, 

139 task_id: str, 

140 session: SessionDep, 

141) -> TaskInstanceResponse: 

142 """Get task instance.""" 

143 query = ( 

144 select(TI) 

145 .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id) 

146 .options(joinedload(TI.rendered_task_instance_fields)) 

147 .options(joinedload(TI.dag_version)) 

148 .options(joinedload(TI.dag_run).options(joinedload(DagRun.dag_model))) 

149 ) 

150 task_instance = session.scalar(query) 

151 

152 if task_instance is None: 

153 raise HTTPException( 

154 status.HTTP_404_NOT_FOUND, 

155 f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}` and task_id: `{task_id}` was not found", 

156 ) 

157 if task_instance.map_index != -1: 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true

158 raise HTTPException( 

159 status.HTTP_404_NOT_FOUND, "Task instance is mapped, add the map_index value to the URL" 

160 ) 

161 

162 return task_instance 

163 

164 

165@task_instances_router.get( 

166 task_instances_prefix + "/{task_id}/listMapped", 

167 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

168 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

169) 

170def get_mapped_task_instances( 

171 dag_id: str, 

172 dag_run_id: str, 

173 task_id: str, 

174 dag_bag: DagBagDep, 

175 run_after_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("run_after", TI))], 

176 logical_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("logical_date", TI))], 

177 start_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("start_date", TI))], 

178 end_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("end_date", TI))], 

179 update_at_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("updated_at", TI))], 

180 duration_range: Annotated[RangeFilter, Depends(float_range_filter_factory("duration", TI))], 

181 state: QueryTIStateFilter, 

182 pool: QueryTIPoolFilter, 

183 pool_name_pattern: QueryTIPoolNamePatternSearch, 

184 pool_name_prefix_pattern: QueryTIPoolNamePrefixPatternSearch, 

185 queue: QueryTIQueueFilter, 

186 queue_name_pattern: QueryTIQueueNamePatternSearch, 

187 queue_name_prefix_pattern: QueryTIQueueNamePrefixPatternSearch, 

188 executor: QueryTIExecutorFilter, 

189 version_number: QueryTIDagVersionFilter, 

190 try_number: QueryTITryNumberFilter, 

191 operator: QueryTIOperatorFilter, 

192 operator_name_pattern: QueryTIOperatorNamePatternSearch, 

193 operator_name_prefix_pattern: QueryTIOperatorNamePrefixPatternSearch, 

194 map_index: QueryTIMapIndexFilter, 

195 rendered_map_index_pattern: QueryTIRenderedMapIndexPatternSearch, 

196 rendered_map_index_prefix_pattern: QueryTIRenderedMapIndexPrefixPatternSearch, 

197 limit: QueryLimit, 

198 offset: QueryOffset, 

199 order_by: Annotated[ 

200 SortParam, 

201 Depends( 

202 SortParam( 

203 [ 

204 "id", 

205 "state", 

206 "duration", 

207 "start_date", 

208 "end_date", 

209 "map_index", 

210 "try_number", 

211 "logical_date", 

212 "run_after", 

213 "data_interval_start", 

214 "data_interval_end", 

215 "rendered_map_index", 

216 "operator", 

217 ], 

218 TI, 

219 to_replace={ 

220 "run_after": DagRun.run_after, 

221 "logical_date": DagRun.logical_date, 

222 "data_interval_start": DagRun.data_interval_start, 

223 "data_interval_end": DagRun.data_interval_end, 

224 # Compound sort: when _rendered_map_index is NULL (no map_index_template), 

225 # all primary values tie and the integer map_index is the effective key, 

226 # giving correct numeric ordering (0, 1, 2, 10…) rather than lexicographic 

227 # ("0", "1", "10", "2"…). When _rendered_map_index is set (map_index_template 

228 # used), TIs are ordered by their human-readable label first, then by 

229 # map_index for identical labels. 

230 "rendered_map_index": [TI._rendered_map_index, TI.map_index], 

231 }, 

232 ).dynamic_depends(default="map_index") 

233 ), 

234 ], 

235 session: SessionDep, 

236) -> TaskInstanceCollectionResponse: 

237 """Get list of mapped task instances.""" 

238 query = eager_load_TI_and_TIH_for_validation( 

239 select(TI).where( 

240 TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id, TI.map_index >= 0 

241 ) 

242 ) 

243 # 0 can mean a mapped TI that expanded to an empty list, so it is not an automatic 404 

244 unfiltered_total_count = get_query_count(query, session=session) 

245 if unfiltered_total_count == 0: 245 ↛ 257line 245 didn't jump to line 257 because the condition on line 245 was always true

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

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

248 try: 

249 task = dag.get_task(task_id) 

250 except TaskNotFound: 

251 error_message = f"Task id {task_id} not found" 

252 raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) 

253 if not task.get_needs_expansion(): 253 ↛ 257line 253 didn't jump to line 257 because the condition on line 253 was always true

254 error_message = f"Task id {task_id} is not mapped" 

255 raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) 

256 

257 task_instance_select, total_entries = paginated_select( 

258 statement=query, 

259 filters=[ 

260 run_after_range, 

261 logical_date_range, 

262 start_date_range, 

263 end_date_range, 

264 update_at_range, 

265 duration_range, 

266 state, 

267 pool, 

268 pool_name_pattern, 

269 pool_name_prefix_pattern, 

270 queue, 

271 queue_name_pattern, 

272 queue_name_prefix_pattern, 

273 executor, 

274 version_number, 

275 try_number, 

276 operator, 

277 operator_name_pattern, 

278 operator_name_prefix_pattern, 

279 map_index, 

280 rendered_map_index_pattern, 

281 rendered_map_index_prefix_pattern, 

282 ], 

283 order_by=order_by, 

284 offset=offset, 

285 limit=limit, 

286 session=session, 

287 ) 

288 task_instances = session.scalars(task_instance_select) 

289 

290 return TaskInstanceCollectionResponse( 

291 task_instances=task_instances, 

292 total_entries=total_entries, 

293 ) 

294 

295 

296@task_instances_router.get( 

297 task_instances_prefix + "/{task_id}/dependencies", 

298 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

299 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

300 operation_id="get_task_instance_dependencies", 

301) 

302@task_instances_router.get( 

303 task_instances_prefix + "/{task_id}/{map_index}/dependencies", 

304 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

305 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

306 operation_id="get_task_instance_dependencies_by_map_index", 

307) 

308def get_task_instance_dependencies( 

309 dag_id: str, 

310 dag_run_id: str, 

311 task_id: str, 

312 session: SessionDep, 

313 dag_bag: DagBagDep, 

314 map_index: int = -1, 

315) -> TaskDependencyCollectionResponse: 

316 """Get dependencies blocking task from getting scheduled.""" 

317 query = select(TI).where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id) 

318 query = query.where(TI.map_index == map_index) 

319 

320 result = session.execute(query).one_or_none() 

321 

322 if result is None: 

323 error_message = ( 

324 f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", 

325 ) 

326 raise HTTPException(status.HTTP_404_NOT_FOUND, error_message) 

327 

328 ti = result[0] 

329 deps = [] 

330 

331 if ti.state in [None, TaskInstanceState.SCHEDULED]: 331 ↛ 332line 331 didn't jump to line 332 because the condition on line 331 was never true

332 dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == ti.dag_id, DagRun.run_id == ti.run_id)) 

333 if dag_run: 

334 dag = dag_bag.get_dag_for_run(dag_run, session=session) 

335 else: 

336 dag = None 

337 

338 if dag: 

339 try: 

340 ti.task = dag.get_task(ti.task_id) 

341 except TaskNotFound: 

342 pass 

343 else: 

344 dep_context = DepContext(SCHEDULER_QUEUED_DEPS) 

345 deps = sorted( 

346 [ 

347 {"name": dep.dep_name, "reason": dep.reason} 

348 for dep in ti.get_failed_dep_statuses(dep_context=dep_context, session=session) 

349 ], 

350 key=lambda x: x["name"], 

351 ) 

352 

353 return TaskDependencyCollectionResponse.model_validate({"dependencies": deps}) 

354 

355 

356@task_instances_router.get( 

357 task_instances_prefix + "/{task_id}/tries", 

358 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

359 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

360) 

361def get_task_instance_tries( 

362 dag_id: str, 

363 dag_run_id: str, 

364 task_id: str, 

365 session: SessionDep, 

366 map_index: int = -1, 

367) -> TaskInstanceHistoryCollectionResponse: 

368 """Get list of task instances history.""" 

369 

370 def _query(orm_object: Base) -> Select: 

371 query = eager_load_TI_and_TIH_for_validation( 

372 select(orm_object).where( 

373 orm_object.dag_id == dag_id, 

374 orm_object.run_id == dag_run_id, 

375 orm_object.task_id == task_id, 

376 orm_object.map_index == map_index, 

377 ), 

378 orm_model=orm_object, 

379 ).options(joinedload(orm_object.hitl_detail)) 

380 return query 

381 

382 # Exclude TaskInstance with state UP_FOR_RETRY since they have been recorded in TaskInstanceHistory 

383 tis = session.scalars( 

384 _query(TI).where(or_(TI.state != TaskInstanceState.UP_FOR_RETRY, TI.state.is_(None))) 

385 ).all() 

386 task_instances = list(session.scalars(_query(TIH)).all()) + list(tis) 

387 

388 if not task_instances: 

389 raise HTTPException( 

390 status.HTTP_404_NOT_FOUND, 

391 f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found", 

392 ) 

393 return TaskInstanceHistoryCollectionResponse( 

394 task_instances=cast("list[TaskInstanceHistoryResponse]", task_instances), 

395 total_entries=len(task_instances), 

396 ) 

397 

398 

399@task_instances_router.get( 

400 task_instances_prefix + "/{task_id}/{map_index}/tries", 

401 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

402 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

403) 

404def get_mapped_task_instance_tries( 

405 dag_id: str, 

406 dag_run_id: str, 

407 task_id: str, 

408 session: SessionDep, 

409 map_index: int, 

410) -> TaskInstanceHistoryCollectionResponse: 

411 return get_task_instance_tries( 

412 dag_id=dag_id, 

413 dag_run_id=dag_run_id, 

414 task_id=task_id, 

415 map_index=map_index, 

416 session=session, 

417 ) 

418 

419 

420@task_instances_router.get( 

421 task_instances_prefix + "/{task_id}/{map_index}", 

422 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

423 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

424) 

425def get_mapped_task_instance( 

426 dag_id: str, 

427 dag_run_id: str, 

428 task_id: str, 

429 map_index: int, 

430 session: SessionDep, 

431) -> TaskInstanceResponse: 

432 """Get task instance.""" 

433 query = ( 

434 select(TI) 

435 .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id, TI.map_index == map_index) 

436 .options(joinedload(TI.rendered_task_instance_fields)) 

437 .options(joinedload(TI.dag_version)) 

438 .options(joinedload(TI.dag_run).options(joinedload(DagRun.dag_model))) 

439 ) 

440 task_instance = session.scalar(query) 

441 

442 if task_instance is None: 442 ↛ 448line 442 didn't jump to line 448 because the condition on line 442 was always true

443 raise HTTPException( 

444 status.HTTP_404_NOT_FOUND, 

445 f"The Mapped Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}`, and map_index: `{map_index}` was not found", 

446 ) 

447 

448 return task_instance 

449 

450 

451@task_instances_router.get( 

452 task_instances_prefix, 

453 responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_404_NOT_FOUND]), 

454 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

455) 

456def get_task_instances( 

457 dag_id: str, 

458 dag_run_id: str, 

459 dag_bag: DagBagDep, 

460 task_id: Annotated[FilterParam[str | None], Depends(filter_param_factory(TI.task_id, str | None))], 

461 run_after_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("run_after", TI))], 

462 logical_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("logical_date", TI))], 

463 start_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("start_date", TI))], 

464 end_date_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("end_date", TI))], 

465 update_at_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("updated_at", TI))], 

466 duration_range: Annotated[RangeFilter, Depends(float_range_filter_factory("duration", TI))], 

467 task_display_name_pattern: QueryTITaskDisplayNamePatternSearch, 

468 task_display_name_prefix_pattern: QueryTITaskDisplayNamePrefixPatternSearch, 

469 task_group_id: QueryTITaskGroupFilter, 

470 dag_id_pattern: Annotated[_SearchParam, Depends(search_param_factory(TI.dag_id, "dag_id_pattern"))], 

471 dag_id_prefix_pattern: Annotated[ 

472 _PrefixSearchParam, 

473 Depends(prefix_search_param_factory(TI.dag_id, "dag_id_prefix_pattern")), 

474 ], 

475 run_id_pattern: Annotated[_SearchParam, Depends(search_param_factory(TI.run_id, "run_id_pattern"))], 

476 run_id_prefix_pattern: Annotated[ 

477 _PrefixSearchParam, 

478 Depends(prefix_search_param_factory(TI.run_id, "run_id_prefix_pattern")), 

479 ], 

480 state: QueryTIStateFilter, 

481 pool: QueryTIPoolFilter, 

482 pool_name_pattern: QueryTIPoolNamePatternSearch, 

483 pool_name_prefix_pattern: QueryTIPoolNamePrefixPatternSearch, 

484 queue: QueryTIQueueFilter, 

485 queue_name_pattern: QueryTIQueueNamePatternSearch, 

486 queue_name_prefix_pattern: QueryTIQueueNamePrefixPatternSearch, 

487 executor: QueryTIExecutorFilter, 

488 version_number: QueryTIDagVersionFilter, 

489 try_number: QueryTITryNumberFilter, 

490 operator: QueryTIOperatorFilter, 

491 operator_name_pattern: QueryTIOperatorNamePatternSearch, 

492 operator_name_prefix_pattern: QueryTIOperatorNamePrefixPatternSearch, 

493 map_index: QueryTIMapIndexFilter, 

494 rendered_map_index_pattern: QueryTIRenderedMapIndexPatternSearch, 

495 rendered_map_index_prefix_pattern: QueryTIRenderedMapIndexPrefixPatternSearch, 

496 limit: QueryLimit, 

497 offset: QueryOffset, 

498 order_by: Annotated[ 

499 SortParam, 

500 Depends( 

501 SortParam( 

502 [ 

503 "id", 

504 "state", 

505 "duration", 

506 "start_date", 

507 "end_date", 

508 "map_index", 

509 "try_number", 

510 "logical_date", 

511 "run_after", 

512 "data_interval_start", 

513 "data_interval_end", 

514 "rendered_map_index", 

515 "operator", 

516 ], 

517 TI, 

518 to_replace={ 

519 "logical_date": DagRun.logical_date, 

520 "run_after": DagRun.run_after, 

521 "data_interval_start": DagRun.data_interval_start, 

522 "data_interval_end": DagRun.data_interval_end, 

523 # Compound sort: see the listMapped endpoint comment for rationale. 

524 "rendered_map_index": [TI._rendered_map_index, TI.map_index], 

525 }, 

526 ).dynamic_depends(default="map_index") 

527 ), 

528 ], 

529 readable_ti_filter: ReadableTIFilterDep, 

530 session: SessionDep, 

531 cursor: str | None = Query( 

532 None, 

533 description="Cursor for keyset-based pagination. " 

534 "Pass an empty string for the first page, then use ``next_cursor`` from the response. " 

535 "When ``cursor`` is provided, ``offset`` is ignored.", 

536 ), 

537) -> TaskInstanceCollectionResponse: 

538 """ 

539 Get list of task instances. 

540 

541 This endpoint allows specifying `~` as the dag_id, dag_run_id 

542 to retrieve task instances for all Dags and Dag runs. 

543 

544 Supports two pagination modes: 

545 

546 **Offset (default):** use `limit` and `offset` query parameters. Returns `total_entries`. 

547 

548 **Cursor:** pass `cursor` (empty string for the first page, then `next_cursor` from the response). 

549 When `cursor` is provided, `offset` is ignored and `total_entries` is capped at 

550 `total_entries_limit` (a value equal to that limit means at least that many task instances 

551 match). ``next_cursor`` is ``null`` when there are no more pages; ``previous_cursor`` is 

552 ``null`` on the first page. 

553 """ 

554 use_cursor = cursor is not None 

555 dag_run = None 

556 query = eager_load_TI_and_TIH_for_validation(select(TI)) 

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

558 if dag_id == "~": 

559 raise HTTPException( 

560 status.HTTP_400_BAD_REQUEST, 

561 "dag_id is required when dag_run_id is specified", 

562 ) 

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

564 if not dag_run: 

565 raise HTTPException( 

566 status.HTTP_404_NOT_FOUND, 

567 f"DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found", 

568 ) 

569 query = query.where(TI.run_id == dag_run_id) 

570 if dag_id != "~": 570 ↛ 576line 570 didn't jump to line 576 because the condition on line 570 was always true

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

572 query = query.where(TI.dag_id == dag_id) 

573 if dag: 573 ↛ 576line 573 didn't jump to line 576 because the condition on line 573 was always true

574 task_group_id.dag = dag 

575 

576 filters: list[OrmClause] = [ 

577 run_after_range, 

578 logical_date_range, 

579 start_date_range, 

580 end_date_range, 

581 update_at_range, 

582 duration_range, 

583 state, 

584 pool, 

585 pool_name_pattern, 

586 pool_name_prefix_pattern, 

587 queue, 

588 queue_name_pattern, 

589 queue_name_prefix_pattern, 

590 executor, 

591 task_id, 

592 task_display_name_pattern, 

593 task_display_name_prefix_pattern, 

594 task_group_id, 

595 dag_id_pattern, 

596 dag_id_prefix_pattern, 

597 run_id_pattern, 

598 run_id_prefix_pattern, 

599 version_number, 

600 readable_ti_filter, 

601 try_number, 

602 operator, 

603 operator_name_pattern, 

604 operator_name_prefix_pattern, 

605 map_index, 

606 rendered_map_index_pattern, 

607 rendered_map_index_prefix_pattern, 

608 ] 

609 

610 if use_cursor: 

611 # Fetch one extra row so we can detect whether a next page exists. 

612 page_limit = cast( 

613 "int", limit.value 

614 ) # LimitFilter value is guaranteed to be set to the default value of QueryLimit 

615 cursor_limit = LimitFilter().set_value(page_limit + 1) 

616 task_instance_select = apply_filters_to_select(statement=query, filters=[*filters, cursor_limit]) 

617 task_instance_select = order_by.to_orm(task_instance_select) 

618 

619 is_backward = False 

620 if cursor: 

621 token, is_backward = parse_cursor(cursor) 

622 if is_backward: 

623 task_instance_select = order_by.to_orm(task_instance_select, reversed=True) 

624 task_instance_select = apply_cursor_filter( 

625 task_instance_select, 

626 token, 

627 order_by, 

628 session.get_bind().dialect.name, 

629 is_backward=is_backward, 

630 ) 

631 

632 fetched = list(session.scalars(task_instance_select)) 

633 has_more = len(fetched) > page_limit 

634 task_instances = fetched[:page_limit] 

635 

636 if is_backward: 

637 task_instances.reverse() 

638 has_prev = has_more 

639 has_next = True 

640 else: 

641 has_prev = bool(cursor) 

642 has_next = has_more 

643 

644 total_entries, total_entries_limit = bounded_total_entries( 

645 statement=query, filters=filters, session=session 

646 ) 

647 return TaskInstanceCollectionResponse( 

648 task_instances=task_instances, 

649 total_entries=total_entries, 

650 total_entries_limit=total_entries_limit, 

651 next_cursor=( 

652 encode_cursor(task_instances[-1], order_by) if has_next and task_instances else None 

653 ), 

654 previous_cursor=( 

655 make_backward_cursor(encode_cursor(task_instances[0], order_by)) 

656 if has_prev and task_instances 

657 else None 

658 ), 

659 ) 

660 

661 task_instance_select, total_entries = paginated_select( 

662 statement=query, 

663 filters=filters, 

664 order_by=order_by, 

665 offset=offset, 

666 limit=limit, 

667 session=session, 

668 ) 

669 task_instances = list(session.scalars(task_instance_select)) 

670 return TaskInstanceCollectionResponse( 

671 task_instances=task_instances, 

672 total_entries=total_entries, 

673 ) 

674 

675 

676@task_instances_router.post( 

677 task_instances_prefix + "/list", 

678 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

679 dependencies=[ 

680 Depends(action_logging()), 

681 Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE)), 

682 ], 

683) 

684def get_task_instances_batch( 

685 dag_id: Literal["~"], 

686 dag_run_id: Literal["~"], 

687 body: TaskInstancesBatchBody, 

688 readable_ti_filter: ReadableTIFilterDep, 

689 session: SessionDep, 

690) -> TaskInstanceCollectionResponse: 

691 """Get list of task instances.""" 

692 dag_ids = FilterParam(TI.dag_id, body.dag_ids, FilterOptionEnum.IN) # type: ignore[arg-type] 

693 dag_run_ids = FilterParam(TI.run_id, body.dag_run_ids, FilterOptionEnum.IN) # type: ignore[arg-type] 

694 task_ids = FilterParam(TI.task_id, body.task_ids, FilterOptionEnum.IN) # type: ignore[arg-type] 

695 run_after = RangeFilter( 

696 Range( 

697 lower_bound_gte=body.run_after_gte, 

698 lower_bound_gt=body.run_after_gt, 

699 upper_bound_lte=body.run_after_lte, 

700 upper_bound_lt=body.run_after_lt, 

701 ), 

702 attribute=DagRun.run_after, 

703 ) 

704 logical_date = RangeFilter( 

705 Range( 

706 lower_bound_gte=body.logical_date_gte, 

707 lower_bound_gt=body.logical_date_gt, 

708 upper_bound_lte=body.logical_date_lte, 

709 upper_bound_lt=body.logical_date_lt, 

710 ), 

711 attribute=DagRun.logical_date, 

712 ) 

713 start_date = RangeFilter( 

714 Range( 

715 lower_bound_gte=body.start_date_gte, 

716 lower_bound_gt=body.start_date_gt, 

717 upper_bound_lte=body.start_date_lte, 

718 upper_bound_lt=body.start_date_lt, 

719 ), 

720 attribute=TI.start_date, # type: ignore[arg-type] 

721 ) 

722 end_date = RangeFilter( 

723 Range( 

724 lower_bound_gte=body.end_date_gte, 

725 lower_bound_gt=body.end_date_gt, 

726 upper_bound_lte=body.end_date_lte, 

727 upper_bound_lt=body.end_date_lt, 

728 ), 

729 attribute=TI.end_date, # type: ignore[arg-type] 

730 ) 

731 duration = RangeFilter( 

732 Range( 

733 lower_bound_gte=body.duration_gte, 

734 lower_bound_gt=body.duration_gt, 

735 upper_bound_lte=body.duration_lte, 

736 upper_bound_lt=body.duration_lt, 

737 ), 

738 attribute=TI.duration, # type: ignore[arg-type] 

739 ) 

740 state = FilterParam(TI.state, body.state, FilterOptionEnum.ANY_EQUAL) # type: ignore[arg-type] 

741 pool = FilterParam(TI.pool, body.pool, FilterOptionEnum.ANY_EQUAL) # type: ignore[arg-type] 

742 queue = FilterParam(TI.queue, body.queue, FilterOptionEnum.ANY_EQUAL) # type: ignore[arg-type] 

743 executor = FilterParam(TI.executor, body.executor, FilterOptionEnum.ANY_EQUAL) # type: ignore[arg-type] 

744 

745 offset = OffsetFilter(body.page_offset) 

746 limit = LimitFilter(body.page_limit) 

747 

748 order_by = SortParam( 

749 ["id", "state", "duration", "start_date", "end_date", "map_index"], 

750 TI, 

751 ).set_value([body.order_by] if body.order_by else None) 

752 

753 query = eager_load_TI_and_TIH_for_validation(select(TI)) 

754 task_instance_select, total_entries = paginated_select( 

755 statement=query, 

756 filters=[ 

757 dag_ids, 

758 dag_run_ids, 

759 task_ids, 

760 run_after, 

761 logical_date, 

762 start_date, 

763 end_date, 

764 duration, 

765 state, 

766 pool, 

767 queue, 

768 executor, 

769 readable_ti_filter, 

770 ], 

771 order_by=order_by, 

772 offset=offset, 

773 limit=limit, 

774 session=session, 

775 ) 

776 task_instances = session.scalars(task_instance_select) 

777 

778 return TaskInstanceCollectionResponse( 

779 task_instances=task_instances, 

780 total_entries=total_entries, 

781 ) 

782 

783 

784@task_instances_router.get( 

785 task_instances_prefix + "/{task_id}/tries/{task_try_number}", 

786 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), 

787 dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.TASK_INSTANCE))], 

788) 

789def get_task_instance_try_details( 

790 dag_id: str, 

791 dag_run_id: str, 

792 task_id: str, 

793 task_try_number: int, 

794 session: SessionDep, 

795 map_index: int = -1, 

796) -> TaskInstanceHistoryResponse: 

797 """Get task instance details by try number.""" 

798 

799 def _query(orm_object: Base) -> TI | TIH | None: 

800 query = select(orm_object).where( 

801 orm_object.dag_id == dag_id, 

802 orm_object.run_id == dag_run_id, 

803 orm_object.task_id == task_id, 

804 orm_object.try_number == task_try_number, 

805 orm_object.map_index == map_index, 

806 ) 

807 

808 ti_or_tih = session.scalar(query) 

809 return ti_or_tih 

810 

811 ti_or_tih = _query(TI) or _query(TIH) 

812 if ti_or_tih is None: 

813 raise HTTPException( 

814 status.HTTP_404_NOT_FOUND, 

815 f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}`, try_number: `{task_try_number}` and map_index: `{map_index}` was not found", 

816 ) 

817 return ti_or_tih 

818 

819 

820@task_instances_router.get( 

821 task_instances_prefix + "/{task_id}/{map_index}/tries/{task_try_number}", 

822 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),