feat: verify ClickHouse events table and query customer event stats (#8116)

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

Author: Zaimwa9

api/app/settings/common.py

@@ -367,6 +367,7 @@
         "user": USER_THROTTLE_RATE,
         "influx_query": "5/min",
         "warehouse_connection_write": "10/min",
+        "warehouse_connection_read": "60/min",
     },
     "DEFAULT_FILTER_BACKENDS": ["django_filters.rest_framework.DjangoFilterBackend"],
     "DEFAULT_RENDERER_CLASSES": [

api/app/settings/test.py

@@ -28,6 +28,7 @@
     "master_api_key": "100000/day",
     "influx_query": "50/min",
     "warehouse_connection_write": "1000/min",
+    "warehouse_connection_read": "1000/min",
 }
 
 AWS_SSE_LOGS_BUCKET_NAME = "test_bucket"

api/experimentation/services.py

@@ -9,6 +9,7 @@
 from clickhouse_driver import errors as clickhouse_errors
 from clickhouse_driver.util.helpers import parse_url
 from django.conf import settings
+from django.core.cache import cache
 from django.db import transaction
 from django.db.models import Q
 from django.utils import timezone
@@ -95,6 +96,9 @@
 CLICKHOUSE_CONNECT_TIMEOUT_SECONDS = 5
 CLICKHOUSE_QUERY_TIMEOUT_SECONDS = 30
 CLICKHOUSE_VERIFY_TIMEOUT_SECONDS = 5
+CUSTOMER_EVENT_STATS_CACHE_SECONDS = 60
+
+_CUSTOMER_EVENT_STATS_UNAVAILABLE = "unavailable"
 
 
 def is_warehouse_feature_enabled(organisation: Organisation) -> bool:
@@ -140,9 +144,8 @@ def get_unique_event_names(environment_key: str) -> list[str]:
     return [row[0] for row in rows]
 
 
