feat: stateless warehouse connection test endpoint (#8060)

Co-authored-by: flagsmith-engineering[bot] <flagsmith-engineering[bot]@users.noreply.github.com>

Author: Zaimwa9

api/app/settings/common.py

@@ -601,6 +601,9 @@
         "WebhookScopeTypeEnum": ["organisation", "environment"],
         "SegmentRuleTypeEnum": "segments.models.SegmentRule.RULE_TYPES",
         "FeatureValueTypeEnum": ["integer", "string", "boolean"],
+        "WarehouseConnectionStatusEnum": (
+            "experimentation.models.WarehouseConnectionStatus.choices"
+        ),
     },
     "COMPONENT_NO_READ_ONLY_REQUIRED": True,
 }

api/experimentation/services.py

@@ -818,9 +818,13 @@ def _describe_verification_error(error: Exception) -> str:
     return "Verification failed."
 
 
-def verify_clickhouse_connection(connection: WarehouseConnection) -> None:
+def verify_clickhouse_connection(
+    connection: WarehouseConnection,
+    persist: bool = True,
+) -> None:
     """Run SELECT 1 against the customer's ClickHouse and set the status to
-    connected or errored; never raises."""
+    connected or errored; never raises. With persist=False, the status is only
+    set on the in-memory instance, allowing unsaved connections to be tested."""
     log = logger.bind(environment__id=connection.environment_id)
     try:
         log = log.bind(organisation__id=connection.environment.project.organisation_id)
@@ -847,7 +851,8 @@ def verify_clickhouse_connection(connection: WarehouseConnection) -> None:
     except Exception as error:
         connection.status = WarehouseConnectionStatus.ERRORED
         connection.status_detail = _describe_verification_error(error)
-        connection.save(update_fields=["status", "status_detail"])
+        if persist:
+            connection.save(update_fields=["status", "status_detail"])
         flagsmith_experimentation_warehouse_connection_verifications_total.labels(
             result="failure"
         ).inc()
@@ -856,7 +861,8 @@ def verify_clickhouse_connection(connection: WarehouseConnection) -> None:
 
     connection.status = WarehouseConnectionStatus.CONNECTED
     connection.status_detail = None
-    connection.save(update_fields=["status", "status_detail"])
+    if persist:
+        connection.save(update_fields=["status", "status_detail"])
     flagsmith_experimentation_warehouse_connection_verifications_total.labels(
         result="success"
     ).inc()

api/experimentation/views.py

@@ -7,7 +7,11 @@
 from django.shortcuts import get_object_or_404
 from django.utils import timezone
 from django.utils.decorators import method_decorator
-from drf_spectacular.utils import OpenApiParameter, extend_schema
+from drf_spectacular.utils import (
+    OpenApiParameter,
+    extend_schema,
+    inline_serializer,
+)
 from rest_framework import mixins, serializers, status
 from rest_framework.decorators import action
 from rest_framework.exceptions import Throttled, ValidationError
@@ -34,6 +38,7 @@
     ExperimentStatus,
     Metric,
     WarehouseConnection,
+    WarehouseConnectionStatus,
     WarehouseType,
 )
 from experimentation.permissions import (
@@ -97,6 +102,7 @@ def get_throttles(self) -> list[BaseThrottle]:
             "update",
             "partial_update",
             "test_warehouse_connection",
+            "test_warehouse_connection_config",
         ):
             self.throttle_scope = "warehouse_connection_write"
             return [*super().get_throttles(), ScopedRateThrottle()]
@@ -167,6 +173,52 @@ def test_warehouse_connection(self, request: Request, **kwargs: object) -> Respo
         serializer = self.get_serializer(connection)
         return Response(serializer.data)
 
+    @extend_schema(
+        operation_id=(
+            "api_v1_environments_warehouse_connections_"
+            "test_warehouse_connection_config_create"
+        ),
+        responses={
+            200: inline_serializer(
+                name="WarehouseConnectionTestResult",
+                fields={
+                    "status": serializers.ChoiceField(
+                        choices=WarehouseConnectionStatus.choices
+                    ),
+                    "status_detail": serializers.CharField(allow_null=True),
+                },
+            )
+        },
+    )
+    @action(
+        detail=False,
+        methods=["post"],
+        url_path="test-warehouse-connection",
+        url_name="test-warehouse-connection-config",
+    )
+    def test_warehouse_connection_config(
+        self, request: Request, **kwargs: object
+    ) -> Response:
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        if serializer.validated_data.get("warehouse_type") != WarehouseType.CLICKHOUSE:
+            return Response(
+                {
+                    "detail": "Connection testing is not supported for this warehouse type."
+                },
+                status=status.HTTP_400_BAD_REQUEST,
+            )
+        connection = WarehouseConnection(
+            environment=self._get_environment(),
+            warehouse_type=serializer.validated_data["warehouse_type"],
+            config=serializer.validated_data.get("config"),
+            credentials=serializer.validated_data.get("credentials"),
+        )
+        verify_clickhouse_connection(connection, persist=False)
+        return Response(
+            {"status": connection.status, "status_detail": connection.status_detail}
+        )
+
     def create(self, request: Request, *args: object, **kwargs: object) -> Response:
         environment = self._get_environment()
         serializer = self.get_serializer(data=request.data)

api/tests/unit/experimentation/test_views.py

@@ -1,6 +1,7 @@
 import socket
 
 import pytest
+from clickhouse_driver import errors as clickhouse_errors
 from django.urls import reverse
 from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
@@ -1090,7 +1091,13 @@ def test_test_warehouse_connection__non_admin__returns_403(
 
 @pytest.mark.parametrize(
     "action",
-    ["create", "update", "partial_update", "test_warehouse_connection"],
+    [
+        "create",
+        "update",
+        "partial_update",
+        "test_warehouse_connection",
+        "test_warehouse_connection_config",
+    ],
 )
 def test_get_throttles__write_actions__returns_scoped_throttle(action: str) -> None:
     # Given
@@ -1666,6 +1673,74 @@ def test_test_warehouse_connection__clickhouse__reverifies_and_returns_status(
     assert clickhouse_connection.status == WarehouseConnectionStatus.CONNECTED
 
 
+CLICKHOUSE_TEST_PAYLOAD = {
+    "warehouse_type": "clickhouse",
+    "config": {"host": "ch.example.com"},
+    "credentials": {"password": "hunter2"},
+}
+
+
+@pytest.mark.parametrize(
+    "data, client_side_effect, expected_status_code, expected_response",
+    [
+        pytest.param(
+            CLICKHOUSE_TEST_PAYLOAD,
+            None,
+            status.HTTP_200_OK,
+            {"status": "connected", "status_detail": None},
+            id="reachable",
+        ),
+        pytest.param(
+            CLICKHOUSE_TEST_PAYLOAD,
+            clickhouse_errors.NetworkError("boom"),
+            status.HTTP_200_OK,
+            {"status": "errored", "status_detail": "Could not connect to the host."},
+            id="unreachable",
+        ),
+        pytest.param(
+            {"warehouse_type": "flagsmith"},
+            None,
+            status.HTTP_400_BAD_REQUEST,
+            {"detail": "Connection testing is not supported for this warehouse type."},
+            id="non_clickhouse_type",
+        ),
+        pytest.param(
+            {"warehouse_type": "clickhouse", "config": {"host": "ch.example.com"}},
+            None,
+            status.HTTP_400_BAD_REQUEST,
+            {"credentials": {"password": "This field is required."}},
+            id="invalid_payload",
+        ),
+    ],
+)
+def test_test_warehouse_connection_config__payload__returns_expected_response(
+    admin_client: APIClient,
+    client_side_effect: Exception | None,
+    data: dict[str, object],
+    enable_features: EnableFeaturesFixture,
+    environment: Environment,
+    expected_response: dict[str, object],
+    expected_status_code: int,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    enable_features("experimentation_warehouse_connection")
+    mocker.patch("experimentation.services.Client", side_effect=client_side_effect)
+    url = reverse(
+        "api-v1:environments:experimentation:"
+        "warehouse-connections-test-warehouse-connection-config",
+        args=[environment.api_key],
+    )
+
+    # When
+    response = admin_client.post(url, data=data, format="json")
+
+    # Then
+    assert response.status_code == expected_status_code
+    assert response.json() == expected_response
+    assert not WarehouseConnection.objects.filter(environment=environment).exists()
+
+
 def test_patch__clickhouse_to_flagsmith__resets_connection_state(
     admin_client: APIClient,
     clickhouse_connection: WarehouseConnection,

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -597,7 +597,7 @@ Attributes:
 ### `warehouse.connection.connected`
 
 Logged at `info` from:
- - `api/experimentation/services.py:878`
+ - `api/experimentation/services.py:884`
 
 Attributes:
  - `environment.id`
@@ -615,7 +615,7 @@ Attributes:
 ### `warehouse.connection.verification_failed`
 
 Logged at `warning` from:
- - `api/experimentation/services.py:854`
+ - `api/experimentation/services.py:859`
 
 Attributes:
  - `environment.id`
@@ -625,7 +625,7 @@ Attributes:
 ### `warehouse.connection.verification_succeeded`
 
 Logged at `info` from:
- - `api/experimentation/services.py:863`
+ - `api/experimentation/services.py:869`
 
 Attributes:
  - `environment.id`

openapi.yaml

@@ -7090,6 +7090,40 @@ paths:
         - Master API Key: []
       tags:
         - Environments
+  '/api/v1/environments/{environment_api_key}/warehouse-connections/test-warehouse-connection/':
+    post:
+      operationId: api_v1_environments_warehouse_connections_test_warehouse_connection_config_create
+      description: Manage data warehouse connections for an environment.
+      parameters:
+        - name: environment_api_key
+          in: path
+          required: true
+          schema:
+            type: string
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/WarehouseConnection'
+          application/x-www-form-urlencoded:
+            schema:
+              $ref: '#/components/schemas/WarehouseConnection'
+          multipart/form-data:
+            schema:
+              $ref: '#/components/schemas/WarehouseConnection'
+      responses:
+        '200':
+          description: ''
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/WarehouseConnectionTestResult'
+      security:
+        - tokenAuth: []
+        - Master API Key: []
+      tags:
+        - Environments
   '/api/v1/environments/{environment_api_key}/webhooks/':
     get:
       operationId: api_v1_environments_webhooks_list
@@ -27862,6 +27896,18 @@ components:
         - pending_connection
         - connected
         - errored
+    WarehouseConnectionTestResult:
+      type: object
+      properties:
+        status:
+          $ref: '#/components/schemas/WarehouseConnectionStatusEnum'
+        status_detail:
+          type:
+            - string
+            - 'null'
+      required:
+        - status
+        - status_detail
     WarehouseTypeEnum:
       description: |-
         * `flagsmith` - Flagsmith