AIP-84 Migrate patch dag to FastAPI API

airflow/api_connexion/endpoints/dag_endpoint.py

@@ -141,6 +141,7 @@ def get_dags(
         raise BadRequest("DAGCollectionSchema error", detail=str(e))
 
 
+@mark_fastapi_migration_done
 @security.requires_access_dag("PUT")
 @action_logging
 @provide_session

airflow/api_fastapi/openapi/v1-generated.yaml

@@ -123,13 +123,56 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /public/dags/{dag_id}:
+    patch:
+      tags:
+      - DAG
+      summary: Patch Dag
+      description: Update the specific DAG.
+      operationId: patch_dag_public_dags__dag_id__patch
+      parameters:
+      - name: dag_id
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Dag Id
+      - name: update_mask
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: array
+            items:
+              type: string
+          - type: 'null'
+          title: Update Mask
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/DAGPatchBody'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DAGResponse'
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     DAGCollectionResponse:
       properties:
         dags:
           items:
-            $ref: '#/components/schemas/DAGModelResponse'
+            $ref: '#/components/schemas/DAGResponse'
           type: array
           title: Dags
         total_entries:
@@ -141,7 +184,17 @@ components:
       - total_entries
       title: DAGCollectionResponse
       description: DAG Collection serializer for responses.
-    DAGModelResponse:
+    DAGPatchBody:
+      properties:
+        is_paused:
+          type: boolean
+          title: Is Paused
+      type: object
+      required:
+      - is_paused
+      title: DAGPatchBody
+      description: Dag Serializer for updatable body.
+    DAGResponse:
       properties:
         dag_id:
           type: string
@@ -292,7 +345,7 @@ components:
       - next_dagrun_create_after
       - owners
       - file_token
-      title: DAGModelResponse
+      title: DAGResponse
       description: DAG serializer for responses.
     DagTagPydantic:
       properties:

airflow/api_fastapi/serializers/dags.py

@@ -31,7 +31,7 @@
 from airflow.serialization.pydantic.dag import DagTagPydantic
 
 
-class DAGModelResponse(BaseModel):
+class DAGResponse(BaseModel):
     """DAG serializer for responses."""
 
     dag_id: str
@@ -82,8 +82,14 @@ def file_token(self) -> str:
         return serializer.dumps(self.fileloc)
 
 
+class DAGPatchBody(BaseModel):
+    """Dag Serializer for updatable body."""
+
+    is_paused: bool
+
+
 class DAGCollectionResponse(BaseModel):
     """DAG Collection serializer for responses."""
 
-    dags: list[DAGModelResponse]
+    dags: list[DAGResponse]
     total_entries: int

airflow/api_fastapi/views/public/dags.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, Query
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 from typing_extensions import Annotated
@@ -34,7 +34,7 @@
     QueryTagsFilter,
     SortParam,
 )
-from airflow.api_fastapi.serializers.dags import DAGCollectionResponse, DAGModelResponse
+from airflow.api_fastapi.serializers.dags import DAGCollectionResponse, DAGPatchBody, DAGResponse
 from airflow.models import DagModel
 from airflow.utils.db import get_query_count
 
@@ -43,7 +43,6 @@
 
 @dags_router.get("/dags")
 async def get_dags(
-    *,
     limit: QueryLimit,
     offset: QueryOffset,
     tags: QueryTagsFilter,
@@ -74,8 +73,35 @@ async def get_dags(
 
     try:
         return DAGCollectionResponse(
-            dags=[DAGModelResponse.model_validate(dag, from_attributes=True) for dag in dags],
+            dags=[DAGResponse.model_validate(dag, from_attributes=True) for dag in dags],
             total_entries=total_entries,
         )
     except ValueError as e:
         raise HTTPException(400, f"DAGCollectionSchema error: {str(e)}")
+
+
+@dags_router.patch("/dags/{dag_id}")
+async def patch_dag(
+    dag_id: str,
+    patch_body: DAGPatchBody,
+    session: Annotated[Session, Depends(get_session)],
+    update_mask: list[str] | None = Query(None),
+) -> DAGResponse:
+    """Update the specific DAG."""
+    dag = session.get(DagModel, dag_id)
+
+    if dag is None:
+        raise HTTPException(404, f"Dag with id: {dag_id} was not found")
+
+    if update_mask:
+        if update_mask != ["is_paused"]:
+            raise HTTPException(400, "Only `is_paused` field can be updated through the REST API")
+
+    else:
+        update_mask = ["is_paused"]
+
+    for attr_name in update_mask:
+        attr_value = getattr(patch_body, attr_name)
+        setattr(dag, attr_name, attr_value)
+
+    return DAGResponse.model_validate(dag, from_attributes=True)

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

@@ -72,3 +72,6 @@ export const UseDagServiceGetDagsPublicDagsGetKeyFn = (
     },
   ]),
 ];
+export type DagServicePatchDagPublicDagsDagIdPatchMutationResult = Awaited<
+  ReturnType<typeof DagService.patchDagPublicDagsDagIdPatch>
+>;

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

@@ -1,7 +1,13 @@
 // generated with @7nohe/openapi-react-query-codegen@1.6.0