-def get_warehouse_event_stats(environment_key: str) -> WarehouseEventStats:
-    """Return event counts recorded for `environment_key` in the warehouse."""
-    rows = _get_clickhouse_client().execute(
+def _query_event_stats(client: Client, environment_key: str) -> WarehouseEventStats:
+    rows = client.execute(
         "SELECT count() AS total, uniqExact(event) AS unique "
         "FROM events WHERE environment_key = %(environment_key)s",
         {"environment_key": environment_key},
@@ -154,6 +157,11 @@ def get_warehouse_event_stats(environment_key: str) -> WarehouseEventStats:
     )
 
 
+def get_warehouse_event_stats(environment_key: str) -> WarehouseEventStats:
+    """Return event counts recorded for `environment_key` in the warehouse."""
+    return _query_event_stats(_get_clickhouse_client(), environment_key)
+
+
 EXPOSURE_BUCKETS_QUERY = (
     _EXPOSURES_CTE
     + """
@@ -799,6 +807,31 @@ class InternalAddressError(Exception):
     pass
 
 
+class MissingEventsTableError(Exception):
+    pass
+
+
+def _connect_customer_clickhouse(connection: WarehouseConnection) -> Client:
+    """Build a client for the customer's ClickHouse instance from stored
+    connection details, re-checking the host against internal address ranges
+    first: DNS may resolve differently than it did at validation time, and
+    rows may predate host validation."""
+    config = typing.cast(ClickHouseConfig, connection.config or {})
+    credentials = typing.cast(ClickHouseCredentials, connection.credentials or {})
+    if is_internal_address(config["host"], include_shared=True):
+        raise InternalAddressError(config["host"])
+    return Client(
+        config["host"],
+        port=config["port"],
+        user=config["username"],
+        password=credentials["password"],
+        database=config["database"],
+        secure=config["secure"],
+        connect_timeout=CLICKHOUSE_CONNECT_TIMEOUT_SECONDS,
+        send_receive_timeout=CLICKHOUSE_VERIFY_TIMEOUT_SECONDS,
+    )
+
+
 def _describe_verification_error(error: Exception) -> str:
     if isinstance(error, clickhouse_errors.ServerException):
         # 516 = AUTHENTICATION_FAILED (not in clickhouse_driver.errors.ErrorCodes)
@@ -813,6 +846,11 @@ def _describe_verification_error(error: Exception) -> str:
         return "Could not connect to the host."
     if isinstance(error, InternalAddressError):
         return "Host must not target internal or private network addresses."
+    if isinstance(error, MissingEventsTableError):
+        return (
+            "Events table not found in the configured database. "
+            "Run the setup SQL to create it."
+        )
     if isinstance(error, KeyError):
         return "Stored connection details are incomplete."
     return "Verification failed."
@@ -822,30 +860,19 @@ 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. With persist=False, the status is only
-    set on the in-memory instance, allowing unsaved connections to be tested."""
+    """Run SELECT 1 against the customer's ClickHouse, check the events table
+    exists, and set the status to 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)
-        config = typing.cast(ClickHouseConfig, connection.config or {})
-        credentials = typing.cast(ClickHouseCredentials, connection.credentials or {})
-        # Re-check right before connecting: DNS may resolve differently than it
-        # did at validation time, and rows may predate host validation.
-        if is_internal_address(config["host"], include_shared=True):
-            raise InternalAddressError(config["host"])
-        client = Client(
-            config["host"],
-            port=config["port"],
-            user=config["username"],
-            password=credentials["password"],
-            database=config["database"],
-            secure=config["secure"],
-            connect_timeout=CLICKHOUSE_CONNECT_TIMEOUT_SECONDS,
-            send_receive_timeout=CLICKHOUSE_VERIFY_TIMEOUT_SECONDS,
-        )
+        client = _connect_customer_clickhouse(connection)
         try:
             client.execute("SELECT 1")
+            rows = client.execute("EXISTS TABLE events")
+            if not rows[0][0]:
+                raise MissingEventsTableError()
         finally:
             client.disconnect()
     except Exception as error:
@@ -893,9 +920,16 @@ def annotate_warehouse_event_stats(
     connection: WarehouseConnection,
     environment_key: str,
 ) -> None:
-    """Attach live warehouse event stats to a flagsmith connection. No-op for
-    non-flagsmith connections or when no warehouse is configured; leaves stats
-    unset when the warehouse is unreachable. Read-only: never changes status."""
+    """Attach live warehouse event stats to a connection — from the managed
+    warehouse for flagsmith connections, from the customer's instance for
+    clickhouse ones. No-op for other types or when no warehouse is configured;
+    leaves stats unset when the warehouse is unreachable. Read-only: never
+    changes status."""
+    if connection.warehouse_type == WarehouseType.CLICKHOUSE:
+        stats = _get_customer_warehouse_event_stats_cached(connection, environment_key)
+        if stats is not None:
+            connection.event_stats = stats
+        return
     if (
         connection.warehouse_type != WarehouseType.FLAGSMITH
         or not settings.EXPERIMENTATION_CLICKHOUSE_URL
@@ -905,3 +939,39 @@ def annotate_warehouse_event_stats(
         connection.event_stats = get_warehouse_event_stats(environment_key)
     except Exception:
         return
+
+
+def _get_customer_warehouse_event_stats_cached(
+    connection: WarehouseConnection,
+    environment_key: str,
+) -> WarehouseEventStats | None:
+    """Return event counts recorded for `environment_key` in the customer's
+    ClickHouse instance, or None when it's unreachable. Results — including
+    failures — are cached briefly so read endpoints don't open a connection to
+    the customer's host on every request."""
+    cache_key = f"experimentation:customer_event_stats:{connection.id}"
+    cached = cache.get(cache_key)
+    if isinstance(cached, WarehouseEventStats):
+        return cached
+    if cached == _CUSTOMER_EVENT_STATS_UNAVAILABLE:
+        return None
+    try:
+        client = _connect_customer_clickhouse(connection)
+        try:
+            stats = _query_event_stats(client, environment_key)
+        finally:
+            client.disconnect()
+    except Exception:
+        cache.set(
+            cache_key,
+            _CUSTOMER_EVENT_STATS_UNAVAILABLE,
+            CUSTOMER_EVENT_STATS_CACHE_SECONDS,
+        )
+        logger.warning(
+            "connection.event_stats_failed",
+            environment__id=connection.environment_id,
+            exc_info=True,
+        )
+        return None
+    cache.set(cache_key, stats, CUSTOMER_EVENT_STATS_CACHE_SECONDS)
+    return stats

api/experimentation/types.py

@@ -57,7 +57,7 @@ class ClickHouseConfig(TypedDict):
 CLICKHOUSE_DEFAULTS: ClickHouseConfig = {
     "host": "",
     "port": 9440,
-    "database": "flagsmith",
+    "database": "flagsmith_exp",
     "username": "default",
     "secure": True,
 }

api/experimentation/views.py

@@ -106,6 +106,9 @@ def get_throttles(self) -> list[BaseThrottle]:
         ):
             self.throttle_scope = "warehouse_connection_write"
             return [*super().get_throttles(), ScopedRateThrottle()]
+        if self.action in ("list", "retrieve"):
+            self.throttle_scope = "warehouse_connection_read"
+            return [*super().get_throttles(), ScopedRateThrottle()]
         return super().get_throttles()
 
     def perform_create(self, serializer: BaseSerializer[WarehouseConnection]) -> None:

api/tests/unit/experimentation/test_services.py

@@ -40,7 +40,9 @@
 from experimentation.results_query import _MetricSlot
 from experimentation.services import (
     InternalAddressError,
+    MissingEventsTableError,
     _describe_verification_error,
+    annotate_warehouse_event_stats,
     verify_clickhouse_connection,
 )
 from experimentation.stats import VariantStats
@@ -2031,6 +2033,7 @@ def test_verify_clickhouse_connection__reachable__sets_connected(
 ) -> None:
     # Given
     mock_client = mocker.patch("experimentation.services.Client")
+    mock_client.return_value.execute.return_value = [(1,)]
     success_count_before = _verification_count("success")
     clickhouse_connection.status_detail = "stale detail"
     clickhouse_connection.save()
@@ -2052,7 +2055,10 @@ def test_verify_clickhouse_connection__reachable__sets_connected(
         connect_timeout=5,
         send_receive_timeout=5,
     )
-    mock_client.return_value.execute.assert_called_once_with("SELECT 1")
+    assert mock_client.return_value.execute.call_args_list == [
+        mocker.call("SELECT 1"),
+        mocker.call("EXISTS TABLE events"),
+    ]
     mock_client.return_value.disconnect.assert_called_once_with()
     assert _verification_count("success") == success_count_before + 1
     assert {
@@ -2071,14 +2077,20 @@ def test_verify_clickhouse_connection__reachable__sets_connected(
             Exception("connection refused"),
             "Verification failed.",
         ),
+        (
+            {"password": "hunter2"},
+            [[(1,)], [(0,)]],
+            "Events table not found in the configured database. "
+            "Run the setup SQL to create it.",
+        ),
         (None, None, "Stored connection details are incomplete."),
     ],
-    ids=["driver_error", "missing_credentials"],
+    ids=["driver_error", "missing_events_table", "missing_credentials"],
 )
 def test_verify_clickhouse_connection__failure__sets_errored_with_detail(
     clickhouse_connection: WarehouseConnection,
     credentials: dict[str, str] | None,
-    execute_side_effect: Exception | None,
+    execute_side_effect: Exception | list[list[tuple[int]]] | None,
     expected_detail: str,
     log: StructuredLogCapture,
     mocker: MockerFixture,
@@ -2161,6 +2173,11 @@ def test_verify_clickhouse_connection__internal_host__sets_errored_without_conne
             InternalAddressError("10.0.0.5"),
             "Host must not target internal or private network addresses.",
         ),
