AIP-84 Delete Pool (#43165)

Author: pierrejeambrun

airflow/api_connexion/endpoints/pool_endpoint.py

@@ -30,6 +30,7 @@
 from airflow.api_connexion.parameters import apply_sorting, check_limit, format_parameters
 from airflow.api_connexion.schemas.pool_schema import PoolCollection, pool_collection_schema, pool_schema
 from airflow.models.pool import Pool
+from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.www.decorators import action_logging
 
@@ -39,6 +40,7 @@
     from airflow.api_connexion.types import APIResponse, UpdateMask
 
 
+@mark_fastapi_migration_done
 @security.requires_access_pool("DELETE")
 @action_logging
 @provide_session
@@ -118,9 +120,11 @@ def patch_pool(
             # there is no way field is None here (UpdateMask is a List[str])
             # so if pool_schema.declared_fields[field].attribute is None file is returned
             update_mask = [
-                pool_schema.declared_fields[field].attribute  # type: ignore[misc]
-                if pool_schema.declared_fields[field].attribute
-                else field
+                (
+                    pool_schema.declared_fields[field].attribute  # type: ignore[misc]
+                    if pool_schema.declared_fields[field].attribute
+                    else field
+                )
                 for field in update_mask
             ]
         except KeyError as err:

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -1010,6 +1010,53 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HealthInfoSchema'
+  /public/pools/{pool_name}:
+    delete:
+      tags:
+      - Pool
+      summary: Delete Pool
+      description: Delete a pool entry.
+      operationId: delete_pool
+      parameters:
+      - name: pool_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Pool Name
+      responses:
+        '204':
+          description: Successful Response
+        '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'
 components:
   schemas:
     BaseInfoSchema:

airflow/api_fastapi/core_api/routes/public/__init__.py

@@ -22,6 +22,7 @@
 from airflow.api_fastapi.core_api.routes.public.dag_run import dag_run_router
 from airflow.api_fastapi.core_api.routes.public.dags import dags_router
 from airflow.api_fastapi.core_api.routes.public.monitor import monitor_router
+from airflow.api_fastapi.core_api.routes.public.pools import pools_router
 from airflow.api_fastapi.core_api.routes.public.variables import variables_router
 
 public_router = AirflowRouter(prefix="/public")
@@ -32,3 +33,4 @@
 public_router.include_router(variables_router)
 public_router.include_router(dag_run_router)
 public_router.include_router(monitor_router)
+public_router.include_router(pools_router)

airflow/api_fastapi/core_api/routes/public/pools.py

@@ -0,0 +1,48 @@
+# 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
+
+from fastapi import Depends, HTTPException
+from sqlalchemy import delete
+from sqlalchemy.orm import Session
+from typing_extensions import Annotated
+
+from airflow.api_fastapi.common.db.common import get_session
+from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.models.pool import Pool
+
+pools_router = AirflowRouter(tags=["Pool"], prefix="/pools")
+
+
+@pools_router.delete(
+    "/{pool_name}",
+    status_code=204,
+    responses=create_openapi_http_exception_doc([400, 401, 403, 404]),
+)
+async def delete_pool(
+    pool_name: str,
+    session: Annotated[Session, Depends(get_session)],
+):
+    """Delete a pool entry."""
+    if pool_name == "default_pool":
+        raise HTTPException(400, "Default Pool can't be deleted")
+
+    affected_count = session.execute(delete(Pool).where(Pool.pool == pool_name)).rowcount
+
+    if affected_count == 0:
+        raise HTTPException(404, f"The Pool with name: `{pool_name}` was not found")

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

@@ -8,6 +8,7 @@ import {
   DagService,
   DashboardService,
   MonitorService,
+  PoolService,
   VariableService,
 } from "../requests/services.gen";
 import { DagRunState } from "../requests/types.gen";
@@ -269,3 +270,6 @@ export type VariableServiceDeleteVariableMutationResult = Awaited<
 export type DagRunServiceDeleteDagRunMutationResult = Awaited<
   ReturnType<typeof DagRunService.deleteDagRun>
 >;
+export type PoolServiceDeletePoolMutationResult = Awaited<
+  ReturnType<typeof PoolService.deletePool>
+>;

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

@@ -13,6 +13,7 @@ import {
   DagService,
   DashboardService,
   MonitorService,
+  PoolService,
   VariableService,
 } from "../requests/services.gen";
 import { DAGPatchBody, DagRunState, VariableBody } from "../requests/types.gen";
@@ -763,3 +764,40 @@ export const useDagRunServiceDeleteDagRun = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Delete Pool
+ * Delete a pool entry.
+ * @param data The data for the request.
+ * @param data.poolName
+ * @returns void Successful Response
+ * @throws ApiError
+ */
+export const usePoolServiceDeletePool = <
+  TData = Common.PoolServiceDeletePoolMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        poolName: string;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      poolName: string;
+    },
+    TContext
+  >({
+    mutationFn: ({ poolName }) =>
+      PoolService.deletePool({ poolName }) as unknown as Promise<TData>,
+    ...options,
+  });

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

@@ -40,6 +40,8 @@ import type {
   DeleteDagRunData,
   DeleteDagRunResponse,
   GetHealthResponse,
+  DeletePoolData,
+  DeletePoolResponse,
 } from "./types.gen";
 
 export class AssetService {
@@ -592,3 +594,32 @@ export class MonitorService {
     });
   }
 }
