Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions api/app/settings/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down
14 changes: 10 additions & 4 deletions api/experimentation/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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()
Expand All @@ -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()
Expand Down
54 changes: 53 additions & 1 deletion api/experimentation/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -34,6 +38,7 @@
ExperimentStatus,
Metric,
WarehouseConnection,
WarehouseConnectionStatus,
WarehouseType,
)
from experimentation.permissions import (
Expand Down Expand Up @@ -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()]
Expand Down Expand Up @@ -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),
},
)
},
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
@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)
Expand Down
77 changes: 76 additions & 1 deletion api/tests/unit/experimentation/test_views.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ Attributes:
### `warehouse.connection.connected`

Logged at `info` from:
- `api/experimentation/services.py:878`
- `api/experimentation/services.py:884`

Attributes:
- `environment.id`
Expand All @@ -578,7 +578,7 @@ Attributes:
### `warehouse.connection.verification_failed`

Logged at `warning` from:
- `api/experimentation/services.py:854`
- `api/experimentation/services.py:859`

Attributes:
- `environment.id`
Expand All @@ -588,7 +588,7 @@ Attributes:
### `warehouse.connection.verification_succeeded`

Logged at `info` from:
- `api/experimentation/services.py:863`
- `api/experimentation/services.py:869`

Attributes:
- `environment.id`
Expand Down
46 changes: 46 additions & 0 deletions openapi.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading