Skip to content

Commit fe92cfd

Browse files
pierrejeambrunpierrejeambrun
authored
API: Add bulk update to mark Dag runs as success/failed (#67948)
* API: Add bulk update to mark Dag runs as success/failed * Address review: factor out dag-run patch helpers and align bulk update with task-instance service * Adjusments on action_on_existence behavior * Make cross-module patch_dag_run helpers public (drop leading underscore) patch_dag_run_state / patch_dag_run_note are imported by the route module, so the leading underscore was misleading. Also rename the bulk-delete `keys` set to `to_delete_keys` to mirror the bulk-update handler.
1 parent 837bfde commit fe92cfd

10 files changed

Lines changed: 454 additions & 215 deletions

File tree

‎airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,10 +51,12 @@ class DAGRunPatchBody(StrictBaseModel):
5151

5252

5353
class BulkDAGRunBody(StrictBaseModel):
54-
"""Request body for bulk delete operations on Dag Runs."""
54+
"""Request body for bulk operations on Dag Runs."""
5555

5656
dag_run_id: str
5757
dag_id: str | None = None
58+
state: DagRunMutableStates | None = None
59+
note: str | None = Field(None, max_length=1000)
5860

5961

6062
class BaseDAGRunClear(StrictBaseModel):

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2153,7 +2153,7 @@ paths:
21532153
tags:
21542154
- DagRun
21552155
summary: Bulk Dag Runs
2156-
description: Bulk delete Dag Runs.
2156+
description: Bulk update or delete Dag Runs.
21572157
operationId: bulk_dag_runs
21582158
security:
21592159
- OAuth2PasswordBearer: []
@@ -12122,12 +12122,22 @@ components:
1212212122
- type: string
1212312123
- type: 'null'
1212412124
title: Dag Id
12125+
state:
12126+
anyOf:
12127+
- $ref: '#/components/schemas/DagRunMutableStates'
12128+
- type: 'null'
12129+
note:
12130+
anyOf:
12131+
- type: string
12132+
maxLength: 1000
12133+
- type: 'null'
12134+
title: Note
1212512135
additionalProperties: false
1212612136
type: object
1212712137
required:
1212812138
- dag_run_id
1212912139
title: BulkDAGRunBody
12130-
description: Request body for bulk delete operations on Dag Runs.
12140+
description: Request body for bulk operations on Dag Runs.
1213112141
BulkDAGRunClearBody:
1213212142
properties:
1213312143
dry_run:

‎airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py‎

Lines changed: 12 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,13 @@
2020
import textwrap
2121
from typing import Annotated, Literal, cast
2222

23-
import structlog
2423
from fastapi import Depends, HTTPException, Query, Request, status
2524
from fastapi.exceptions import RequestValidationError
2625
from fastapi.responses import StreamingResponse
2726
from pydantic import ValidationError
2827
from sqlalchemy import select
2928
from sqlalchemy.orm import joinedload
3029

31-
from airflow.api.common.mark_tasks import (
32-
set_dag_run_state_to_failed,
33-
set_dag_run_state_to_queued,
34-
set_dag_run_state_to_success,
35-
)
3630
from airflow.api_fastapi.app import get_auth_manager
3731
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
3832
from airflow.api_fastapi.common.cursors import (
@@ -106,19 +100,18 @@
106100
DagRunWaiter,
107101
dry_run_clear_dag_run,
108102
get_dag_run_and_dag_for_clear,
103+
patch_dag_run_note,
104+
patch_dag_run_state,
109105
perform_clear_dag_run,
110106
)
111107
from airflow.api_fastapi.logging.decorators import action_logging
112108
from airflow.exceptions import ParamValidationError
113-
from airflow.listeners.listener import get_listener_manager
114109
from airflow.models import DagModel, DagRun
115110
from airflow.models.asset import AssetEvent
116111
from airflow.models.dag_version import DagVersion
117112
from airflow.utils.state import DagRunState
118113
from airflow.utils.types import DagRunTriggeredByType, DagRunType
119114

120-
log = structlog.get_logger(__name__)
121-
122115
dag_run_router = AirflowRouter(tags=["DagRun"], prefix="/dags/{dag_id}/dagRuns")
123116
dag_run_at_dag_router = AirflowRouter(tags=["DagRun"], prefix="/dags/{dag_id}")
124117

@@ -227,33 +220,12 @@ def patch_dag_run(
227220
data = patch_body.model_dump(include=fields_to_update, by_alias=True)
228221

229222
for attr_name, attr_value_raw in data.items():
230-
if attr_name == "state":
231-
attr_value = getattr(patch_body, "state")
232-
if attr_value == DagRunMutableStates.SUCCESS:
233-
set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
234-
try:
235-
get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="")
236-
except Exception:
237-
log.exception("error calling listener")
238-
239-
# TODO AIP-103: https://github.com/apache/airflow/issues/66755
240-
# Handle clearing states for all task instances in a dagrun when cleared
241-
elif attr_value == DagRunMutableStates.QUEUED:
242-
set_dag_run_state_to_queued(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
243-
# Not notifying on queued - only notifying on RUNNING, this is happening in scheduler
244-
elif attr_value == DagRunMutableStates.FAILED:
245-
set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
246-
try:
247-
get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run, msg="")
248-
except Exception:
249-
log.exception("error calling listener")
223+
if attr_name == "state" and patch_body.state is not None:
224+
patch_dag_run_state(dag=dag, dag_run=dag_run, state=patch_body.state, session=session)
250225
elif attr_name == "note":
251226
updated_dag_run = session.get(DagRun, dag_run.id)
252-
if updated_dag_run and updated_dag_run.dag_run_note is None:
253-
updated_dag_run.note = (attr_value_raw, user.get_id())
254-
elif updated_dag_run:
255-
updated_dag_run.dag_run_note.content = attr_value_raw
256-
updated_dag_run.dag_run_note.user_id = user.get_id()
227+
if updated_dag_run is not None:
228+
patch_dag_run_note(dag_run=updated_dag_run, note=attr_value_raw, user=user)
257229

258230
final_dag_run = session.get(DagRun, dag_run.id)
259231
if not final_dag_run:
@@ -270,9 +242,13 @@ def bulk_dag_runs(
270242
request: BulkBody[BulkDAGRunBody],
271243
session: SessionDep,
272244
dag_id: str,
245+
dag_bag: DagBagDep,
246+
user: GetUserDep,
273247
) -> BulkResponse:
274-
"""Bulk delete Dag Runs."""
275-
return BulkDagRunService(session=session, request=request, dag_id=dag_id).handle_request()
248+
"""Bulk update or delete Dag Runs."""
249+
return BulkDagRunService(
250+
session=session, request=request, dag_id=dag_id, dag_bag=dag_bag, user=user
251+
).handle_request()
276252

277253

278254
@dag_run_router.get(

0 commit comments

Comments
 (0)