AIP-84 Post Pool (apache#43317)

astronomer · Oct 24, 2024 · 6070bb6 · 6070bb6
1 parent 27f3be6
commit 6070bb6
Show file tree

Hide file tree

Showing 10 changed files with 314 additions and 45 deletions.
diff --git a/airflow/api_connexion/endpoints/pool_endpoint.py b/airflow/api_connexion/endpoints/pool_endpoint.py
@@ -147,6 +147,7 @@ def patch_pool(
     return pool_schema.dump(pool)
 
 
+@mark_fastapi_migration_done
 @security.requires_access_pool("POST")
 @action_logging
 @provide_session

diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -1191,7 +1191,7 @@ paths:
         content:
           application/json:
             schema:
-              $ref: '#/components/schemas/PoolBody'
+              $ref: '#/components/schemas/PoolPatchBody'
       responses:
         '200':
           description: Successful Response
@@ -1289,6 +1289,43 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+    post:
+      tags:
+      - Pool
+      summary: Post Pool
+      description: Create a Pool.
+      operationId: post_pool
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/PoolPostBody'
+      responses:
+        '201':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/PoolResponse'
+        '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'
   /public/providers/:
     get:
       tags:
@@ -2290,7 +2327,23 @@ components:
       - timetables
       title: PluginResponse
       description: Plugin serializer.
-    PoolBody:
+    PoolCollectionResponse:
+      properties:
+        pools:
+          items:
+            $ref: '#/components/schemas/PoolResponse'
+          type: array
+          title: Pools
+        total_entries:
+          type: integer
+          title: Total Entries
+      type: object
+      required:
+      - pools
+      - total_entries
+      title: PoolCollectionResponse
+      description: Pool Collection serializer for responses.
+    PoolPatchBody:
       properties:
         pool:
           anyOf:
@@ -2313,24 +2366,31 @@ components:
           - type: 'null'
           title: Include Deferred
       type: object
-      title: PoolBody
-      description: Pool serializer for bodies.
-    PoolCollectionResponse:
+      title: PoolPatchBody
+      description: Pool serializer for patch bodies.
+    PoolPostBody:
       properties:
-        pools:
-          items:
-            $ref: '#/components/schemas/PoolResponse'
-          type: array
-          title: Pools
-        total_entries:
+        name:
+          type: string
+          title: Name
+        slots:
           type: integer
-          title: Total Entries
+          title: Slots
+        description:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Description
+        include_deferred:
+          type: boolean
+          title: Include Deferred
+          default: false
       type: object
       required:
-      - pools
-      - total_entries
-      title: PoolCollectionResponse
-      description: Pool Collection serializer for responses.
+      - name
+      - slots
+      title: PoolPostBody
+      description: Pool serializer for post bodies.
     PoolResponse:
       properties:
         name:

diff --git a/airflow/api_fastapi/core_api/routes/public/pools.py b/airflow/api_fastapi/core_api/routes/public/pools.py
@@ -29,8 +29,9 @@
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.serializers.pools import (
     BasePool,
-    PoolBody,
     PoolCollectionResponse,
+    PoolPatchBody,
+    PoolPostBody,
     PoolResponse,
 )
 from airflow.models.pool import Pool
@@ -107,7 +108,7 @@ async def get_pools(
 @pools_router.patch("/{pool_name}", responses=create_openapi_http_exception_doc([400, 401, 403, 404]))
 async def patch_pool(
     pool_name: str,
-    patch_body: PoolBody,
+    patch_body: PoolPatchBody,
     session: Annotated[Session, Depends(get_session)],
     update_mask: list[str] | None = Query(None),
 ) -> PoolResponse:
@@ -136,3 +137,16 @@ async def patch_pool(
         setattr(pool, key, value)
 
     return PoolResponse.model_validate(pool, from_attributes=True)
+
+
+@pools_router.post("/", status_code=201, responses=create_openapi_http_exception_doc([401, 403]))
+async def post_pool(
+    post_body: PoolPostBody,
+    session: Annotated[Session, Depends(get_session)],
+) -> PoolResponse:
+    """Create a Pool."""
+    pool = Pool(**post_body.model_dump())
+
+    session.add(pool)
+
+    return PoolResponse.model_validate(pool, from_attributes=True)
diff --git a/airflow/api_fastapi/core_api/serializers/pools.py b/airflow/api_fastapi/core_api/serializers/pools.py
@@ -58,12 +58,20 @@ class PoolCollectionResponse(BaseModel):
     total_entries: int
 
 
-class PoolBody(BaseModel):
-    """Pool serializer for bodies."""
+class PoolPatchBody(BaseModel):
+    """Pool serializer for patch bodies."""
 
     model_config = ConfigDict(populate_by_name=True)
 
     name: str | None = Field(default=None, alias="pool")
     slots: int | None = None
     description: str | None = None
     include_deferred: bool | None = None
+
+
+class PoolPostBody(BasePool):
+    """Pool serializer for post bodies."""
+
+    pool: str = Field(alias="name")
+    description: str | None = None
+    include_deferred: bool = False
diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -361,6 +361,9 @@ export const UseVersionServiceGetVersionKeyFn = (queryKey?: Array<unknown>) => [
 export type VariableServicePostVariableMutationResult = Awaited<
   ReturnType<typeof VariableService.postVariable>
 >;
+export type PoolServicePostPoolMutationResult = Awaited<
+  ReturnType<typeof PoolService.postPool>
+>;
 export type DagServicePatchDagsMutationResult = Awaited<
   ReturnType<typeof DagService.patchDags>
 >;

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -22,7 +22,8 @@ import {
 import {
   DAGPatchBody,
   DagRunState,
-  PoolBody,
+  PoolPatchBody,
+  PoolPostBody,
   VariableBody,
 } from "../requests/types.gen";
 import * as Common from "./common";
@@ -621,6 +622,43 @@ export const useVariableServicePostVariable = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Post Pool
+ * Create a Pool.
+ * @param data The data for the request.
+ * @param data.requestBody
+ * @returns PoolResponse Successful Response
+ * @throws ApiError
+ */
+export const usePoolServicePostPool = <
+  TData = Common.PoolServicePostPoolMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        requestBody: PoolPostBody;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      requestBody: PoolPostBody;
+    },
+    TContext
+  >({
+    mutationFn: ({ requestBody }) =>
+      PoolService.postPool({ requestBody }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Patch Dags
  * Patch multiple DAGs.
@@ -822,7 +860,7 @@ export const usePoolServicePatchPool = <
       TError,
       {
         poolName: string;
-        requestBody: PoolBody;
+        requestBody: PoolPatchBody;
         updateMask?: string[];
       },
       TContext
@@ -835,7 +873,7 @@ export const usePoolServicePatchPool = <
     TError,
     {
       poolName: string;
-      requestBody: PoolBody;
+      requestBody: PoolPatchBody;
       updateMask?: string[];
     },
     TContext

diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -1420,7 +1420,27 @@ export const $PluginResponse = {
   description: "Plugin serializer.",
 } as const;
 
-export const $PoolBody = {
+export const $PoolCollectionResponse = {
+  properties: {
+    pools: {
+      items: {
+        $ref: "#/components/schemas/PoolResponse",
+      },
+      type: "array",
+      title: "Pools",
+    },
+    total_entries: {
+      type: "integer",
+      title: "Total Entries",
+    },
+  },
+  type: "object",
+  required: ["pools", "total_entries"],
+  title: "PoolCollectionResponse",
+  description: "Pool Collection serializer for responses.",
+} as const;
+
+export const $PoolPatchBody = {
   properties: {
     pool: {
       anyOf: [
@@ -1468,28 +1488,41 @@ export const $PoolBody = {
     },
   },
   type: "object",
-  title: "PoolBody",
-  description: "Pool serializer for bodies.",
+  title: "PoolPatchBody",
+  description: "Pool serializer for patch bodies.",
 } as const;
 
-export const $PoolCollectionResponse = {
+export const $PoolPostBody = {
   properties: {
-    pools: {
-      items: {
-        $ref: "#/components/schemas/PoolResponse",
-      },
-      type: "array",
-      title: "Pools",
+    name: {
+      type: "string",
+      title: "Name",
     },
-    total_entries: {
+    slots: {
       type: "integer",
-      title: "Total Entries",
+      title: "Slots",
+    },
+    description: {
+      anyOf: [
+        {
+          type: "string",
+        },
+        {
+          type: "null",
+        },
+      ],
+      title: "Description",
+    },
+    include_deferred: {
+      type: "boolean",
+      title: "Include Deferred",
+      default: false,
     },
   },
   type: "object",
-  required: ["pools", "total_entries"],
-  title: "PoolCollectionResponse",
-  description: "Pool Collection serializer for responses.",
+  required: ["name", "slots"],
+  title: "PoolPostBody",
+  description: "Pool serializer for post bodies.",
 } as const;
 
 export const $PoolResponse = {

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -50,6 +50,8 @@ import type {
   PatchPoolResponse,
   GetPoolsData,
   GetPoolsResponse,
+  PostPoolData,
+  PostPoolResponse,
   GetProvidersData,
   GetProvidersResponse,
   GetPluginsData,
@@ -754,6 +756,30 @@ export class PoolService {
       },
     });
   }
+
+  /**
+   * Post Pool
+   * Create a Pool.
+   * @param data The data for the request.
+   * @param data.requestBody
+   * @returns PoolResponse Successful Response
+   * @throws ApiError
+   */
+  public static postPool(
+    data: PostPoolData,
+  ): CancelablePromise<PostPoolResponse> {
+    return __request(OpenAPI, {
+      method: "POST",
+      url: "/public/pools/",
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        401: "Unauthorized",
+        403: "Forbidden",
+        422: "Validation Error",
+      },
+    });
+  }
 }
 
 export class ProviderService {