AIP-81 Implement Create Default Connections Endpoint in REST API (Fas…

…tAPI) (apache#45363) * Add Create Default Connections Endpoint in FastAPI * Include mock db_create_default_connections and ensure it is called at least once
aws-mwaa · Jan 8, 2025 · c300e0e · c300e0e
1 parent 301017d
commit c300e0e
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 0 deletions.
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -1834,6 +1834,28 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /public/connections/defaults:
+    post:
+      tags:
+      - Connection
+      summary: Create Default Connections
+      description: Create default connections.
+      operationId: create_default_connections
+      responses:
+        '204':
+          description: Successful Response
+        '401':
+          description: Unauthorized
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
   /public/dags/{dag_id}/dagRuns/{dag_run_id}:
     get:
       tags:

diff --git a/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow/api_fastapi/core_api/routes/public/connections.py
@@ -38,6 +38,7 @@
 from airflow.configuration import conf
 from airflow.models import Connection
 from airflow.secrets.environment_variables import CONN_ENV_PREFIX
+from airflow.utils.db import create_default_connections as db_create_default_connections
 from airflow.utils.strings import get_random_string
 
 connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
@@ -262,3 +263,14 @@ def test_connection(
         return ConnectionTestResponse.model_validate({"status": test_status, "message": test_message})
     finally:
         os.environ.pop(conn_env_var, None)
+
+
+@connections_router.post(
+    "/defaults",
+    status_code=status.HTTP_204_NO_CONTENT,
+)
+def create_default_connections(
+    session: SessionDep,
+):
+    """Create default connections."""
+    db_create_default_connections(session)
diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -1602,6 +1602,9 @@ export type ConnectionServicePostConnectionMutationResult = Awaited<
 export type ConnectionServiceTestConnectionMutationResult = Awaited<
   ReturnType<typeof ConnectionService.testConnection>
 >;
+export type ConnectionServiceCreateDefaultConnectionsMutationResult = Awaited<
+  ReturnType<typeof ConnectionService.createDefaultConnections>
+>;
 export type DagRunServiceClearDagRunMutationResult = Awaited<ReturnType<typeof DagRunService.clearDagRun>>;
 export type DagRunServiceTriggerDagRunMutationResult = Awaited<
   ReturnType<typeof DagRunService.triggerDagRun>

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -2788,6 +2788,23 @@ export const useConnectionServiceTestConnection = <
       ConnectionService.testConnection({ requestBody }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Create Default Connections
+ * Create default connections.
+ * @returns void Successful Response
+ * @throws ApiError
+ */
+export const useConnectionServiceCreateDefaultConnections = <
+  TData = Common.ConnectionServiceCreateDefaultConnectionsMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<UseMutationOptions<TData, TError, void, TContext>, "mutationFn">,
+) =>
+  useMutation<TData, TError, void, TContext>({
+    mutationFn: () => ConnectionService.createDefaultConnections() as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Clear Dag Run
  * @param data The data for the request.

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -70,6 +70,7 @@ import type {
   PutConnectionsResponse,
   TestConnectionData,
   TestConnectionResponse,
+  CreateDefaultConnectionsResponse,
   GetDagRunData,
   GetDagRunResponse,
   DeleteDagRunData,
@@ -1156,6 +1157,23 @@ export class ConnectionService {
       },
     });
   }
+
+  /**
+   * Create Default Connections
+   * Create default connections.
+   * @returns void Successful Response
+   * @throws ApiError
+   */
+  public static createDefaultConnections(): CancelablePromise<CreateDefaultConnectionsResponse> {
+    return __request(OpenAPI, {
+      method: "POST",
+      url: "/public/connections/defaults",
+      errors: {
+        401: "Unauthorized",
+        403: "Forbidden",
+      },
+    });
+  }
 }
 
 export class DagRunService {

diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -1631,6 +1631,8 @@ export type TestConnectionData = {
 
 export type TestConnectionResponse = ConnectionTestResponse;
 
+export type CreateDefaultConnectionsResponse = void;
+
 export type GetDagRunData = {
   dagId: string;
   dagRunId: string;
@@ -3044,6 +3046,24 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/connections/defaults": {
+    post: {
+      res: {
+        /**
+         * Successful Response
+         */
+        204: void;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+      };
+    };
+  };
   "/public/dags/{dag_id}/dagRuns/{dag_run_id}": {
     get: {
       req: GetDagRunData;

diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py
@@ -1045,3 +1045,16 @@ def test_should_respond_403_by_default(self, test_client, body):
             "detail": "Testing connections is disabled in Airflow configuration. "
             "Contact your deployment admin to enable it."
         }
+
+
+class TestCreateDefaultConnections(TestConnectionEndpoint):
+    def test_should_respond_204(self, test_client):
+        response = test_client.post("/public/connections/defaults")
+        assert response.status_code == 204
+        assert response.content == b""
+
+    @mock.patch("airflow.api_fastapi.core_api.routes.public.connections.db_create_default_connections")
+    def test_should_call_db_create_default_connections(self, mock_db_create_default_connections, test_client):
+        response = test_client.post("/public/connections/defaults")
+        assert response.status_code == 204
+        mock_db_create_default_connections.assert_called_once()