From 9201077390b262e356df52e53145a3d537b220a4 Mon Sep 17 00:00:00 2001 From: pierrejeambrun Date: Fri, 18 Oct 2024 16:48:09 +0200 Subject: [PATCH] AIP-84 Delete Pool --- .../api_connexion/endpoints/pool_endpoint.py | 10 ++- .../core_api/openapi/v1-generated.yaml | 47 +++++++++++ .../core_api/routes/public/__init__.py | 2 + .../core_api/routes/public/pools.py | 48 ++++++++++++ airflow/ui/openapi-gen/queries/common.ts | 4 + airflow/ui/openapi-gen/queries/queries.ts | 38 +++++++++ .../ui/openapi-gen/requests/services.gen.ts | 31 ++++++++ airflow/ui/openapi-gen/requests/types.gen.ts | 37 +++++++++ .../core_api/routes/public/test_pools.py | 78 +++++++++++++++++++ 9 files changed, 292 insertions(+), 3 deletions(-) create mode 100644 airflow/api_fastapi/core_api/routes/public/pools.py create mode 100644 tests/api_fastapi/core_api/routes/public/test_pools.py diff --git a/airflow/api_connexion/endpoints/pool_endpoint.py b/airflow/api_connexion/endpoints/pool_endpoint.py index 553d50c7464b7..6ea28f7457ee6 100644 --- a/airflow/api_connexion/endpoints/pool_endpoint.py +++ b/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: diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 88e292364767f..97aac85644a8c 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/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: diff --git a/airflow/api_fastapi/core_api/routes/public/__init__.py b/airflow/api_fastapi/core_api/routes/public/__init__.py index 3d43a7bbb0efb..12d98e44199da 100644 --- a/airflow/api_fastapi/core_api/routes/public/__init__.py +++ b/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) diff --git a/airflow/api_fastapi/core_api/routes/public/pools.py b/airflow/api_fastapi/core_api/routes/public/pools.py new file mode 100644 index 0000000000000..2b4ffc02632c3 --- /dev/null +++ b/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") diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 7c0ab1c646163..3a31b4c0ad993 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/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 >; +export type PoolServiceDeletePoolMutationResult = Awaited< + ReturnType +>; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 9d3e640fbf98d..6d358a7008a50 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/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, ...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, + ...options, + }); diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 1561c024a83d1..64bf06c4d6f6f 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/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 { + 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", + }, + }); + } +} diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index aae457b13afd1..ef3847b23f2e0 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/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; + }; + }; + }; }; diff --git a/tests/api_fastapi/core_api/routes/public/test_pools.py b/tests/api_fastapi/core_api/routes/public/test_pools.py new file mode 100644 index 0000000000000..4a5c511b8dd50 --- /dev/null +++ b/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"]