AIP-84 Migrate Modify Dag Run endpoint to FastAPI

airflow/api_connexion/endpoints/dag_run_endpoint.py

@@ -373,6 +373,7 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
     raise AlreadyExists(detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists")
 
 
+@mark_fastapi_migration_done
 @security.requires_access_dag("PUT", DagAccessEntity.RUN)
 @provide_session
 @action_logging

airflow/api_fastapi/openapi/v1-generated.yaml

@@ -944,6 +944,68 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+    patch:
+      tags:
+      - DagRun
+      summary: Modify Dag Run
+      description: Modify a DAG Run.
+      operationId: modify_dag_run
+      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/DAGRunPatchBody'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DAGRunResponse'
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '401':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Unauthorized
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/monitor/health:
     get:
       tags:
@@ -1453,6 +1515,23 @@ components:
       - file_token
       title: DAGResponse
       description: DAG serializer for responses.
+    DAGRunModifyFormStates:
+      type: string
+      enum:
+      - queued
+      - success
+      - failed
+      title: DAGRunModifyFormStates
+      description: Enum for DAG Run states when updating a DAG Run.
+    DAGRunPatchBody:
+      properties:
+        state:
+          $ref: '#/components/schemas/DAGRunModifyFormStates'
+      type: object
+      required:
+      - state
+      title: DAGRunPatchBody
+      description: DAG Run Serializer for PATCH requests.
     DAGRunResponse:
       properties:
         run_id:

airflow/api_fastapi/routes/public/dag_run.py

@@ -25,7 +25,7 @@
 from airflow.api_fastapi.db.common import get_session
 from airflow.api_fastapi.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.routes.router import AirflowRouter
-from airflow.api_fastapi.serializers.dag_run import DAGRunResponse
+from airflow.api_fastapi.serializers.dag_run import DAGRunPatchBody, DAGRunResponse
 from airflow.models import DagRun
 
 dag_run_router = AirflowRouter(tags=["DagRun"], prefix="/dags/{dag_id}/dagRuns")
@@ -57,3 +57,19 @@ async def delete_dag_run(dag_id: str, dag_run_id: str, session: Annotated[Sessio
         )
 
     session.delete(dag_run)
+
+
+@dag_run_router.patch("/{dag_run_id}", responses=create_openapi_http_exception_doc([400, 401, 403, 404]))
+async def modify_dag_run(
+    dag_id: str, dag_run_id: str, state: DAGRunPatchBody, session: Annotated[Session, Depends(get_session)]
+) -> DAGRunResponse:
+    """Modify a DAG Run."""
+    dag_run = session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=dag_run_id))
+    if dag_run is None:
+        raise HTTPException(
+            404, f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found"
+        )
+
+    setattr(dag_run, "state", state.state)
+
+    return DAGRunResponse.model_validate(dag_run, from_attributes=True)

airflow/api_fastapi/serializers/dag_run.py

@@ -18,13 +18,28 @@
 from __future__ import annotations
 
 from datetime import datetime
+from enum import Enum
 
 from pydantic import BaseModel, Field
 
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 
+class DAGRunModifyFormStates(str, Enum):
+    """Enum for DAG Run states when updating a DAG Run."""
+
+    QUEUED = DagRunState.QUEUED
+    SUCCESS = DagRunState.SUCCESS
+    FAILED = DagRunState.FAILED
+
+
+class DAGRunPatchBody(BaseModel):
+    """DAG Run Serializer for PATCH requests."""
+
+    state: DAGRunModifyFormStates
+
+
 class DAGRunResponse(BaseModel):
     """DAG Run serializer for responses."""
 

airflow/ui/openapi-gen/queries/common.ts

@@ -234,6 +234,9 @@ export type DagServicePatchDagMutationResult = Awaited<
 export type VariableServicePatchVariableMutationResult = Awaited<
   ReturnType<typeof VariableService.patchVariable>
 >;
+export type DagRunServiceModifyDagRunMutationResult = Awaited<
+  ReturnType<typeof DagRunService.modifyDagRun>
+>;
 export type DagServiceDeleteDagMutationResult = Awaited<
   ReturnType<typeof DagService.deleteDag>
 >;

airflow/ui/openapi-gen/queries/queries.ts

@@ -15,7 +15,12 @@ import {
   MonitorService,
   VariableService,
 } from "../requests/services.gen";
