AIP-81 Implement POST/Insert Multiple Connections in REST API (FastAPI) (#44531)

* Move validate key to pydantic model, leave unique check to database session for post endpoints, include bulk connection insert endpoint

* Fix test naming

Author: bugraoz93

airflow/api_fastapi/core_api/datamodels/connections.py

@@ -79,7 +79,7 @@ class ConnectionTestResponse(BaseModel):
 class ConnectionBody(BaseModel):
     """Connection Serializer for requests body."""
 
-    connection_id: str = Field(serialization_alias="conn_id")
+    connection_id: str = Field(serialization_alias="conn_id", max_length=200, pattern=r"^[\w.-]+$")
     conn_type: str
     description: str | None = Field(default=None)
     host: str | None = Field(default=None)
@@ -88,3 +88,9 @@ class ConnectionBody(BaseModel):
     port: int | None = Field(default=None)
     password: str | None = Field(default=None)
     extra: str | None = Field(default=None)
+
+
+class ConnectionBulkBody(BaseModel):
+    """Connections Serializer for requests body."""
+
+    connections: list[ConnectionBody]

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

@@ -1399,6 +1399,50 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /public/connections/bulk:
+    post:
+      tags:
+      - Connection
+      summary: Post Connections
+      description: Create connection entry.
+      operationId: post_connections
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/ConnectionBulkBody'
+        required: true
+      responses:
+        '201':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ConnectionCollectionResponse'
+        '401':
+          description: Unauthorized
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '409':
+          description: Conflict
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
   /public/connections/test:
     post:
       tags:
@@ -6207,6 +6251,8 @@ components:
       properties:
         connection_id:
           type: string
+          maxLength: 200
+          pattern: ^[\w.-]+$
           title: Connection Id
         conn_type:
           type: string
@@ -6252,6 +6298,18 @@ components:
       - conn_type
       title: ConnectionBody
       description: Connection Serializer for requests body.
+    ConnectionBulkBody:
+      properties:
+        connections:
+          items:
+            $ref: '#/components/schemas/ConnectionBody'
+          type: array
+          title: Connections
+      type: object
+      required:
+      - connections
+      title: ConnectionBulkBody
+      description: Connections Serializer for requests body.
     ConnectionCollectionResponse:
       properties:
         connections:

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

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import os
-from typing import Annotated
+from typing import Annotated, cast
 
 from fastapi import Depends, HTTPException, Query, status
 from sqlalchemy import select
@@ -27,6 +27,7 @@
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.datamodels.connections import (
     ConnectionBody,
+    ConnectionBulkBody,
     ConnectionCollectionResponse,
     ConnectionResponse,
     ConnectionTestResponse,
@@ -35,7 +36,6 @@
 from airflow.configuration import conf
 from airflow.models import Connection
 from airflow.secrets.environment_variables import CONN_ENV_PREFIX
-from airflow.utils import helpers
 from airflow.utils.strings import get_random_string
 
 connections_router = AirflowRouter(tags=["Connection"], prefix="/connections")
@@ -126,24 +126,29 @@ def post_connection(
     session: SessionDep,
 ) -> ConnectionResponse:
     """Create connection entry."""
-    try:
-        helpers.validate_key(post_body.connection_id, max_length=200)
-    except Exception as e:
-        raise HTTPException(status.HTTP_400_BAD_REQUEST, f"{e}")
-
-    connection = session.scalar(select(Connection).filter_by(conn_id=post_body.connection_id))
-    if connection is not None:
-        raise HTTPException(
-            status.HTTP_409_CONFLICT,
-            f"Connection with connection_id: `{post_body.connection_id}` already exists",
-        )
-
     connection = Connection(**post_body.model_dump(by_alias=True))
     session.add(connection)
-
     return connection
 
 
+@connections_router.post(
+    "/bulk",
+    status_code=status.HTTP_201_CREATED,
+    responses=create_openapi_http_exception_doc([status.HTTP_409_CONFLICT]),
+)
+def post_connections(
+    post_body: ConnectionBulkBody,
+    session: SessionDep,
+) -> ConnectionCollectionResponse:
+    """Create connection entry."""
+    connections = [Connection(**body.model_dump(by_alias=True)) for body in post_body.connections]
+    session.add_all(connections)
+    return ConnectionCollectionResponse(
+        connections=cast(list[ConnectionResponse], connections),
+        total_entries=len(connections),
+    )
+
+
 @connections_router.patch(
     "/{connection_id}",
     responses=create_openapi_http_exception_doc(

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

@@ -1594,6 +1594,9 @@ export type BackfillServiceCreateBackfillMutationResult = Awaited<
 export type ConnectionServicePostConnectionMutationResult = Awaited<
   ReturnType<typeof ConnectionService.postConnection>
 >;
+export type ConnectionServicePostConnectionsMutationResult = Awaited<
+  ReturnType<typeof ConnectionService.postConnections>
+>;
 export type ConnectionServiceTestConnectionMutationResult = Awaited<
   ReturnType<typeof ConnectionService.testConnection>
 >;

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

@@ -38,6 +38,7 @@ import {
   BackfillPostBody,
   ClearTaskInstancesBody,
   ConnectionBody,
+  ConnectionBulkBody,
   CreateAssetEventsBody,
   DAGPatchBody,
   DAGRunClearBody,
@@ -2684,6 +2685,45 @@ export const useConnectionServicePostConnection = <
       }) as unknown as Promise<TData>,
     ...options,
   });
+/**
+ * Post Connections
+ * Create connection entry.
+ * @param data The data for the request.
+ * @param data.requestBody
+ * @returns ConnectionCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const useConnectionServicePostConnections = <
+  TData = Common.ConnectionServicePostConnectionsMutationResult,
+  TError = unknown,
+  TContext = unknown,
+>(
+  options?: Omit<
+    UseMutationOptions<
+      TData,
+      TError,
+      {
+        requestBody: ConnectionBulkBody;
+      },
+      TContext
+    >,
+    "mutationFn"
+  >,
+) =>
+  useMutation<
+    TData,
+    TError,
+    {
+      requestBody: ConnectionBulkBody;
+    },
+    TContext
+  >({
+    mutationFn: ({ requestBody }) =>
+      ConnectionService.postConnections({
+        requestBody,
+      }) as unknown as Promise<TData>,
+    ...options,
+  });
 /**
  * Test Connection
  * Test an API connection.

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

@@ -750,6 +750,8 @@ export const $ConnectionBody = {
   properties: {
     connection_id: {
       type: "string",
+      maxLength: 200,
+      pattern: "^[\\w.-]+$",
       title: "Connection Id",
     },
     conn_type: {
@@ -840,6 +842,22 @@ export const $ConnectionBody = {
   description: "Connection Serializer for requests body.",
 } as const;
 
+export const $ConnectionBulkBody = {
+  properties: {
+    connections: {
+      items: {
+        $ref: "#/components/schemas/ConnectionBody",
+      },
+      type: "array",
+      title: "Connections",
+    },
+  },
+  type: "object",
+  required: ["connections"],
+  title: "ConnectionBulkBody",
+  description: "Connections Serializer for requests body.",
+} as const;
+
 export const $ConnectionCollectionResponse = {
   properties: {
     connections: {

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

@@ -58,6 +58,8 @@ import type {
   GetConnectionsResponse,
   PostConnectionData,
   PostConnectionResponse,
+  PostConnectionsData,
+  PostConnectionsResponse,
   TestConnectionData,
   TestConnectionResponse,
   GetDagRunData,
@@ -997,6 +999,31 @@ export class ConnectionService {
     });
   }
 
+  /**
+   * Post Connections
+   * Create connection entry.
+   * @param data The data for the request.
+   * @param data.requestBody
+   * @returns ConnectionCollectionResponse Successful Response
+   * @throws ApiError
+   */
+  public static postConnections(
+    data: PostConnectionsData,
+  ): CancelablePromise<PostConnectionsResponse> {
+    return __request(OpenAPI, {
+      method: "POST",
+      url: "/public/connections/bulk",
+      body: data.requestBody,
+      mediaType: "application/json",
+      errors: {
+        401: "Unauthorized",
+        403: "Forbidden",
+        409: "Conflict",
+        422: "Validation Error",
+      },
+    });
+  }
+
   /**
    * Test Connection
    * Test an API connection.

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

@@ -214,6 +214,13 @@ export type ConnectionBody = {
   extra?: string | null;
 };
 
+/**
+ * Connections Serializer for requests body.
+ */
+export type ConnectionBulkBody = {
+  connections: Array<ConnectionBody>;
+};
+
 /**
  * Connection Collection serializer for responses.
  */
@@ -1499,6 +1506,12 @@ export type PostConnectionData = {
 
 export type PostConnectionResponse = ConnectionResponse;
 
+export type PostConnectionsData = {
+  requestBody: ConnectionBulkBody;
+};
+
+export type PostConnectionsResponse = ConnectionCollectionResponse;
+
 export type TestConnectionData = {
   requestBody: ConnectionBody;
 };
@@ -2766,6 +2779,33 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/connections/bulk": {
+    post: {
+      req: PostConnectionsData;
+      res: {
+        /**
+         * Successful Response
+         */
+        201: ConnectionCollectionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Conflict
+         */
+        409: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
+  };
   "/public/connections/test": {
     post: {
       req: TestConnectionData;