This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f8e29160117 Add TI bulk actions endpoint (#50443)
f8e29160117 is described below
commit f8e29160117a91e2443b7f3f99cabc2485d10271
Author: Guan Ming(Wesley) Chiu <[email protected]>
AuthorDate: Fri May 30 00:35:52 2025 +0800
Add TI bulk actions endpoint (#50443)
---
.../core_api/datamodels/task_instances.py | 7 +
.../core_api/openapi/v2-rest-api-generated.yaml | 175 ++++++++++
.../core_api/routes/public/task_instances.py | 75 ++---
.../core_api/services/public/task_instances.py | 271 ++++++++++++++++
.../src/airflow/ui/openapi-gen/queries/common.ts | 3 +
.../src/airflow/ui/openapi-gen/queries/queries.ts | 44 +++
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 174 ++++++++++
.../ui/openapi-gen/requests/services.gen.ts | 30 ++
.../airflow/ui/openapi-gen/requests/types.gen.ts | 87 +++++
.../core_api/routes/public/test_task_instances.py | 357 +++++++++++++++++++++
.../src/airflowctl/api/datamodels/generated.py | 79 ++++-
11 files changed, 1247 insertions(+), 55 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
index 2eba2eeecf8..392fcc598df 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
@@ -224,3 +224,10 @@ class PatchTaskInstanceBody(StrictBaseModel):
if ns not in valid_states:
raise ValueError(f"'{ns}' is not one of {valid_states}")
return ns
+
+
+class BulkTaskInstanceBody(PatchTaskInstanceBody, StrictBaseModel):
+ """Request body for bulk update, and delete task instances."""
+
+ task_id: str
+ map_index: int | None = None
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 2da6e38a273..3fbe6692b09 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -5568,6 +5568,58 @@ paths:
application/json:
schema:
$ref: '#/components/schemas/HTTPValidationError'
+ patch:
+ tags:
+ - Task Instance
+ summary: Bulk Task Instances
+ description: Bulk update, and delete task instances.
+ operationId: bulk_task_instances
+ security:
+ - OAuth2PasswordBearer: []
+ parameters:
+ - name: dag_id
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Dag Id
+ - name: dag_run_id
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Dag Run Id
+ requestBody:
+ required: true
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/BulkBody_BulkTaskInstanceBody_'
+ responses:
+ '200':
+ description: Successful Response
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/BulkResponse'
+ '401':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/HTTPExceptionResponse'
+ description: Unauthorized
+ '403':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/HTTPExceptionResponse'
+ description: Forbidden
+ '422':
+ description: Validation Error
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/HTTPValidationError'
/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/list:
post:
tags:
@@ -7389,6 +7441,21 @@ components:
This structure helps users understand which key actions succeeded and
which
failed.'
+ BulkBody_BulkTaskInstanceBody_:
+ properties:
+ actions:
+ items:
+ oneOf:
+ - $ref:
'#/components/schemas/BulkCreateAction_BulkTaskInstanceBody_'
+ - $ref:
'#/components/schemas/BulkUpdateAction_BulkTaskInstanceBody_'
+ - $ref:
'#/components/schemas/BulkDeleteAction_BulkTaskInstanceBody_'
+ type: array
+ title: Actions
+ additionalProperties: false
+ type: object
+ required:
+ - actions
+ title: BulkBody[BulkTaskInstanceBody]
BulkBody_ConnectionBody_:
properties:
actions:
@@ -7434,6 +7501,28 @@ components:
required:
- actions
title: BulkBody[VariableBody]
+ BulkCreateAction_BulkTaskInstanceBody_:
+ properties:
+ action:
+ type: string
+ const: create
+ title: Action
+ description: The action to be performed on the entities.
+ entities:
+ items:
+ $ref: '#/components/schemas/BulkTaskInstanceBody'
+ type: array
+ title: Entities
+ description: A list of entities to be created.
+ action_on_existence:
+ $ref: '#/components/schemas/BulkActionOnExistence'
+ default: fail
+ additionalProperties: false
+ type: object
+ required:
+ - action
+ - entities
+ title: BulkCreateAction[BulkTaskInstanceBody]
BulkCreateAction_ConnectionBody_:
properties:
action:
@@ -7500,6 +7589,28 @@ components:
- action
- entities
title: BulkCreateAction[VariableBody]
+ BulkDeleteAction_BulkTaskInstanceBody_:
+ properties:
+ action:
+ type: string
+ const: delete
+ title: Action
+ description: The action to be performed on the entities.
+ entities:
+ items:
+ type: string
+ type: array
+ title: Entities
+ description: A list of entity id/key to be deleted.
+ action_on_non_existence:
+ $ref: '#/components/schemas/BulkActionNotOnExistence'
+ default: fail
+ additionalProperties: false
+ type: object
+ required:
+ - action
+ - entities
+ title: BulkDeleteAction[BulkTaskInstanceBody]
BulkDeleteAction_ConnectionBody_:
properties:
action:
@@ -7599,6 +7710,70 @@ components:
Fields are populated in the response only if the respective action was
part
of the request, else are set None.'
+ BulkTaskInstanceBody:
+ properties:
+ new_state:
+ anyOf:
+ - $ref: '#/components/schemas/TaskInstanceState'
+ - type: 'null'
+ note:
+ anyOf:
+ - type: string
+ maxLength: 1000
+ - type: 'null'
+ title: Note
+ include_upstream:
+ type: boolean
+ title: Include Upstream
+ default: false
+ include_downstream:
+ type: boolean
+ title: Include Downstream
+ default: false
+ include_future:
+ type: boolean
+ title: Include Future
+ default: false
+ include_past:
+ type: boolean
+ title: Include Past
+ default: false
+ task_id:
+ type: string
+ title: Task Id
+ map_index:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Map Index
+ additionalProperties: false
+ type: object
+ required:
+ - task_id
+ title: BulkTaskInstanceBody
+ description: Request body for bulk update, and delete task instances.
+ BulkUpdateAction_BulkTaskInstanceBody_:
+ properties:
+ action:
+ type: string
+ const: update
+ title: Action
+ description: The action to be performed on the entities.
+ entities:
+ items:
+ $ref: '#/components/schemas/BulkTaskInstanceBody'
+ type: array
+ title: Entities
+ description: A list of entities to be updated.
+ action_on_non_existence:
+ $ref: '#/components/schemas/BulkActionNotOnExistence'
+ default: fail
+ additionalProperties: false
+ type: object
+ required:
+ - action
+ - entities
+ title: BulkUpdateAction[BulkTaskInstanceBody]
BulkUpdateAction_ConnectionBody_:
properties:
action:
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index d14ab18c95c..79a0f977916 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -21,8 +21,6 @@ from typing import Annotated, Literal, cast
import structlog
from fastapi import Depends, HTTPException, Query, status
-from fastapi.exceptions import RequestValidationError
-from pydantic import ValidationError
from sqlalchemy import or_, select
from sqlalchemy.orm import joinedload
from sqlalchemy.sql.selectable import Select
@@ -51,7 +49,9 @@ from airflow.api_fastapi.common.parameters import (
float_range_filter_factory,
)
from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.datamodels.common import BulkBody,
BulkResponse
from airflow.api_fastapi.core_api.datamodels.task_instances import (
+ BulkTaskInstanceBody,
ClearTaskInstancesBody,
PatchTaskInstanceBody,
TaskDependencyCollectionResponse,
@@ -63,11 +63,14 @@ from airflow.api_fastapi.core_api.datamodels.task_instances
import (
)
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import GetUserDep,
ReadableTIFilterDep, requires_access_dag
+from airflow.api_fastapi.core_api.services.public.task_instances import (
+ BulkTaskInstanceService,
+ _patch_ti_validate_request,
+)
from airflow.api_fastapi.logging.decorators import action_logging
from airflow.exceptions import TaskNotFound
from airflow.listeners.listener import get_listener_manager
from airflow.models import Base, DagRun
-from airflow.models.dag import DAG
from airflow.models.taskinstance import TaskInstance as TI,
clear_task_instances
from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH
from airflow.ti_deps.dep_context import DepContext
@@ -728,54 +731,6 @@ def post_clear_task_instances(
)
-def _patch_ti_validate_request(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
- dag_bag: DagBagDep,
- body: PatchTaskInstanceBody,
- session: SessionDep,
- map_index: int | None = -1,
- update_mask: list[str] | None = Query(None),
-) -> tuple[DAG, list[TI], dict]:
- dag = dag_bag.get_dag(dag_id)
- if not dag:
- raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not
found")
-
- if not dag.has_task(task_id):
- raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not
found in DAG '{dag_id}'")
-
- query = (
- select(TI)
- .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id ==
task_id)
- .join(TI.dag_run)
- .options(joinedload(TI.rendered_task_instance_fields))
- )
- if map_index is not None:
- query = query.where(TI.map_index == map_index)
- else:
- query = query.order_by(TI.map_index)
-
- tis = session.scalars(query).all()
-
- err_msg_404 = (
- 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",
- )
- if len(tis) == 0:
- raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
-
- fields_to_update = body.model_fields_set
- if update_mask:
- fields_to_update = fields_to_update.intersection(update_mask)
- else:
- try:
- PatchTaskInstanceBody.model_validate(body)
- except ValidationError as e:
- raise RequestValidationError(errors=e.errors())
-
- return dag, list(tis), body.model_dump(include=fields_to_update,
by_alias=True)
-
-
@task_instances_router.patch(
task_instances_prefix + "/{task_id}/dry_run",
responses=create_openapi_http_exception_doc(
@@ -835,6 +790,24 @@ def patch_task_instance_dry_run(
)
+@task_instances_router.patch(
+ task_instances_prefix,
+ dependencies=[Depends(requires_access_dag(method="PUT",
access_entity=DagAccessEntity.TASK_INSTANCE))],
+)
+def bulk_task_instances(
+ request: BulkBody[BulkTaskInstanceBody],
+ session: SessionDep,
+ dag_id: str,
+ dag_bag: DagBagDep,
+ dag_run_id: str,
+ user: GetUserDep,
+) -> BulkResponse:
+ """Bulk update, and delete task instances."""
+ return BulkTaskInstanceService(
+ session=session, request=request, dag_id=dag_id,
dag_run_id=dag_run_id, dag_bag=dag_bag, user=user
+ ).handle_request()
+
+
@task_instances_router.patch(
task_instances_prefix + "/{task_id}",
responses=create_openapi_http_exception_doc(
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
new file mode 100644
index 00000000000..a99e802278f
--- /dev/null
+++
b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
@@ -0,0 +1,271 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import structlog
+from fastapi import HTTPException, Query, status
+from fastapi.exceptions import RequestValidationError
+from pydantic import ValidationError
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload
+from sqlalchemy.orm.session import Session
+
+from airflow.api_fastapi.common.dagbag import DagBagDep
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.core_api.datamodels.common import (
+ BulkActionNotOnExistence,
+ BulkActionResponse,
+ BulkBody,
+ BulkCreateAction,
+ BulkDeleteAction,
+ BulkUpdateAction,
+)
+from airflow.api_fastapi.core_api.datamodels.task_instances import
BulkTaskInstanceBody, PatchTaskInstanceBody
+from airflow.api_fastapi.core_api.security import GetUserDep
+from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.listeners.listener import get_listener_manager
+from airflow.models.dag import DAG
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.utils.state import TaskInstanceState
+
+log = structlog.get_logger(__name__)
+
+
+def _patch_ti_validate_request(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ dag_bag: DagBagDep,
+ body: PatchTaskInstanceBody,
+ session: SessionDep,
+ map_index: int | None = -1,
+ update_mask: list[str] | None = Query(None),
+) -> tuple[DAG, list[TI], dict]:
+ dag = dag_bag.get_dag(dag_id)
+ if not dag:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not
found")
+
+ if not dag.has_task(task_id):
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not
found in DAG '{dag_id}'")
+
+ query = (
+ select(TI)
+ .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id ==
task_id)
+ .join(TI.dag_run)
+ .options(joinedload(TI.rendered_task_instance_fields))
+ )
+ if map_index is not None:
+ query = query.where(TI.map_index == map_index)
+ else:
+ query = query.order_by(TI.map_index)
+
+ tis = session.scalars(query).all()
+
+ err_msg_404 = (
+ 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",
+ )
+ if len(tis) == 0:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
+
+ fields_to_update = body.model_fields_set
+ if update_mask:
+ fields_to_update = fields_to_update.intersection(update_mask)
+ else:
+ try:
+ PatchTaskInstanceBody.model_validate(body)
+ except ValidationError as e:
+ raise RequestValidationError(errors=e.errors())
+
+ return dag, list(tis), body.model_dump(include=fields_to_update,
by_alias=True)
+
+
+class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
+ """Service for handling bulk operations on task instances."""
+
+ def __init__(
+ self,
+ session: Session,
+ request: BulkBody[BulkTaskInstanceBody],
+ dag_id: str,
+ dag_run_id: str,
+ dag_bag: DagBagDep,
+ user: GetUserDep,
+ ):
+ super().__init__(session, request)
+ self.dag_id = dag_id
+ self.dag_run_id = dag_run_id
+ self.dag_bag = dag_bag
+ self.user = user
+
+ def categorize_task_instances(
+ self, task_ids: set[tuple[str, int]]
+ ) -> tuple[dict[tuple[str, int], TI], set[tuple[str, int]], set[tuple[str,
int]]]:
+ """
+ Categorize the given task_ids into matched_task_keys and
not_found_task_keys based on existing task_ids.
+
+ :param task_ids: set of task_ids
+ :return: tuple of (task_instances_map, matched_task_keys,
not_found_task_keys)
+ """
+ query = select(TI).where(
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.dag_run_id,
+ TI.task_id.in_([task_id for task_id, _ in task_ids]),
+ )
+ task_instances = self.session.scalars(query).all()
+ task_instances_map = {
+ (ti.task_id, ti.map_index if ti.map_index is not None else -1): ti
for ti in task_instances
+ }
+ matched_task_keys = {(task_id, map_index) for (task_id, map_index) in
task_instances_map.keys()}
+ not_found_task_keys = {(task_id, map_index) for task_id, map_index in
task_ids} - matched_task_keys
+ return task_instances_map, matched_task_keys, not_found_task_keys
+
+ def _patch_task_instance_state(
+ self,
+ dag: DAG,
+ task_instance_body: BulkTaskInstanceBody,
+ data: dict,
+ ) -> None:
+ map_indexes = None if task_instance_body.map_index is None else
[task_instance_body.map_index]
+
+ updated_tis = dag.set_task_instance_state(
+ task_id=task_instance_body.task_id,
+ run_id=self.dag_run_id,
+ map_indexes=map_indexes,
+ state=data["new_state"],
+ upstream=task_instance_body.include_upstream,
+ downstream=task_instance_body.include_downstream,
+ future=task_instance_body.include_future,
+ past=task_instance_body.include_past,
+ commit=True,
+ session=self.session,
+ )
+ if not updated_tis:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ f"Task id {task_instance_body.task_id} is already in
{data['new_state']} state",
+ )
+ for ti in updated_tis:
+ try:
+ if data["new_state"] == TaskInstanceState.SUCCESS:
+ get_listener_manager().hook.on_task_instance_success(
+ previous_state=None, task_instance=ti
+ )
+ elif data["new_state"] == TaskInstanceState.FAILED:
+ get_listener_manager().hook.on_task_instance_failed(
+ previous_state=None,
+ task_instance=ti,
+ error=f"TaskInstance's state was manually set to
`{TaskInstanceState.FAILED}`.",
+ )
+ except Exception:
+ log.exception("error calling listener")
+
+ def _patch_task_instance_note(self, task_instance_body:
BulkTaskInstanceBody, tis: list[TI]) -> None:
+ for ti in tis:
+ if task_instance_body.note is not None:
+ if ti.task_instance_note is None:
+ ti.note = (task_instance_body.note, self.user.get_id())
+ else:
+ ti.task_instance_note.content = task_instance_body.note
+ ti.task_instance_note.user_id = self.user.get_id()
+
+ def handle_bulk_create(
+ self, action: BulkCreateAction[BulkTaskInstanceBody], results:
BulkActionResponse
+ ) -> None:
+ results.errors.append(
+ {
+ "error": "Task instances bulk create is not supported",
+ "status_code": status.HTTP_405_METHOD_NOT_ALLOWED,
+ }
+ )
+
+ def handle_bulk_update(
+ self, action: BulkUpdateAction[BulkTaskInstanceBody], results:
BulkActionResponse
+ ) -> None:
+ """Bulk Update Task Instances."""
+ to_update_task_keys = {
+ (task_instance.task_id, task_instance.map_index if
task_instance.map_index is not None else -1)
+ for task_instance in action.entities
+ }
+ _, _, not_found_task_keys =
self.categorize_task_instances(to_update_task_keys)
+
+ try:
+ for task_instance_body in action.entities:
+ task_key = (
+ task_instance_body.task_id,
+ task_instance_body.map_index if
task_instance_body.map_index is not None else -1,
+ )
+
+ if task_key in not_found_task_keys:
+ if action.action_on_non_existence ==
BulkActionNotOnExistence.FAIL:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=f"The Task Instance with dag_id:
`{self.dag_id}`, run_id: `{self.dag_run_id}`, task_id:
`{task_instance_body.task_id}` and map_index: `{task_instance_body.map_index}`
was not found",
+ )
+ if action.action_on_non_existence ==
BulkActionNotOnExistence.SKIP:
+ continue
+
+ dag, tis, data = _patch_ti_validate_request(
+ dag_id=self.dag_id,
+ dag_run_id=self.dag_run_id,
+ task_id=task_instance_body.task_id,
+ dag_bag=self.dag_bag,
+ body=task_instance_body,
+ session=self.session,
+ map_index=task_instance_body.map_index,
+ update_mask=None,
+ )
+
+ for key, _ in data.items():
+ if key == "new_state":
+ self._patch_task_instance_state(
+ dag=dag,
+ task_instance_body=task_instance_body,
+ data=data,
+ )
+ elif key == "note":
+
self._patch_task_instance_note(task_instance_body=task_instance_body, tis=tis)
+
+ results.success.append(task_instance_body.task_id)
+ except ValidationError as e:
+ results.errors.append({"error": f"{e.errors()}"})
+ except HTTPException as e:
+ results.errors.append({"error": f"{e.detail}", "status_code":
e.status_code})
+
+ def handle_bulk_delete(
+ self, action: BulkDeleteAction[BulkTaskInstanceBody], results:
BulkActionResponse
+ ) -> None:
+ """Bulk delete task instances."""
+ to_delete_task_keys = set((task_id, -1) for task_id in action.entities)
+ _, matched_task_keys, not_found_task_keys =
self.categorize_task_instances(to_delete_task_keys)
+ not_found_task_ids = [task_id for task_id, _ in not_found_task_keys]
+
+ try:
+ if action.action_on_non_existence == BulkActionNotOnExistence.FAIL
and not_found_task_keys:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=f"The task instances with these task_ids:
{not_found_task_ids} were not found",
+ )
+
+ for task_id, _ in matched_task_keys:
+ existing_task_instance =
self.session.scalar(select(TI).where(TI.task_id == task_id).limit(1))
+ if existing_task_instance:
+ self.session.delete(existing_task_instance)
+ results.success.append(task_id)
+
+ except HTTPException as e:
+ results.errors.append({"error": f"{e.detail}", "status_code":
e.status_code})
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
index bafbac9fb23..8410dde0ba3 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -1912,6 +1912,9 @@ export type
TaskInstanceServicePatchTaskInstanceMutationResult = Awaited<
export type TaskInstanceServicePatchTaskInstanceByMapIndexMutationResult =
Awaited<
ReturnType<typeof TaskInstanceService.patchTaskInstanceByMapIndex>
>;
+export type TaskInstanceServiceBulkTaskInstancesMutationResult = Awaited<
+ ReturnType<typeof TaskInstanceService.bulkTaskInstances>
+>;
export type TaskInstanceServicePatchTaskInstanceDryRunByMapIndexMutationResult
= Awaited<
ReturnType<typeof TaskInstanceService.patchTaskInstanceDryRunByMapIndex>
>;
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index 19c242bd493..8e2e3a810a6 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -37,6 +37,7 @@ import {
} from "../requests/services.gen";
import {
BackfillPostBody,
+ BulkBody_BulkTaskInstanceBody_,
BulkBody_ConnectionBody_,
BulkBody_PoolBody_,
BulkBody_VariableBody_,
@@ -4174,6 +4175,49 @@ export const
useTaskInstanceServicePatchTaskInstanceByMapIndex = <
}) as unknown as Promise<TData>,
...options,
});
+/**
+ * Bulk Task Instances
+ * Bulk update, and delete task instances.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.dagRunId
+ * @param data.requestBody
+ * @returns BulkResponse Successful Response
+ * @throws ApiError
+ */
+export const useTaskInstanceServiceBulkTaskInstances = <
+ TData = Common.TaskInstanceServiceBulkTaskInstancesMutationResult,
+ TError = unknown,
+ TContext = unknown,
+>(
+ options?: Omit<
+ UseMutationOptions<
+ TData,
+ TError,
+ {
+ dagId: string;
+ dagRunId: string;
+ requestBody: BulkBody_BulkTaskInstanceBody_;
+ },
+ TContext
+ >,
+ "mutationFn"
+ >,
+) =>
+ useMutation<
+ TData,
+ TError,
+ {
+ dagId: string;
+ dagRunId: string;
+ requestBody: BulkBody_BulkTaskInstanceBody_;
+ },
+ TContext
+ >({
+ mutationFn: ({ dagId, dagRunId, requestBody }) =>
+ TaskInstanceService.bulkTaskInstances({ dagId, dagRunId, requestBody })
as unknown as Promise<TData>,
+ ...options,
+ });
/**
* Patch Task Instance Dry Run
* Update a task instance dry_run mode.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index e41c93b11f4..7ffda75f426 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -582,6 +582,32 @@ The response includes a list of successful keys and any
errors encountered durin
This structure helps users understand which key actions succeeded and which
failed.`,
} as const;
+export const $BulkBody_BulkTaskInstanceBody_ = {
+ properties: {
+ actions: {
+ items: {
+ oneOf: [
+ {
+ $ref:
"#/components/schemas/BulkCreateAction_BulkTaskInstanceBody_",
+ },
+ {
+ $ref:
"#/components/schemas/BulkUpdateAction_BulkTaskInstanceBody_",
+ },
+ {
+ $ref:
"#/components/schemas/BulkDeleteAction_BulkTaskInstanceBody_",
+ },
+ ],
+ },
+ type: "array",
+ title: "Actions",
+ },
+ },
+ additionalProperties: false,
+ type: "object",
+ required: ["actions"],
+ title: "BulkBody[BulkTaskInstanceBody]",
+} as const;
+
export const $BulkBody_ConnectionBody_ = {
properties: {
actions: {
@@ -660,6 +686,33 @@ export const $BulkBody_VariableBody_ = {
title: "BulkBody[VariableBody]",
} as const;
+export const $BulkCreateAction_BulkTaskInstanceBody_ = {
+ properties: {
+ action: {
+ type: "string",
+ const: "create",
+ title: "Action",
+ description: "The action to be performed on the entities.",
+ },
+ entities: {
+ items: {
+ $ref: "#/components/schemas/BulkTaskInstanceBody",
+ },
+ type: "array",
+ title: "Entities",
+ description: "A list of entities to be created.",
+ },
+ action_on_existence: {
+ $ref: "#/components/schemas/BulkActionOnExistence",
+ default: "fail",
+ },
+ },
+ additionalProperties: false,
+ type: "object",
+ required: ["action", "entities"],
+ title: "BulkCreateAction[BulkTaskInstanceBody]",
+} as const;
+
export const $BulkCreateAction_ConnectionBody_ = {
properties: {
action: {
@@ -741,6 +794,33 @@ export const $BulkCreateAction_VariableBody_ = {
title: "BulkCreateAction[VariableBody]",
} as const;
+export const $BulkDeleteAction_BulkTaskInstanceBody_ = {
+ properties: {
+ action: {
+ type: "string",
+ const: "delete",
+ title: "Action",
+ description: "The action to be performed on the entities.",
+ },
+ entities: {
+ items: {
+ type: "string",
+ },
+ type: "array",
+ title: "Entities",
+ description: "A list of entity id/key to be deleted.",
+ },
+ action_on_non_existence: {
+ $ref: "#/components/schemas/BulkActionNotOnExistence",
+ default: "fail",
+ },
+ },
+ additionalProperties: false,
+ type: "object",
+ required: ["action", "entities"],
+ title: "BulkDeleteAction[BulkTaskInstanceBody]",
+} as const;
+
export const $BulkDeleteAction_ConnectionBody_ = {
properties: {
action: {
@@ -867,6 +947,100 @@ Each action (if requested) is represented as a field
containing details about su
Fields are populated in the response only if the respective action was part of
the request, else are set None.`,
} as const;
+export const $BulkTaskInstanceBody = {
+ properties: {
+ new_state: {
+ anyOf: [
+ {
+ $ref: "#/components/schemas/TaskInstanceState",
+ },
+ {
+ type: "null",
+ },
+ ],
+ },
+ note: {
+ anyOf: [
+ {
+ type: "string",
+ maxLength: 1000,
+ },
+ {
+ type: "null",
+ },
+ ],
+ title: "Note",
+ },
+ include_upstream: {
+ type: "boolean",
+ title: "Include Upstream",
+ default: false,
+ },
+ include_downstream: {
+ type: "boolean",
+ title: "Include Downstream",
+ default: false,
+ },
+ include_future: {
+ type: "boolean",
+ title: "Include Future",
+ default: false,
+ },
+ include_past: {
+ type: "boolean",
+ title: "Include Past",
+ default: false,
+ },
+ task_id: {
+ type: "string",
+ title: "Task Id",
+ },
+ map_index: {
+ anyOf: [
+ {
+ type: "integer",
+ },
+ {
+ type: "null",
+ },
+ ],
+ title: "Map Index",
+ },
+ },
+ additionalProperties: false,
+ type: "object",
+ required: ["task_id"],
+ title: "BulkTaskInstanceBody",
+ description: "Request body for bulk update, and delete task instances.",
+} as const;
+
+export const $BulkUpdateAction_BulkTaskInstanceBody_ = {
+ properties: {
+ action: {
+ type: "string",
+ const: "update",
+ title: "Action",
+ description: "The action to be performed on the entities.",
+ },
+ entities: {
+ items: {
+ $ref: "#/components/schemas/BulkTaskInstanceBody",
+ },
+ type: "array",
+ title: "Entities",
+ description: "A list of entities to be updated.",
+ },
+ action_on_non_existence: {
+ $ref: "#/components/schemas/BulkActionNotOnExistence",
+ default: "fail",
+ },
+ },
+ additionalProperties: false,
+ type: "object",
+ required: ["action", "entities"],
+ title: "BulkUpdateAction[BulkTaskInstanceBody]",
+} as const;
+
export const $BulkUpdateAction_ConnectionBody_ = {
properties: {
action: {
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index d5b88a7bf2b..cf819bbc4ee 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -134,6 +134,8 @@ import type {
PatchTaskInstanceByMapIndexResponse,
GetTaskInstancesData,
GetTaskInstancesResponse,
+ BulkTaskInstancesData,
+ BulkTaskInstancesResponse,
GetTaskInstancesBatchData,
GetTaskInstancesBatchResponse,
GetTaskInstanceTryDetailsData,
@@ -2337,6 +2339,34 @@ export class TaskInstanceService {
});
}
+ /**
+ * Bulk Task Instances
+ * Bulk update, and delete task instances.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.dagRunId
+ * @param data.requestBody
+ * @returns BulkResponse Successful Response
+ * @throws ApiError
+ */
+ public static bulkTaskInstances(data: BulkTaskInstancesData):
CancelablePromise<BulkTaskInstancesResponse> {
+ return __request(OpenAPI, {
+ method: "PATCH",
+ url: "/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances",
+ path: {
+ dag_id: data.dagId,
+ dag_run_id: data.dagRunId,
+ },
+ body: data.requestBody,
+ mediaType: "application/json",
+ errors: {
+ 401: "Unauthorized",
+ 403: "Forbidden",
+ 422: "Validation Error",
+ },
+ });
+ }
+
/**
* Get Task Instances Batch
* Get list of task instances.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 200008de784..924dfa45916 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -173,6 +173,14 @@ export type BulkActionResponse = {
}>;
};
+export type BulkBody_BulkTaskInstanceBody_ = {
+ actions: Array<
+ | BulkCreateAction_BulkTaskInstanceBody_
+ | BulkUpdateAction_BulkTaskInstanceBody_
+ | BulkDeleteAction_BulkTaskInstanceBody_
+ >;
+};
+
export type BulkBody_ConnectionBody_ = {
actions: Array<
BulkCreateAction_ConnectionBody_ | BulkUpdateAction_ConnectionBody_ |
BulkDeleteAction_ConnectionBody_
@@ -189,6 +197,18 @@ export type BulkBody_VariableBody_ = {
>;
};
+export type BulkCreateAction_BulkTaskInstanceBody_ = {
+ /**
+ * The action to be performed on the entities.
+ */
+ action: "create";
+ /**
+ * A list of entities to be created.
+ */
+ entities: Array<BulkTaskInstanceBody>;
+ action_on_existence?: BulkActionOnExistence;
+};
+
export type BulkCreateAction_ConnectionBody_ = {
/**
* The action to be performed on the entities.
@@ -225,6 +245,18 @@ export type BulkCreateAction_VariableBody_ = {
action_on_existence?: BulkActionOnExistence;
};
+export type BulkDeleteAction_BulkTaskInstanceBody_ = {
+ /**
+ * The action to be performed on the entities.
+ */
+ action: "delete";
+ /**
+ * A list of entity id/key to be deleted.
+ */
+ entities: Array<string>;
+ action_on_non_existence?: BulkActionNotOnExistence;
+};
+
export type BulkDeleteAction_ConnectionBody_ = {
/**
* The action to be performed on the entities.
@@ -283,6 +315,32 @@ export type BulkResponse = {
delete?: BulkActionResponse | null;
};
+/**
+ * Request body for bulk update, and delete task instances.
+ */
+export type BulkTaskInstanceBody = {
+ new_state?: TaskInstanceState | null;
+ note?: string | null;
+ include_upstream?: boolean;
+ include_downstream?: boolean;
+ include_future?: boolean;
+ include_past?: boolean;
+ task_id: string;
+ map_index?: number | null;
+};
+
+export type BulkUpdateAction_BulkTaskInstanceBody_ = {
+ /**
+ * The action to be performed on the entities.
+ */
+ action: "update";
+ /**
+ * A list of entities to be updated.
+ */
+ entities: Array<BulkTaskInstanceBody>;
+ action_on_non_existence?: BulkActionNotOnExistence;
+};
+
export type BulkUpdateAction_ConnectionBody_ = {
/**
* The action to be performed on the entities.
@@ -2406,6 +2464,14 @@ export type GetTaskInstancesData = {
export type GetTaskInstancesResponse = TaskInstanceCollectionResponse;
+export type BulkTaskInstancesData = {
+ dagId: string;
+ dagRunId: string;
+ requestBody: BulkBody_BulkTaskInstanceBody_;
+};
+
+export type BulkTaskInstancesResponse = BulkResponse;
+
export type GetTaskInstancesBatchData = {
dagId: "~";
dagRunId: "~";
@@ -4584,6 +4650,27 @@ export type $OpenApiTs = {
422: HTTPValidationError;
};
};
+ patch: {
+ req: BulkTaskInstancesData;
+ res: {
+ /**
+ * Successful Response
+ */
+ 200: BulkResponse;
+ /**
+ * Unauthorized
+ */
+ 401: HTTPExceptionResponse;
+ /**
+ * Forbidden
+ */
+ 403: HTTPExceptionResponse;
+ /**
+ * Validation Error
+ */
+ 422: HTTPValidationError;
+ };
+ };
};
"/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/list": {
post: {
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index bfe8d049dcf..1e2a4b64601 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -4307,3 +4307,360 @@ class TestDeleteTaskInstance(TestTaskInstanceEndpoint):
remaining_tis = base_query.filter(TaskInstance.map_index !=
-1).all()
if expected_remaining is not None:
assert set(ti.map_index for ti in remaining_tis) ==
expected_remaining
+
+
+class TestBulkTaskInstances(TestTaskInstanceEndpoint):
+ DAG_ID = "example_python_operator"
+ TASK_ID = "print_the_context"
+ RUN_ID = "TEST_DAG_RUN_ID"
+ ENDPOINT_URL = f"/dags/{DAG_ID}/dagRuns/{RUN_ID}/taskInstances"
+
+ @pytest.mark.parametrize(
+ "default_ti, actions, expected_results",
+ [
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.SUCCESS}],
+ {
+ "actions": [
+ {
+ "action": "delete",
+ "entities": [
+ TASK_ID,
+ ],
+ }
+ ]
+ },
+ {
+ "delete": {
+ "success": [TASK_ID],
+ "errors": [],
+ }
+ },
+ id="delete-success",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.SUCCESS}],
+ {
+ "actions": [
+ {
+ "action": "delete",
+ "entities": [
+ "non_existent_task",
+ ],
+ "action_on_non_existence": "skip",
+ }
+ ]
+ },
+ {
+ "delete": {
+ "success": [],
+ "errors": [],
+ }
+ },
+ id="delete-skip",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.SUCCESS}],
+ {
+ "actions": [
+ {
+ "action": "delete",
+ "entities": [
+ "non_existent_task",
+ ],
+ }
+ ]
+ },
+ {
+ "delete": {
+ "success": [],
+ "errors": [
+ {
+ "error": "The task instances with these
task_ids: ['non_existent_task'] were not found",
+ "status_code": 404,
+ }
+ ],
+ }
+ },
+ id="delete-failure",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": TASK_ID,
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ },
+ ],
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [TASK_ID],
+ "errors": [],
+ }
+ },
+ id="update-success",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": "non_existent_task",
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ },
+ ],
+ "action_on_non_existence": "skip",
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [],
+ "errors": [],
+ }
+ },
+ id="update-skip",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING, "map_index":
100}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": TASK_ID,
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ "map_index": 100,
+ },
+ ],
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [TASK_ID],
+ "errors": [],
+ }
+ },
+ id="update-success-mapped-task",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": "non_existent_task",
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ },
+ ],
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [],
+ "errors": [
+ {
+ "error": "The Task Instance with dag_id:
`example_python_operator`, run_id: `TEST_DAG_RUN_ID`, task_id:
`non_existent_task` and map_index: `None` was not found",
+ "status_code": 404,
+ }
+ ],
+ }
+ },
+ id="update-failure",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": TASK_ID,
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ "map_index": -100,
+ },
+ ],
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [],
+ "errors": [
+ {
+ "error": "The Task Instance with dag_id:
`example_python_operator`, run_id: `TEST_DAG_RUN_ID`, task_id:
`print_the_context` and map_index: `-100` was not found",
+ "status_code": 404,
+ }
+ ],
+ }
+ },
+ id="update-failure-mapped-task",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.RUNNING}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": TASK_ID,
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ "map_index": -100,
+ },
+ ],
+ "action_on_non_existence": "skip",
+ }
+ ]
+ },
+ {
+ "update": {
+ "success": [],
+ "errors": [],
+ }
+ },
+ id="update-failure-mapped-task-with-skip",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.SUCCESS}],
+ {
+ "actions": [
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": "non_existent_task",
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ },
+ ],
+ },
+ {
+ "action": "update",
+ "entities": [
+ {
+ "task_id": TASK_ID,
+ "new_state": "failed",
+ "note": "test",
+ "include_upstream": True,
+ "include_downstream": True,
+ "include_future": True,
+ "include_past": True,
+ },
+ ],
+ },
+ {"action": "delete", "entities": [TASK_ID]},
+ {"action": "delete", "entities":
["non_existent_task"]},
+ ],
+ },
+ {
+ "update": {
+ "success": [TASK_ID],
+ "errors": [
+ {
+ "error": "The Task Instance with dag_id:
`example_python_operator`, run_id: `TEST_DAG_RUN_ID`, task_id:
`non_existent_task` and map_index: `None` was not found",
+ "status_code": 404,
+ }
+ ],
+ },
+ "delete": {
+ "success": [TASK_ID],
+ "errors": [
+ {
+ "error": "The task instances with these
task_ids: ['non_existent_task'] were not found",
+ "status_code": 404,
+ }
+ ],
+ },
+ },
+ id="update-delete-success",
+ ),
+ pytest.param(
+ [{"task_id": TASK_ID, "state": State.SUCCESS}],
+ {
+ "actions": [{"action": "create", "entities": []}],
+ },
+ {
+ "create": {
+ "success": [],
+ "errors": [
+ {"error": "Task instances bulk create is not
supported", "status_code": 405}
+ ],
+ }
+ },
+ id="create-failure",
+ ),
+ ],
+ )
+ def test_bulk_task_instances(self, test_client, session, default_ti,
actions, expected_results):
+ self.create_task_instances(session, task_instances=default_ti)
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json=actions,
+ )
+ assert response.status_code == 200
+ response_data = response.json()
+ for task_id, value in expected_results.items():
+ assert response_data[task_id] == value
+
+ def test_should_respond_401(self, unauthenticated_test_client):
+ response = unauthenticated_test_client.patch(self.ENDPOINT_URL,
json={})
+ assert response.status_code == 401
+
+ def test_should_respond_403(self, unauthorized_test_client):
+ response = unauthorized_test_client.patch(self.ENDPOINT_URL, json={})
+ assert response.status_code == 403
+
+ def test_should_respond_422(self, test_client):
+ response = test_client.patch(self.ENDPOINT_URL, json={})
+ assert response.status_code == 422
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
index 4dbcd4463f9..63056402f6a 100644
--- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
+++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -98,6 +98,19 @@ class BulkActionResponse(BaseModel):
] = []
+class BulkDeleteActionBulkTaskInstanceBody(BaseModel):
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ action: Annotated[
+ Literal["delete"], Field(description="The action to be performed on
the entities.", title="Action")
+ ]
+ entities: Annotated[
+ list[str], Field(description="A list of entity id/key to be deleted.",
title="Entities")
+ ]
+ action_on_non_existence: BulkActionNotOnExistence | None = "fail"
+
+
class BulkDeleteActionConnectionBody(BaseModel):
model_config = ConfigDict(
extra="forbid",
@@ -160,6 +173,10 @@ class BulkResponse(BaseModel):
] = None
+class Note(RootModel[str]):
+ root: Annotated[str, Field(max_length=1000, title="Note")]
+
+
class TaskIds(RootModel[list]):
root: Annotated[list, Field(max_length=2, min_length=2)]
@@ -293,10 +310,6 @@ class DAGRunClearBody(BaseModel):
only_failed: Annotated[bool | None, Field(title="Only Failed")] = False
-class Note(RootModel[str]):
- root: Annotated[str, Field(max_length=1000, title="Note")]
-
-
class DAGRunPatchStates(str, Enum):
"""
Enum for DAG Run states when updating a DAG Run.
@@ -1096,6 +1109,37 @@ class BulkCreateActionVariableBody(BaseModel):
action_on_existence: BulkActionOnExistence | None = "fail"
+class BulkTaskInstanceBody(BaseModel):
+ """
+ Request body for bulk update, and delete task instances.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ new_state: TaskInstanceState | None = None
+ note: Annotated[Note | None, Field(title="Note")] = None
+ include_upstream: Annotated[bool | None, Field(title="Include Upstream")]
= False
+ include_downstream: Annotated[bool | None, Field(title="Include
Downstream")] = False
+ include_future: Annotated[bool | None, Field(title="Include Future")] =
False
+ include_past: Annotated[bool | None, Field(title="Include Past")] = False
+ task_id: Annotated[str, Field(title="Task Id")]
+ map_index: Annotated[int | None, Field(title="Map Index")] = None
+
+
+class BulkUpdateActionBulkTaskInstanceBody(BaseModel):
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ action: Annotated[
+ Literal["update"], Field(description="The action to be performed on
the entities.", title="Action")
+ ]
+ entities: Annotated[
+ list[BulkTaskInstanceBody], Field(description="A list of entities to
be updated.", title="Entities")
+ ]
+ action_on_non_existence: BulkActionNotOnExistence | None = "fail"
+
+
class BulkUpdateActionConnectionBody(BaseModel):
model_config = ConfigDict(
extra="forbid",
@@ -1641,6 +1685,19 @@ class BulkBodyVariableBody(BaseModel):
]
+class BulkCreateActionBulkTaskInstanceBody(BaseModel):
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ action: Annotated[
+ Literal["create"], Field(description="The action to be performed on
the entities.", title="Action")
+ ]
+ entities: Annotated[
+ list[BulkTaskInstanceBody], Field(description="A list of entities to
be created.", title="Entities")
+ ]
+ action_on_existence: BulkActionOnExistence | None = "fail"
+
+
class DAGCollectionResponse(BaseModel):
"""
DAG Collection serializer for responses.
@@ -1702,3 +1759,17 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: Annotated[list[TaskInstanceHistoryResponse],
Field(title="Task Instances")]
total_entries: Annotated[int, Field(title="Total Entries")]
+
+
+class BulkBodyBulkTaskInstanceBody(BaseModel):
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ actions: Annotated[
+ list[
+ BulkCreateActionBulkTaskInstanceBody
+ | BulkUpdateActionBulkTaskInstanceBody
+ | BulkDeleteActionBulkTaskInstanceBody
+ ],
+ Field(title="Actions"),
+ ]