-import { useQuery, UseQueryOptions } from "@tanstack/react-query";
+import {
+  useMutation,
+  UseMutationOptions,
+  useQuery,
+  UseQueryOptions,
+} from "@tanstack/react-query";
 
 import { DagService, DatasetService } from "../requests/services.gen";
+import { DAGPatchBody } from "../requests/types.gen";
 import * as Common from "./common";
 
 /**
@@ -110,3 +116,50 @@ export const useDagServiceGetDagsPublicDagsGet = <
       }) as TData,
     ...options,
   });
+/**
+ * Patch Dag
+ * Update the specific DAG.
+ * @param data The data for the request.
+ * @param data.dagId
+ * @param data.requestBody
+ * @param data.updateMask
+ * @returns DAGResponse Successful Response
+ * @throws ApiError
+ */
+export const useDagServicePatchDagPublicDagsDagIdPatch = <
+  TData = Common.DagServicePatchDagPublicDagsDagIdPatchMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        dagId: string;
+        requestBody: DAGPatchBody;
+        updateMask?: string[];
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      dagId: string;
+      requestBody: DAGPatchBody;
+      updateMask?: string[];
+    },
+    TContext
+  >({
+    mutationFn: ({ dagId, requestBody, updateMask }) =>
+      DagService.patchDagPublicDagsDagIdPatch({
+        dagId,
+        requestBody,
+        updateMask,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });

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

@@ -4,7 +4,7 @@ export const $DAGCollectionResponse = {
   properties: {
     dags: {
       items: {
-        $ref: "#/components/schemas/DAGModelResponse",
+        $ref: "#/components/schemas/DAGResponse",
       },
       type: "array",
       title: "Dags",
@@ -20,7 +20,20 @@ export const $DAGCollectionResponse = {
   description: "DAG Collection serializer for responses.",
 } as const;
 
-export const $DAGModelResponse = {
+export const $DAGPatchBody = {
+  properties: {
+    is_paused: {
+      type: "boolean",
+      title: "Is Paused",
+    },
+  },
+  type: "object",
+  required: ["is_paused"],
+  title: "DAGPatchBody",
+  description: "Dag Serializer for updatable body.",
+} as const;
+
+export const $DAGResponse = {
   properties: {
     dag_id: {
       type: "string",
@@ -271,7 +284,7 @@ export const $DAGModelResponse = {
     "owners",
     "file_token",
   ],
-  title: "DAGModelResponse",
+  title: "DAGResponse",
   description: "DAG serializer for responses.",
 } as const;
 

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

@@ -7,6 +7,8 @@ import type {
   NextRunDatasetsUiNextRunDatasetsDagIdGetResponse,
   GetDagsPublicDagsGetData,
   GetDagsPublicDagsGetResponse,
+  PatchDagPublicDagsDagIdPatchData,
+  PatchDagPublicDagsDagIdPatchResponse,
 } from "./types.gen";
 
 export class DatasetService {
@@ -72,4 +74,34 @@ export class DagService {
       },
     });
   }
+
+  /**
+   * Patch Dag
+   * Update the specific DAG.
+   * @param data The data for the request.
+   * @param data.dagId
+   * @param data.requestBody
+   * @param data.updateMask
+   * @returns DAGResponse Successful Response
+   * @throws ApiError
+   */
+  public static patchDagPublicDagsDagIdPatch(
+    data: PatchDagPublicDagsDagIdPatchData,
+  ): CancelablePromise<PatchDagPublicDagsDagIdPatchResponse> {
+    return __request(OpenAPI, {
+      method: "PATCH",
+      url: "/public/dags/{dag_id}",
+      path: {
+        dag_id: data.dagId,
+      },
+      query: {
+        update_mask: data.updateMask,
+      },
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        422: "Validation Error",
+      },
+    });
+  }
 }

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

@@ -4,14 +4,21 @@
  * DAG Collection serializer for responses.
  */
 export type DAGCollectionResponse = {
-  dags: Array<DAGModelResponse>;
+  dags: Array<DAGResponse>;
   total_entries: number;
 };
 
+/**
+ * Dag Serializer for updatable body.
+ */
+export type DAGPatchBody = {
+  is_paused: boolean;
+};
+
 /**
  * DAG serializer for responses.
  */
-export type DAGModelResponse = {
+export type DAGResponse = {
   dag_id: string;
   dag_display_name: string;
   is_paused: boolean;
@@ -83,6 +90,14 @@ export type GetDagsPublicDagsGetData = {
 
 export type GetDagsPublicDagsGetResponse = DAGCollectionResponse;
 
+export type PatchDagPublicDagsDagIdPatchData = {
+  dagId: string;
+  requestBody: DAGPatchBody;
+  updateMask?: Array<string> | null;
+};
+
+export type PatchDagPublicDagsDagIdPatchResponse = DAGResponse;
+
 export type $OpenApiTs = {
   "/ui/next_run_datasets/{dag_id}": {
     get: {
@@ -116,4 +131,19 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/dags/{dag_id}": {
+    patch: {
+      req: PatchDagPublicDagsDagIdPatchData;
+      res: {
+        /**
+         * Successful Response
+         */
+        200: DAGResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
+  };
 };