+        (
+            MissingEventsTableError(),
+            "Events table not found in the configured database. "
+            "Run the setup SQL to create it.",
+        ),
         (
             KeyError("host"),
             "Stored connection details are incomplete.",
@@ -2178,6 +2195,7 @@ def test_verify_clickhouse_connection__internal_host__sets_errored_without_conne
         "builtin_timeout_error",
         "network_error",
         "internal_address_error",
+        "missing_events_table_error",
         "key_error",
         "generic_exception",
     ],
@@ -2191,3 +2209,50 @@ def test_describe_verification_error__known_error_types__returns_expected_detail
 
     # Then
     assert detail == expected_detail
+
+
+@pytest.mark.parametrize(
+    "execute_side_effect, expected_stats",
+    [
+        (
+            [[(42, 7)]],
+            WarehouseEventStats(total_events_received=42, unique_events_count=7),
+        ),
+        (Exception("connection refused"), None),
+    ],
+    ids=["reachable", "unreachable"],
+)
+def test_annotate_warehouse_event_stats__clickhouse_connection__queries_customer_instance(
+    clickhouse_connection: WarehouseConnection,
+    reset_cache: None,
+    execute_side_effect: Exception | list[list[tuple[int, int]]],
+    expected_stats: WarehouseEventStats | None,
+    log: StructuredLogCapture,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mock_client = mocker.patch("experimentation.services.Client")
+    mock_client.return_value.execute.side_effect = execute_side_effect
+
+    # When
+    annotate_warehouse_event_stats(clickhouse_connection, "test-env-key")
+
+    # Then
+    assert getattr(clickhouse_connection, "event_stats", None) == expected_stats
+    mock_client.return_value.execute.assert_called_once_with(
+        "SELECT count() AS total, uniqExact(event) AS unique "
+        "FROM events WHERE environment_key = %(environment_key)s",
+        {"environment_key": "test-env-key"},
+    )
+    mock_client.return_value.disconnect.assert_called_once_with()
+    assert any(
+        event["event"] == "connection.event_stats_failed" for event in log.events
+    ) == (expected_stats is None)
+
+    # When — the outcome is cached, so a second request doesn't reconnect
+    fresh_connection = WarehouseConnection.objects.get(id=clickhouse_connection.id)
+    annotate_warehouse_event_stats(fresh_connection, "test-env-key")
+
+    # Then
+    mock_client.assert_called_once()
+    assert getattr(fresh_connection, "event_stats", None) == expected_stats