-import { DAGPatchBody, DagRunState, VariableBody } from "../requests/types.gen";
+import {
+  DAGPatchBody,
+  DAGRunPatchBody,
+  DagRunState,
+  VariableBody,
+} from "../requests/types.gen";
 import * as Common from "./common";
 
 /**
@@ -569,6 +574,53 @@ export const useVariableServicePatchVariable = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Modify Dag Run
+ * Modify a DAG Run.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.dagRunId
+ * @param data.requestBody
+ * @returns DAGRunResponse Successful Response
+ * @throws ApiError
+ */
+export const useDagRunServiceModifyDagRun = <
+  TData = Common.DagRunServiceModifyDagRunMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        dagId: string;
+        dagRunId: string;
+        requestBody: DAGRunPatchBody;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      dagId: string;
+      dagRunId: string;
+      requestBody: DAGRunPatchBody;
+    },
+    TContext
+  >({
+    mutationFn: ({ dagId, dagRunId, requestBody }) =>
+      DagRunService.modifyDagRun({
+        dagId,
+        dagRunId,
+        requestBody,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Delete Dag
  * Delete the specific DAG.

airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -824,6 +824,25 @@ export const $DAGResponse = {
   description: "DAG serializer for responses.",
 } as const;
 
+export const $DAGRunModifyFormStates = {
+  type: "string",
+  enum: ["queued", "success", "failed"],
+  title: "DAGRunModifyFormStates",
+  description: "Enum for DAG Run states when updating a DAG Run.",
+} as const;
+
+export const $DAGRunPatchBody = {
+  properties: {
+    state: {
+      $ref: "#/components/schemas/DAGRunModifyFormStates",
+    },
+  },
+  type: "object",
+  required: ["state"],
+  title: "DAGRunPatchBody",
+  description: "DAG Run Serializer for PATCH requests.",
+} as const;
+
 export const $DAGRunResponse = {
   properties: {
     run_id: {

airflow/ui/openapi-gen/requests/services.gen.ts

@@ -37,6 +37,8 @@ import type {
   GetDagRunResponse,
   DeleteDagRunData,
   DeleteDagRunResponse,
+  ModifyDagRunData,
+  ModifyDagRunResponse,
   GetHealthResponse,
 } from "./types.gen";
 
@@ -546,6 +548,38 @@ export class DagRunService {
       },
     });
   }
+
+  /**
+   * Modify Dag Run
+   * Modify a DAG Run.
+   * @param data The data for the request.
+   * @param data.dagId
+   * @param data.dagRunId
+   * @param data.requestBody
+   * @returns DAGRunResponse Successful Response
+   * @throws ApiError
+   */
+  public static modifyDagRun(
+    data: ModifyDagRunData,
+  ): CancelablePromise<ModifyDagRunResponse> {
+    return __request(OpenAPI, {
+      method: "PATCH",
+      url: "/public/dags/{dag_id}/dagRuns/{dag_run_id}",
+      path: {
+        dag_id: data.dagId,
+        dag_run_id: data.dagRunId,
+      },
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        400: "Bad Request",
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
 }
 
 export class MonitorService {

airflow/ui/openapi-gen/requests/types.gen.ts

@@ -135,6 +135,18 @@ export type DAGResponse = {
   readonly file_token: string;
 };
 
+/**
+ * Enum for DAG Run states when updating a DAG Run.
+ */
+export type DAGRunModifyFormStates = "queued" | "success" | "failed";
+
+/**
+ * DAG Run Serializer for PATCH requests.
+ */
+export type DAGRunPatchBody = {
+  state: DAGRunModifyFormStates;
+};
+
 /**
  * DAG Run serializer for responses.
  */
@@ -448,6 +460,14 @@ export type DeleteDagRunData = {
 
 export type DeleteDagRunResponse = void;
 
+export type ModifyDagRunData = {
+  dagId: string;
+  dagRunId: string;
+  requestBody: DAGRunPatchBody;
+};
+
+export type ModifyDagRunResponse = DAGRunResponse;
+
 export type GetHealthResponse = HealthInfoSchema;
 
 export type $OpenApiTs = {
@@ -889,6 +909,35 @@ export type $OpenApiTs = {
         422: HTTPValidationError;
       };
     };
+    patch: {
+      req: ModifyDagRunData;
+      res: {
+        /**
+         * Successful Response
+         */
+        200: DAGRunResponse;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
   };
   "/public/monitor/health": {
     get: {