Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/public/task_instances.py: 59%
240 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
20from collections.abc import Sequence
21from typing import Literal
23import structlog
24from fastapi import HTTPException, status
25from fastapi.exceptions import RequestValidationError
26from pydantic import ValidationError
27from sqlalchemy import select, tuple_
28from sqlalchemy.orm import joinedload
29from sqlalchemy.orm.session import Session
31from airflow._shared.state import TaskScope
32from airflow.api_fastapi.app import get_auth_manager
33from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
34from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
35from airflow.api_fastapi.common.db.common import SessionDep
36from airflow.api_fastapi.core_api.datamodels.common import (
37 BulkActionNotOnExistence,
38 BulkActionResponse,
39 BulkBody,
40 BulkCreateAction,
41 BulkDeleteAction,
42 BulkUpdateAction,
43)
44from airflow.api_fastapi.core_api.datamodels.task_instances import (
45 BulkTaskInstanceBody,
46 ClearTaskInstancesBody,
47 PatchTaskInstanceBody,
48)
49from airflow.api_fastapi.core_api.security import GetUserDep
50from airflow.api_fastapi.core_api.services.public.common import BulkService
51from airflow.configuration import conf
52from airflow.listeners.listener import get_listener_manager
53from airflow.models.dag import DagModel
54from airflow.models.taskinstance import TaskInstance as TI
55from airflow.serialization.definitions.dag import SerializedDAG
56from airflow.state.metastore import _get_db_backend
57from airflow.utils.state import TaskInstanceState
59log = structlog.get_logger(__name__)
62def _clear_task_state_store_on_success(tis: Sequence[TI], session: Session) -> None:
63 """Clear task state store rows for each TI if clear_on_success is enabled."""
64 if not conf.getboolean("state_store", "clear_on_success", fallback=False): 64 ↛ 66line 64 didn't jump to line 66 because the condition on line 64 was always true
65 return
66 backend = _get_db_backend()
67 for ti in tis:
68 scope = TaskScope(
69 dag_id=ti.dag_id,
70 run_id=ti.run_id,
71 task_id=ti.task_id,
72 map_index=ti.map_index if ti.map_index is not None else -1,
73 )
74 try:
75 backend.clear(scope=scope, session=session)
76 log.info(
77 "Cleared task state on success",
78 dag_id=ti.dag_id,
79 run_id=ti.run_id,
80 task_id=ti.task_id,
81 map_index=ti.map_index,
82 )
83 except Exception:
84 log.warning(
85 "Failed to clear task state on success",
86 dag_id=ti.dag_id,
87 run_id=ti.run_id,
88 task_id=ti.task_id,
89 )
92def _validate_patch_task_instance_body(
93 body: PatchTaskInstanceBody,
94 update_mask: list[str] | None,
95) -> dict:
96 """Validate the patch body and return the fields to update as a dict."""
97 fields_to_update = body.model_fields_set
98 if update_mask:
99 fields_to_update = fields_to_update.intersection(update_mask)
100 else:
101 try:
102 PatchTaskInstanceBody.model_validate(body)
103 except ValidationError as e:
104 raise RequestValidationError(errors=e.errors())
106 return body.model_dump(include=fields_to_update, by_alias=True)
109def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInstanceState) -> None:
110 """Fire listener hooks for the given TIs based on their new state. Listener errors are logged."""
111 for ti in updated_tis:
112 try:
113 if new_state == TaskInstanceState.SUCCESS:
114 get_listener_manager().hook.on_task_instance_success(previous_state=None, task_instance=ti)
115 elif new_state == TaskInstanceState.FAILED: 115 ↛ 121line 115 didn't jump to line 121 because the condition on line 115 was always true
116 get_listener_manager().hook.on_task_instance_failed(
117 previous_state=None,
118 task_instance=ti,
119 error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
120 )
121 elif new_state == TaskInstanceState.SKIPPED:
122 get_listener_manager().hook.on_task_instance_skipped(previous_state=None, task_instance=ti)
123 except Exception:
124 log.exception("error calling listener")
127def _reload_tis_with_rendered_fields(tis: list[TI], session: Session) -> list[TI]:
128 """
129 Re-load TIs with ``rendered_task_instance_fields`` eagerly loaded.
131 ``set_task_instance_state`` / ``set_task_group_state`` return TIs without this relationship
132 loaded; we re-query so they can be serialized without lazy loads.
133 ``populate_existing=True`` ensures the joinedload updates TIs already in the identity map.
134 """
135 if not tis:
136 return tis
137 return list(
138 session.scalars(
139 select(TI)
140 .options(joinedload(TI.rendered_task_instance_fields))
141 .where(TI.id.in_([ti.id for ti in tis]))
142 .execution_options(populate_existing=True)
143 ).all()
144 )
147def _patch_ti_validate_request(
148 dag_id: str,
149 dag_run_id: str,
150 task_id: str,
151 dag_bag: DagBagDep,
152 body: PatchTaskInstanceBody,
153 session: SessionDep,
154 map_index: int | None = -1,
155 update_mask: list[str] | None = None,
156) -> tuple[SerializedDAG, list[TI], dict]:
157 dag = get_latest_version_of_dag(dag_bag, dag_id, session)
158 if not dag.has_task(task_id):
159 raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in Dag '{dag_id}'")
161 query = (
162 select(TI)
163 .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
164 .options(joinedload(TI.rendered_task_instance_fields))
165 )
166 if map_index is not None:
167 query = query.where(TI.map_index == map_index)
168 else:
169 query = query.order_by(TI.map_index)
171 tis = session.scalars(query).all()
173 err_msg_404 = (
174 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",
175 )
176 if len(tis) == 0:
177 raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
179 data = _validate_patch_task_instance_body(body, update_mask)
180 return dag, list(tis), data
183def _get_task_group_task_instances(
184 dag_id: str,
185 dag_run_id: str,
186 task_group_id: str,
187 dag: SerializedDAG,
188 session: Session,
189) -> list[TI]:
190 """Get all task instances in a task group for a specific DAG run."""
191 task_group = dag.task_group_dict.get(task_group_id)
192 if not task_group: 192 ↛ 197line 192 didn't jump to line 197 because the condition on line 192 was always true
193 raise HTTPException(
194 status.HTTP_404_NOT_FOUND, f"Task group '{task_group_id}' not found in DAG '{dag_id}'"
195 )
197 task_ids = [task.task_id for task in task_group.iter_tasks()]
199 query = (
200 select(TI)
201 .where(
202 TI.dag_id == dag_id,
203 TI.run_id == dag_run_id,
204 TI.task_id.in_(task_ids),
205 )
206 .order_by(TI.task_id, TI.map_index)
207 )
209 group_tis = list(session.scalars(query).all())
210 if not group_tis:
211 raise HTTPException(
212 status.HTTP_404_NOT_FOUND,
213 f"No task instances found for task group '{task_group_id}' in dag run '{dag_run_id}'",
214 )
216 return group_tis
219def _patch_ti_group_validate_request(
220 dag_id: str,
221 dag_run_id: str,
222 task_group_id: str,
223 dag_bag: DagBagDep,
224 body: PatchTaskInstanceBody,
225 session: SessionDep,
226 update_mask: list[str] | None = None,
227) -> tuple[SerializedDAG, list[TI], dict]:
228 """Validate and prepare data for task group patch request."""
229 dag = get_latest_version_of_dag(dag_bag, dag_id, session)
230 tis = _get_task_group_task_instances(dag_id, dag_run_id, task_group_id, dag, session)
232 data = _validate_patch_task_instance_body(body, update_mask)
233 return dag, tis, data
236def _patch_task_instance_state(
237 task_id: str,
238 dag_run_id: str,
239 dag: SerializedDAG,
240 task_instance_body: BulkTaskInstanceBody | PatchTaskInstanceBody,
241 data: dict,
242 session: Session,
243) -> list[TI]:
244 map_index = getattr(task_instance_body, "map_index", None)
245 map_indexes = None if map_index is None else [map_index]
247 updated_tis = dag.set_task_instance_state(
248 task_id=task_id,
249 run_id=dag_run_id,
250 map_indexes=map_indexes,
251 state=data["new_state"],
252 upstream=task_instance_body.include_upstream,
253 downstream=task_instance_body.include_downstream,
254 future=task_instance_body.include_future,
255 past=task_instance_body.include_past,
256 commit=True,
257 session=session,
258 )
259 if not updated_tis: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 raise HTTPException(
261 status.HTTP_409_CONFLICT,
262 f"Task id {task_id} is already in {data['new_state']} state",
263 )
265 if data["new_state"] == TaskInstanceState.SUCCESS:
266 _clear_task_state_store_on_success(updated_tis, session)
268 _emit_state_listener_hooks(updated_tis, data["new_state"])
270 return updated_tis
273def _patch_task_group_state(
274 group_id: str,
275 dag_run_id: str,
276 dag: SerializedDAG,
277 body: PatchTaskInstanceBody,
278 data: dict,
279 *,
280 session: Session,
281) -> list[TI]:
282 """Update the state of all task instances in a task group."""
283 updated_tis = dag.set_task_group_state(
284 group_id=group_id,
285 run_id=dag_run_id,
286 state=data["new_state"],
287 upstream=body.include_upstream,
288 downstream=body.include_downstream,
289 future=body.include_future,
290 past=body.include_past,
291 commit=True,
292 session=session,
293 )
294 if not updated_tis:
295 raise HTTPException(
296 status.HTTP_409_CONFLICT,
297 f"All task instances in the group are already in {data['new_state']} state",
298 )
300 if data["new_state"] == TaskInstanceState.SUCCESS:
301 _clear_task_state_store_on_success(updated_tis, session)
303 _emit_state_listener_hooks(updated_tis, data["new_state"])
305 return updated_tis
308def _patch_task_instance_note(
309 task_instance_body: BulkTaskInstanceBody | ClearTaskInstancesBody | PatchTaskInstanceBody,
310 tis: list[TI],
311 user: GetUserDep,
312 update_mask: list[str] | None = None,
313) -> None:
314 for ti in tis:
315 if update_mask or task_instance_body.note is not None: 315 ↛ 314line 315 didn't jump to line 314 because the condition on line 315 was always true
316 if task_instance_body.note == "": 316 ↛ 317line 316 didn't jump to line 317 because the condition on line 316 was never true
317 ti.task_instance_note = None
318 elif ti.task_instance_note is None:
319 ti.note = (task_instance_body.note, user.get_id())
320 else:
321 ti.task_instance_note.content = task_instance_body.note
322 ti.task_instance_note.user_id = user.get_id()
325class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
326 """Service for handling bulk operations on task instances."""
328 def __init__(
329 self,
330 session: Session,
331 request: BulkBody[BulkTaskInstanceBody],
332 dag_id: str,
333 dag_run_id: str,
334 dag_bag: DagBagDep,
335 user: GetUserDep,
336 ):
337 super().__init__(session, request)
338 self.dag_id = dag_id
339 self.dag_run_id = dag_run_id
340 self.dag_bag = dag_bag
341 self.user = user
343 def _extract_task_identifiers(
344 self, entity: str | BulkTaskInstanceBody
345 ) -> tuple[str, str, str, int | None]:
346 """
347 Extract task identifiers from an id or entity object.
349 :param entity: Task identifier as string or BulkTaskInstanceBody object
350 :return: tuple of (dag_id, dag_run_id, task_id, map_index)
351 """
352 if isinstance(entity, str):
353 dag_id = self.dag_id
354 dag_run_id = self.dag_run_id
355 task_id = entity
356 map_index = None
357 else:
358 dag_id = entity.dag_id if entity.dag_id else self.dag_id
359 dag_run_id = entity.dag_run_id if entity.dag_run_id else self.dag_run_id
360 task_id = entity.task_id
361 map_index = entity.map_index
363 return dag_id, dag_run_id, task_id, map_index
365 def _categorize_entities(
366 self,
367 entities: Sequence[str | BulkTaskInstanceBody],
368 results: BulkActionResponse,
369 method: Literal["PUT", "DELETE"],
370 action_name: str,
371 ) -> tuple[set[tuple[str, str, str, int]], set[tuple[str, str, str]]]:
372 """
373 Validate entities and categorize them into specific and all map index update sets.
375 :param entities: Sequence of entities to validate
376 :param results: BulkActionResponse object to track errors
377 :return: tuple of (specific_map_index_task_keys, all_map_index_task_keys)
378 """
379 specific_map_index_task_keys = set()
380 all_map_index_task_keys = set()
381 dag_authorization_cache: dict[str, bool] = {}
383 for entity in entities:
384 dag_id, dag_run_id, task_id, map_index = self._extract_task_identifiers(entity)
386 # Validate that we have specific values, not wildcards
387 if dag_id == "~" or dag_run_id == "~":
388 if isinstance(entity, str):
389 error_msg = f"When using wildcard in path, dag_id and dag_run_id must be specified in BulkTaskInstanceBody object, not as string for task_id: {entity}"
390 else:
391 error_msg = f"When using wildcard in path, dag_id and dag_run_id must be specified in request body for task_id: {entity.task_id}"
392 results.errors.append(
393 {
394 "error": error_msg,
395 "status_code": status.HTTP_400_BAD_REQUEST,
396 }
397 )
398 continue
400 if dag_id not in dag_authorization_cache:
401 team_name = DagModel.get_team_name(dag_id, session=self.session)
402 dag_authorization_cache[dag_id] = get_auth_manager().is_authorized_dag(
403 method=method,
404 access_entity=DagAccessEntity.TASK_INSTANCE,
405 details=DagDetails(id=dag_id, team_name=team_name),
406 user=self.user,
407 )
408 if not dag_authorization_cache[dag_id]: 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true
409 results.errors.append(
410 {
411 "error": f"User is not authorized to {action_name} task instances for DAG '{dag_id}'",
412 "status_code": status.HTTP_403_FORBIDDEN,
413 }
414 )
415 continue
417 # Separate logic for "update all" vs "update specific"
418 if map_index is not None: 418 ↛ 419line 418 didn't jump to line 419 because the condition on line 418 was never true
419 specific_map_index_task_keys.add((dag_id, dag_run_id, task_id, map_index))
420 else:
421 all_map_index_task_keys.add((dag_id, dag_run_id, task_id))
423 return specific_map_index_task_keys, all_map_index_task_keys
425 def _categorize_task_instances(
426 self, task_keys: set[tuple[str, str, str, int]]
427 ) -> tuple[
428 dict[tuple[str, str, str, int], TI], set[tuple[str, str, str, int]], set[tuple[str, str, str, int]]
429 ]:
430 """
431 Categorize the given task_keys into matched and not_found based on existing task instances.
433 :param task_keys: set of task_keys (tuple of dag_id, dag_run_id, task_id, and map_index)
434 :return: tuple of (task_instances_map, matched_task_keys, not_found_task_keys)
435 """
436 # Filter at database level using exact tuple matching instead of fetching all combinations
437 # and filtering in Python
438 task_keys_list = list(task_keys)
439 query = select(TI).where(tuple_(TI.dag_id, TI.run_id, TI.task_id, TI.map_index).in_(task_keys_list))
441 task_instances = self.session.scalars(query).all()
442 task_instances_map = {
443 (ti.dag_id, ti.run_id, ti.task_id, ti.map_index if ti.map_index is not None else -1): ti
444 for ti in task_instances
445 }
446 matched_task_keys = set(task_instances_map.keys())
447 not_found_task_keys = task_keys - matched_task_keys
448 return task_instances_map, matched_task_keys, not_found_task_keys
450 def _perform_update(
451 self,
452 entity: BulkTaskInstanceBody,
453 dag_id: str,
454 dag_run_id: str,
455 task_id: str,
456 map_index: int,
457 results: BulkActionResponse,
458 update_mask: list[str] | None = None,
459 ) -> None:
460 dag, tis, data = _patch_ti_validate_request(
461 dag_id=dag_id,
462 dag_run_id=dag_run_id,
463 task_id=task_id,
464 dag_bag=self.dag_bag,
465 body=entity,
466 session=self.session,
467 map_index=map_index,
468 update_mask=update_mask,
469 )
471 # Apply "note" before "state" so listeners fired inside _patch_task_instance_state() see the updated note.
472 if "note" in data:
473 _patch_task_instance_note(
474 task_instance_body=entity,
475 tis=tis,
476 user=self.user,
477 )
478 if "new_state" in data:
479 _patch_task_instance_state(
480 task_id=task_id,
481 dag_run_id=dag_run_id,
482 dag=dag,
483 task_instance_body=entity,
484 session=self.session,
485 data=data,
486 )
488 results.success.append(f"{dag_id}.{dag_run_id}.{task_id}[{map_index}]")
490 def handle_bulk_create(
491 self, action: BulkCreateAction[BulkTaskInstanceBody], results: BulkActionResponse
492 ) -> None:
493 results.errors.append(
494 {
495 "error": "Task instances bulk create is not supported",
496 "status_code": status.HTTP_405_METHOD_NOT_ALLOWED,
497 }
498 )
500 def handle_bulk_update(
501 self, action: BulkUpdateAction[BulkTaskInstanceBody], results: BulkActionResponse
502 ) -> None:
503 """Bulk Update Task Instances."""
504 # Validate and categorize entities into specific and all map index update sets
505 update_specific_map_index_task_keys, update_all_map_index_task_keys = self._categorize_entities(
506 action.entities, results, method="PUT", action_name=action.action.value
507 )
509 try:
510 specific_entity_map = {
511 self._extract_task_identifiers(entity): entity
512 for entity in action.entities
513 if entity.map_index is not None
514 }
515 all_map_entity_map = {
516 self._extract_task_identifiers(entity)[:3]: entity
517 for entity in action.entities
518 if entity.map_index is None
519 }
521 # Handle updates for specific map_index task instances
522 if update_specific_map_index_task_keys: 522 ↛ 523line 522 didn't jump to line 523 because the condition on line 522 was never true
523 _, matched_task_keys, not_found_task_keys = self._categorize_task_instances(
524 update_specific_map_index_task_keys
525 )
527 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_task_keys:
528 not_found_task_ids = [
529 {"dag_id": dag_id, "dag_run_id": run_id, "task_id": task_id, "map_index": map_index}
530 for dag_id, run_id, task_id, map_index in not_found_task_keys
531 ]
532 raise HTTPException(
533 status_code=status.HTTP_404_NOT_FOUND,
534 detail=f"The task instances with these identifiers: {not_found_task_ids} were not found",
535 )
537 for dag_id, dag_run_id, task_id, map_index in matched_task_keys:
538 entity = specific_entity_map.get((dag_id, dag_run_id, task_id, map_index))
540 if entity is not None:
541 self._perform_update(
542 dag_id=dag_id,
543 dag_run_id=dag_run_id,
544 task_id=task_id,
545 map_index=map_index,
546 entity=entity,
547 results=results,
548 update_mask=action.update_mask,
549 )
551 # Handle updates for all map indexes
552 if update_all_map_index_task_keys: 552 ↛ 553line 552 didn't jump to line 553 because the condition on line 552 was never true
553 all_dag_ids = {dag_id for dag_id, _, _ in update_all_map_index_task_keys}
554 all_run_ids = {run_id for _, run_id, _ in update_all_map_index_task_keys}
555 all_task_ids = {task_id for _, _, task_id in update_all_map_index_task_keys}
557 batch_task_instances = self.session.scalars(
558 select(TI).where(
559 TI.dag_id.in_(all_dag_ids),
560 TI.run_id.in_(all_run_ids),
561 TI.task_id.in_(all_task_ids),
562 )
563 ).all()
565 # Group task instances by (dag_id, run_id, task_id)
566 task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
567 for ti in batch_task_instances:
568 key = (ti.dag_id, ti.run_id, ti.task_id)
569 task_instances_by_key.setdefault(key, []).append(ti)
571 for dag_id, run_id, task_id in update_all_map_index_task_keys:
572 all_task_instances = task_instances_by_key.get((dag_id, run_id, task_id), [])
574 if (
575 not all_task_instances
576 and action.action_on_non_existence == BulkActionNotOnExistence.FAIL
577 ):
578 raise HTTPException(
579 status_code=status.HTTP_404_NOT_FOUND,
580 detail=f"No task instances found for dag_id: {dag_id}, run_id: {run_id}, task_id: {task_id}",
581 )
583 entity = all_map_entity_map.get((dag_id, run_id, task_id))
585 if entity is not None:
586 for ti in all_task_instances:
587 self._perform_update(
588 dag_id=dag_id,
589 dag_run_id=run_id,
590 task_id=task_id,
591 map_index=ti.map_index if ti.map_index is not None else -1,
592 entity=entity,
593 results=results,
594 update_mask=action.update_mask,
595 )
597 except ValidationError as e:
598 results.errors.append({"error": f"{e.errors()}"})
599 except HTTPException as e:
600 results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})
602 def handle_bulk_delete(
603 self, action: BulkDeleteAction[BulkTaskInstanceBody], results: BulkActionResponse
604 ) -> None:
605 """Bulk delete task instances."""
606 # Validate and categorize entities into specific and all map index delete sets
607 delete_specific_map_index_task_keys, delete_all_map_index_task_keys = self._categorize_entities(
608 action.entities, results, method="DELETE", action_name=action.action.value
609 )
611 try:
612 # Handle deletion of specific (dag_id, dag_run_id, task_id, map_index) tuples
613 if delete_specific_map_index_task_keys: 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true
614 task_instances_map, matched_task_keys, not_found_task_keys = self._categorize_task_instances(
615 delete_specific_map_index_task_keys
616 )
617 not_found_task_ids = [
618 {"dag_id": dag_id, "dag_run_id": run_id, "task_id": task_id, "map_index": map_index}
619 for dag_id, run_id, task_id, map_index in not_found_task_keys
620 ]
622 if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_task_keys:
623 raise HTTPException(
624 status_code=status.HTTP_404_NOT_FOUND,
625 detail=f"The task instances with these identifiers: {not_found_task_ids} were not found",
626 )
628 for task_key in matched_task_keys:
629 dag_id, run_id, task_id, map_index = task_key
630 self.session.delete(task_instances_map[task_key])
631 results.success.append(f"{dag_id}.{run_id}.{task_id}[{map_index}]")
633 # Handle deletion of all map indexes for certain (dag_id, dag_run_id, task_id) tuples
634 if delete_all_map_index_task_keys:
635 all_dag_ids = {dag_id for dag_id, _, _ in delete_all_map_index_task_keys}
636 all_run_ids = {run_id for _, run_id, _ in delete_all_map_index_task_keys}
637 all_task_ids = {task_id for _, _, task_id in delete_all_map_index_task_keys}
639 batch_task_instances = self.session.scalars(
640 select(TI).where(
641 TI.dag_id.in_(all_dag_ids),
642 TI.run_id.in_(all_run_ids),
643 TI.task_id.in_(all_task_ids),
644 )
645 ).all()
647 # Group task instances by (dag_id, run_id, task_id) for efficient lookup
648 task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
649 for ti in batch_task_instances: 649 ↛ 650line 649 didn't jump to line 650 because the loop on line 649 never started
650 key = (ti.dag_id, ti.run_id, ti.task_id)
651 task_instances_by_key.setdefault(key, []).append(ti)
653 for dag_id, run_id, task_id in delete_all_map_index_task_keys:
654 all_task_instances = task_instances_by_key.get((dag_id, run_id, task_id), [])
656 if (
657 not all_task_instances
658 and action.action_on_non_existence == BulkActionNotOnExistence.FAIL
659 ):
660 raise HTTPException(
661 status_code=status.HTTP_404_NOT_FOUND,
662 detail=f"No task instances found for dag_id: {dag_id}, run_id: {run_id}, task_id: {task_id}",
663 )
665 for ti in all_task_instances: 665 ↛ 666line 665 didn't jump to line 666 because the loop on line 665 never started
666 self.session.delete(ti)
667 results.success.append(f"{dag_id}.{run_id}.{task_id}[{ti.map_index}]")
669 except HTTPException as e:
670 results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})