+
+export class PoolService {
+  /**
+   * Delete Pool
+   * Delete a pool entry.
+   * @param data The data for the request.
+   * @param data.poolName
+   * @returns void Successful Response
+   * @throws ApiError
+   */
+  public static deletePool(
+    data: DeletePoolData,
+  ): CancelablePromise<DeletePoolResponse> {
+    return __request(OpenAPI, {
+      method: "DELETE",
+      url: "/public/pools/{pool_name}",
+      path: {
+        pool_name: data.poolName,
+      },
+      errors: {
+        400: "Bad Request",
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
+}

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

@@ -466,6 +466,12 @@ export type DeleteDagRunResponse = void;
 
 export type GetHealthResponse = HealthInfoSchema;
 
+export type DeletePoolData = {
+  poolName: string;
+};
+
+export type DeletePoolResponse = void;
+
 export type $OpenApiTs = {
   "/ui/next_run_assets/{dag_id}": {
     get: {
@@ -937,4 +943,35 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/pools/{pool_name}": {
+    delete: {
+      req: DeletePoolData;
+      res: {
+        /**
+         * Successful Response
+         */
+        204: void;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
+  };
 };

tests/api_fastapi/core_api/routes/public/test_pools.py

@@ -0,0 +1,78 @@
+# 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 pytest
+
+from airflow.models.pool import Pool
+from airflow.utils.session import provide_session
+
+from tests_common.test_utils.db import clear_db_pools
+
+pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode]
+
+POOL1_NAME = "pool1"
+POOL1_SLOT = 3
+POOL1_INCLUDE_DEFERRED = True
+
+
+POOL2_NAME = "pool2"
+POOL2_SLOT = 10
+POOL2_INCLUDE_DEFERRED = False
+POOL2_DESCRIPTION = "Some Description"
+
+
+@provide_session
+def _create_pools(session) -> None:
+    pool1 = Pool(pool=POOL1_NAME, slots=POOL1_SLOT, include_deferred=POOL1_INCLUDE_DEFERRED)
+    pool2 = Pool(pool=POOL2_NAME, slots=POOL2_SLOT, include_deferred=POOL2_INCLUDE_DEFERRED)
+    session.add_all([pool1, pool2])
+
+
+class TestPools:
+    @pytest.fixture(autouse=True)
+    def setup(self) -> None:
+        clear_db_pools()
+
+    def teardown_method(self) -> None:
+        clear_db_pools()
+
+    def create_pools(self):
+        _create_pools()
+
+
+class TestDeletePool(TestPools):
+    def test_delete_should_respond_204(self, test_client, session):
+        self.create_pools()
+        pools = session.query(Pool).all()
+        assert len(pools) == 3
+        response = test_client.delete(f"/public/pools/{POOL1_NAME}")
+        assert response.status_code == 204
+        pools = session.query(Pool).all()
+        assert len(pools) == 2
+
+    def test_delete_should_respond_400(self, test_client):
+        response = test_client.delete("/public/pools/default_pool")
+        assert response.status_code == 400
+        body = response.json()
+        assert "Default Pool can't be deleted" == body["detail"]
+
+    def test_delete_should_respond_404(self, test_client):
+        response = test_client.delete(f"/public/pools/{POOL1_NAME}")
+        assert response.status_code == 404
+        body = response.json()
+        assert f"The Pool with name: `{POOL1_NAME}` was not found" == body["detail"]