Add async connection testing via workers for security isolation

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``1d6611b6ab7c`` (head) | ``888b59e02a5b`` | ``3.2.0``         | Add bundle_name to callback table.                           |
+| ``a7e6d4c3b2f1`` (head) | ``1d6611b6ab7c`` | ``3.2.0``         | Add connection_test_request table for async connection       |
+|                         |                  |                   | testing.                                                     |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``1d6611b6ab7c``        | ``888b59e02a5b`` | ``3.2.0``         | Add bundle_name to callback table.                           |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``888b59e02a5b``        | ``6222ce48e289`` | ``3.2.0``         | Fix migration file ORM inconsistencies.                      |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py

@@ -19,6 +19,7 @@
 
 import json
 from collections.abc import Iterable, Mapping
+from datetime import datetime
 from typing import Annotated, Any
 
 from pydantic import Field, field_validator
@@ -72,12 +73,46 @@ class ConnectionCollectionResponse(BaseModel):
 
 
 class ConnectionTestResponse(BaseModel):
-    """Connection Test serializer for responses."""
+    """Connection Test serializer for synchronous test responses."""
 
     status: bool
     message: str
 
 
+class ConnectionTestRequestBody(StrictBaseModel):
+    """Request body for async connection test."""
+
+    connection_id: str
+    conn_type: str
+    host: str | None = None
+    login: str | None = None
+    schema_: str | None = Field(None, alias="schema")
+    port: int | None = None
+    password: str | None = None
+    extra: str | None = None
+    commit_on_success: bool = False
+    executor: str | None = None
+    queue: str | None = None
+
+
+class ConnectionTestQueuedResponse(BaseModel):
+    """Response returned when an async connection test is queued."""
+
+    token: str
+    connection_id: str
+    state: str
+
+
+class ConnectionTestStatusResponse(BaseModel):
+    """Response returned when polling for async connection test status."""
+
+    token: str
+    connection_id: str
+    state: str
+    result_message: str | None = None
+    created_at: datetime
+
+
 class ConnectionHookFieldBehavior(BaseModel):
     """A class to store the behavior of each standard field of a Hook."""
 

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -1335,6 +1335,12 @@ paths:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
           description: Not Found
+        '409':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Conflict
         '422':
           description: Validation Error
           content:
@@ -1681,6 +1687,110 @@ paths:
       security:
       - OAuth2PasswordBearer: []
       - HTTPBearer: []
+  /api/v2/connections/test-async:
+    post:
+      tags:
+      - Connection
+      summary: Test Connection Async
+      description: 'Queue an async connection test to be executed on a worker.
+
+
+        The connection data is stored in the test request table and the worker
+
+        reads from there. Returns a token to poll for the result via
+
+        GET /connections/test-async/{token}.'
+      operationId: test_connection_async
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/ConnectionTestRequestBody'
+        required: true
+      responses:
+        '202':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ConnectionTestQueuedResponse'
+        '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'
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+  /api/v2/connections/test-async/{connection_test_token}:
+    get:
+      tags:
+      - Connection
+      summary: Get Connection Test
+      description: "Poll for the status of an async connection test.\n\nKnowledge\
+        \ of the token serves as authorization \u2014 only the client\nthat initiated\
+        \ the test knows the crypto-random token."
+      operationId: get_connection_test
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: connection_test_token
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Connection Test Token
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ConnectionTestStatusResponse'
+        '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'
   /api/v2/connections/defaults:
     post:
       tags:
@@ -10096,6 +10206,83 @@ components:
       - team_name
       title: ConnectionResponse
       description: Connection serializer for responses.
+    ConnectionTestQueuedResponse:
+      properties:
+        token:
+          type: string
+          title: Token
+        connection_id:
+          type: string
+          title: Connection Id
+        state:
+          type: string
+          title: State
+      type: object
+      required:
+      - token
+      - connection_id
+      - state
+      title: ConnectionTestQueuedResponse
+      description: Response returned when an async connection test is queued.
+    ConnectionTestRequestBody:
+      properties:
+        connection_id:
+          type: string
+          title: Connection Id
+        conn_type:
+          type: string
+          title: Conn Type
+        host:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Host
+        login:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Login
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Schema
+        port:
+          anyOf:
+          - type: integer
+          - type: 'null'
+          title: Port
+        password:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Password
+        extra:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Extra
+        commit_on_success:
+          type: boolean
+          title: Commit On Success
+          default: false
+        executor:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Executor
+        queue:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Queue
+      additionalProperties: false
+      type: object
+      required:
+      - connection_id
+      - conn_type
+      title: ConnectionTestRequestBody
+      description: Request body for async connection test.
     ConnectionTestResponse:
       properties:
         status:
@@ -10109,7 +10296,35 @@ components:
       - status
       - message
       title: ConnectionTestResponse
-      description: Connection Test serializer for responses.
+      description: Connection Test serializer for synchronous test responses.
+    ConnectionTestStatusResponse:
+      properties:
+        token:
+          type: string
+          title: Token
+        connection_id:
+          type: string
+          title: Connection Id
+        state:
+          type: string
+          title: State
+        result_message:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Result Message
+        created_at:
+          type: string
+          format: date-time
+          title: Created At
+      type: object
+      required:
+      - token
+      - connection_id
+      - state
+      - created_at
+      title: ConnectionTestStatusResponse
+      description: Response returned when polling for async connection test status.
     CreateAssetEventsBody:
       properties:
         asset_id: