From 9bd026d213a2ab4e8b5d11541ea135ac4ad3dbfa Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Wed, 2 Oct 2024 10:47:56 +0200 Subject: [PATCH] AIP-84 Migrate delete a connection to FastAPI API (#42571) * Include connections router and migrate delete a connection endpoint to fastapi * Mark tests as db_test * Use only pyfixture session * make method async * setup method to setup_attrs * Convert APIRouter tags, make setup method unified * Use AirflowRouter over fastapi.APIRouter --- .../endpoints/connection_endpoint.py | 2 + airflow/api_fastapi/openapi/v1-generated.yaml | 41 ++++++++++++ airflow/api_fastapi/views/public/__init__.py | 2 + .../api_fastapi/views/public/connections.py | 47 ++++++++++++++ airflow/ui/openapi-gen/queries/common.ts | 9 ++- airflow/ui/openapi-gen/queries/queries.ts | 45 ++++++++++++- .../ui/openapi-gen/requests/services.gen.ts | 30 +++++++++ airflow/ui/openapi-gen/requests/types.gen.ts | 33 ++++++++++ .../views/public/test_connections.py | 63 +++++++++++++++++++ 9 files changed, 270 insertions(+), 2 deletions(-) create mode 100644 airflow/api_fastapi/views/public/connections.py create mode 100644 tests/api_fastapi/views/public/test_connections.py diff --git a/airflow/api_connexion/endpoints/connection_endpoint.py b/airflow/api_connexion/endpoints/connection_endpoint.py index c17a9280d78f..b28c9dfcafa7 100644 --- a/airflow/api_connexion/endpoints/connection_endpoint.py +++ b/airflow/api_connexion/endpoints/connection_endpoint.py @@ -40,6 +40,7 @@ from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.security import permissions from airflow.utils import helpers +from airflow.utils.api_migration import mark_fastapi_migration_done from airflow.utils.log.action_logger import action_event_from_permission from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.strings import get_random_string @@ -53,6 +54,7 @@ RESOURCE_EVENT_PREFIX = "connection" +@mark_fastapi_migration_done @security.requires_access_connection("DELETE") @provide_session @action_logging( diff --git a/airflow/api_fastapi/openapi/v1-generated.yaml b/airflow/api_fastapi/openapi/v1-generated.yaml index b08ef42c16df..a54e0e4ca57d 100644 --- a/airflow/api_fastapi/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/openapi/v1-generated.yaml @@ -319,6 +319,47 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /public/connections/{connection_id}: + delete: + tags: + - Connection + summary: Delete Connection + description: Delete a connection entry. + operationId: delete_connection + parameters: + - name: connection_id + in: path + required: true + schema: + type: string + title: Connection Id + responses: + '204': + description: Successful Response + '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: DAGCollectionResponse: diff --git a/airflow/api_fastapi/views/public/__init__.py b/airflow/api_fastapi/views/public/__init__.py index 1c2511fc82ac..9c0eefebb875 100644 --- a/airflow/api_fastapi/views/public/__init__.py +++ b/airflow/api_fastapi/views/public/__init__.py @@ -17,6 +17,7 @@ from __future__ import annotations +from airflow.api_fastapi.views.public.connections import connections_router from airflow.api_fastapi.views.public.dags import dags_router from airflow.api_fastapi.views.router import AirflowRouter @@ -24,3 +25,4 @@ public_router.include_router(dags_router) +public_router.include_router(connections_router) diff --git a/airflow/api_fastapi/views/public/connections.py b/airflow/api_fastapi/views/public/connections.py new file mode 100644 index 000000000000..d418e1002679 --- /dev/null +++ b/airflow/api_fastapi/views/public/connections.py @@ -0,0 +1,47 @@ +# 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 select +from sqlalchemy.orm import Session +from typing_extensions import Annotated + +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.views.router import AirflowRouter +from airflow.models import Connection + +connections_router = AirflowRouter(tags=["Connection"]) + + +@connections_router.delete( + "/connections/{connection_id}", + status_code=204, + responses=create_openapi_http_exception_doc([401, 403, 404]), +) +async def delete_connection( + connection_id: str, + session: Annotated[Session, Depends(get_session)], +): + """Delete a connection entry.""" + connection = session.scalar(select(Connection).filter_by(conn_id=connection_id)) + + if connection is None: + raise HTTPException(404, f"The Connection with connection_id: `{connection_id}` was not found") + + session.delete(connection) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 96e49cc6d767..fcddded7dc12 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -1,7 +1,11 @@ // generated with @7nohe/openapi-react-query-codegen@1.6.0 import { UseQueryResult } from "@tanstack/react-query"; -import { AssetService, DagService } from "../requests/services.gen"; +import { + AssetService, + ConnectionService, + DagService, +} from "../requests/services.gen"; import { DagRunState } from "../requests/types.gen"; export type AssetServiceNextRunAssetsDefaultResponse = Awaited< @@ -76,3 +80,6 @@ export type DagServicePatchDagsMutationResult = Awaited< export type DagServicePatchDagMutationResult = Awaited< ReturnType >; +export type ConnectionServiceDeleteConnectionMutationResult = Awaited< + ReturnType +>; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 985bf952e3eb..f83c151b91e2 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -6,7 +6,11 @@ import { UseQueryOptions, } from "@tanstack/react-query"; -import { AssetService, DagService } from "../requests/services.gen"; +import { + AssetService, + ConnectionService, + DagService, +} from "../requests/services.gen"; import { DAGPatchBody, DagRunState } from "../requests/types.gen"; import * as Common from "./common"; @@ -247,3 +251,42 @@ export const useDagServicePatchDag = < }) as unknown as Promise, ...options, }); +/** + * Delete Connection + * Delete a connection entry. + * @param data The data for the request. + * @param data.connectionId + * @returns void Successful Response + * @throws ApiError + */ +export const useConnectionServiceDeleteConnection = < + TData = Common.ConnectionServiceDeleteConnectionMutationResult, + TError = unknown, + TContext = unknown, +>( + options?: Omit< + UseMutationOptions< + TData, + TError, + { + connectionId: string; + }, + TContext + >, + "mutationFn" + >, +) => + useMutation< + TData, + TError, + { + connectionId: string; + }, + TContext + >({ + mutationFn: ({ connectionId }) => + ConnectionService.deleteConnection({ + connectionId, + }) 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 be216bd534c6..24c960d2b7d5 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -11,6 +11,8 @@ import type { PatchDagsResponse, PatchDagData, PatchDagResponse, + DeleteConnectionData, + DeleteConnectionResponse, } from "./types.gen"; export class AssetService { @@ -159,3 +161,31 @@ export class DagService { }); } } + +export class ConnectionService { + /** + * Delete Connection + * Delete a connection entry. + * @param data The data for the request. + * @param data.connectionId + * @returns void Successful Response + * @throws ApiError + */ + public static deleteConnection( + data: DeleteConnectionData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "DELETE", + url: "/public/connections/{connection_id}", + path: { + connection_id: data.connectionId, + }, + errors: { + 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 e1db8310a1dc..b38d5c00a69f 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -134,6 +134,12 @@ export type PatchDagData = { export type PatchDagResponse = DAGResponse; +export type DeleteConnectionData = { + connectionId: string; +}; + +export type DeleteConnectionResponse = void; + export type $OpenApiTs = { "/ui/next_run_datasets/{dag_id}": { get: { @@ -227,4 +233,31 @@ export type $OpenApiTs = { }; }; }; + "/public/connections/{connection_id}": { + delete: { + req: DeleteConnectionData; + res: { + /** + * Successful Response + */ + 204: void; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; }; diff --git a/tests/api_fastapi/views/public/test_connections.py b/tests/api_fastapi/views/public/test_connections.py new file mode 100644 index 000000000000..cfdca1d67984 --- /dev/null +++ b/tests/api_fastapi/views/public/test_connections.py @@ -0,0 +1,63 @@ +# 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 import Connection +from airflow.utils.session import provide_session +from tests.test_utils.db import clear_db_connections + +pytestmark = pytest.mark.db_test + +TEST_CONN_ID = "test_connection_id" +TEST_CONN_TYPE = "test_type" + + +@provide_session +def _create_connection(session) -> None: + connection_model = Connection(conn_id=TEST_CONN_ID, conn_type=TEST_CONN_TYPE) + session.add(connection_model) + + +class TestConnectionEndpoint: + @pytest.fixture(autouse=True) + def setup(self) -> None: + clear_db_connections(False) + + def teardown_method(self) -> None: + clear_db_connections() + + def create_connection(self): + _create_connection() + + +class TestDeleteConnection(TestConnectionEndpoint): + def test_delete_should_respond_204(self, test_client, session): + self.create_connection() + conns = session.query(Connection).all() + assert len(conns) == 1 + response = test_client.delete(f"/public/connections/{TEST_CONN_ID}") + assert response.status_code == 204 + connection = session.query(Connection).all() + assert len(connection) == 0 + + def test_delete_should_respond_404(self, test_client): + response = test_client.delete(f"/public/connections/{TEST_CONN_ID}") + assert response.status_code == 404 + body = response.json() + assert f"The Connection with connection_id: `{TEST_CONN_ID}` was not found" == body["detail"]