feat: push delivery of buffered events to external ClickHouse warehouses

api/experimentation/constants.py

@@ -14,6 +14,8 @@
 
 CONTROL_VARIANT_KEY = "control"
 
+DELIVERY_INTERVAL = timedelta(minutes=10)
+
 # Below these per-variant floors a metric shows "collecting data" rather than
 # inference; sample-ratio is only checked once there is enough traffic to judge.
 RESULTS_MIN_IDENTITIES_PER_VARIANT = 50

api/experimentation/metrics.py

@@ -8,3 +8,20 @@
         ["result"],
     )
 )
+
+flagsmith_experimentation_warehouse_delivery_runs_total = prometheus_client.Counter(
+    "flagsmith_experimentation_warehouse_delivery_runs_total",
+    "Outcomes of per-connection runs delivering buffered event objects to "
+    "customers' own data warehouses. `result` label is either `success` or "
+    "`failure`; a failed run delivers nothing and is retried on the next tick.",
+    ["result"],
+)
+
+flagsmith_experimentation_warehouse_delivery_objects_total = prometheus_client.Counter(
+    "flagsmith_experimentation_warehouse_delivery_objects_total",
+    "Buffered S3 event objects processed by delivery to customers' own data "
+    "warehouses. `result` label is `delivered` for objects inserted and "
+    "archived, or `rejected` for objects the warehouse refused, which are "
+    "moved aside and never retried.",
+    ["result"],
+)

api/experimentation/services.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import time
 import typing
 from dataclasses import replace
 from functools import lru_cache
 
 import structlog
+from clickhouse_connect.driver.exceptions import ClickHouseError
 from clickhouse_driver import Client
 from clickhouse_driver import errors as clickhouse_errors
 from clickhouse_driver.util.helpers import parse_url
@@ -21,6 +23,7 @@
 from core.dataclasses import AuthorData
 from core.network import is_internal_address
 from environments.tasks import rebuild_environment_document
+from experimentation import warehouse_delivery_service
 from experimentation.constants import (
     CONTROL_VARIANT_KEY,
     EXPERIMENT_FLAG,
@@ -45,6 +48,8 @@
 )
 from experimentation.metrics import (
     flagsmith_experimentation_warehouse_connection_verifications_total,
+    flagsmith_experimentation_warehouse_delivery_objects_total,
+    flagsmith_experimentation_warehouse_delivery_runs_total,
 )
 from experimentation.models import (
     VALID_STATUS_TRANSITIONS,
@@ -84,6 +89,8 @@
     from collections.abc import Sequence
     from datetime import datetime
 
+    from clickhouse_connect.driver.client import Client as ClickHouseHTTPClient
+
     from experimentation.models import Metric, WarehouseConnection
     from experimentation.types import ExposureGranularity
     from features.feature_states.models import FeatureValueType
@@ -100,6 +107,10 @@
 
 _CUSTOMER_EVENT_STATS_UNAVAILABLE = "unavailable"
 
+# A delivery run stops taking on new objects after this long, leaving room for
+# the slowest possible in-flight insert to still land inside the task timeout.
+DELIVERY_TIME_BUDGET_SECONDS = 210
+
 
 def is_warehouse_feature_enabled(organisation: Organisation) -> bool:
     return get_openfeature_client().get_boolean_value(
@@ -803,6 +814,159 @@ def mark_warehouse_pending_connection(
     return connection
 
 
+def mark_warehouse_delivery_failed(
+    connection: WarehouseConnection,
+    detail: str,
+) -> None:
+    connection.status = WarehouseConnectionStatus.ERRORED
+    connection.status_detail = detail[:255]
+    connection.save(update_fields=["status", "status_detail"])
+
+
+def mark_warehouse_delivery_succeeded(connection: WarehouseConnection) -> None:
+    if connection.status == WarehouseConnectionStatus.CONNECTED:
+        return
+
+    connection.status = WarehouseConnectionStatus.CONNECTED
+    connection.status_detail = None
+    connection.save(update_fields=["status", "status_detail"])
+
+
+def _deliver_pending_objects(
+    client: ClickHouseHTTPClient,
+    *,
+    bucket_name: str,
+    pending: list[str],
+    connection: WarehouseConnection,
+) -> tuple[int, int, int]:
+    log = logger.bind(
+        connection__id=connection.id,
+        environment__id=connection.environment_id,
+        organisation__id=connection.environment.project.organisation_id,
+    )
+    # A run that outlives the task timeout is retried while its own thread
+    # keeps delivering, so it must finish first: whatever is left is picked up
+    # on the next tick.
+    deadline = time.monotonic() + DELIVERY_TIME_BUDGET_SECONDS
+    delivered_count = rejected_count = rows_count = 0
+    for index, s3_key in enumerate(pending):
+        if time.monotonic() > deadline:
+            log.info(
+                "delivery.budget_exhausted",
+                objects__remaining_count=len(pending) - index,
+            )
+            break
+        try:
+            rows_count += warehouse_delivery_service.deliver_object(
+                client,
+                bucket_name,
+                s3_key,
+            )
+        except warehouse_delivery_service.ObjectRejectedError:
+            # This object's contents are the problem; the ones behind it are
+            # still deliverable.
+            warehouse_delivery_service.move_object(
+                bucket_name,
+                s3_key,
+                to_prefix=warehouse_delivery_service.FAILED_PREFIX,
+            )
+            rejected_count += 1
+            flagsmith_experimentation_warehouse_delivery_objects_total.labels(
+                result="rejected"
+            ).inc()
+            log.error(
+                "delivery.object_rejected",
+                s3__key=s3_key,
+                exc_info=True,
+            )
+            continue
+        warehouse_delivery_service.move_object(
+            bucket_name,
+            s3_key,
+            to_prefix=warehouse_delivery_service.ARCHIVE_PREFIX,
+        )
+        delivered_count += 1
+        flagsmith_experimentation_warehouse_delivery_objects_total.labels(
+            result="delivered"
+        ).inc()
+    return delivered_count, rejected_count, rows_count
+
+
+def deliver_warehouse_events(
+    connection: WarehouseConnection,
+    *,
+    bucket_name: str,
+) -> None:
+    """Deliver the environment's pending event objects to the connection's
+    warehouse, surfacing the outcome on the connection's status."""
+    log = logger.bind(
+        connection__id=connection.id,
+        environment__id=connection.environment_id,
+        organisation__id=connection.environment.project.organisation_id,
+    )
+    pending = warehouse_delivery_service.list_pending_objects(
+        bucket_name,
+        environment_key=connection.environment.api_key,
+    )
+    if not pending:
+        return
+
+    try:
+        with warehouse_delivery_service.delivery_client(connection) as client:
+            delivered_count, rejected_count, rows_count = _deliver_pending_objects(
+                client,
+                bucket_name=bucket_name,
+                pending=pending,
+                connection=connection,
+            )
+    except (warehouse_delivery_service.DeliveryConfigError, ClickHouseError) as exc:
+        # The warehouse itself is unusable; deliver nothing, leave every
+        # remaining object in place for the next run, and surface the
+        # breakage on the connection. Anything else — an S3 failure, a bug
+        # here — is ours, so it propagates and fails the task instead of
+        # blaming the customer's warehouse.
+        mark_warehouse_delivery_failed(
+            connection,
+            detail=warehouse_delivery_service.describe_delivery_error(exc),
+        )
+        flagsmith_experimentation_warehouse_delivery_runs_total.labels(
+            result="failure"
+        ).inc()
+        log.error("delivery.failed", exc_info=exc)
+        return
+
+    if delivered_count == 0 and rejected_count:
+        # Records all come from one ingestion pipeline, so every object being
+        # rejected points at the table's schema rather than the objects.
+        mark_warehouse_delivery_failed(
+            connection,
+            detail=(
+                f"The warehouse rejected every event object. Check that the "
+                f"`{warehouse_delivery_service.EVENTS_TABLE_NAME}` table "
+                f"matches the expected schema."
+            ),
+        )
+        flagsmith_experimentation_warehouse_delivery_runs_total.labels(
+            result="failure"
+        ).inc()
+        log.error(
+            "delivery.all_objects_rejected",
+            objects__rejected_count=rejected_count,
+        )
+        return
+
+    mark_warehouse_delivery_succeeded(connection)
+    flagsmith_experimentation_warehouse_delivery_runs_total.labels(
+        result="success"
+    ).inc()
+    log.info(
+        "delivery.completed",
+        objects__count=delivered_count,
+        objects__rejected_count=rejected_count,
+        rows__count=rows_count,
+    )
+
+
 class InternalAddressError(Exception):
     pass
 

api/experimentation/tasks.py

@@ -1,15 +1,31 @@
+from datetime import timedelta
+
 import structlog
 from django.utils import timezone
-from task_processor.decorators import register_task_handler
+from task_processor.decorators import (
+    register_recurring_task,
+    register_task_handler,
+)
 
 from environments.models import Environment, EnvironmentAPIKey
 from experimentation import ingestion_sync_service
-from experimentation.models import Experiment, ExperimentExposures, ExperimentResults
+from experimentation.constants import DELIVERY_INTERVAL
+from experimentation.models import (
+    Experiment,
+    ExperimentExposures,
+    ExperimentResults,
+    WarehouseConnection,
+    WarehouseType,
+)
 from experimentation.organisation_ingestion_service import (
     disable_ingestion_for_organisation,
     enable_ingestion_for_organisation,
 )
-from experimentation.services import compute_exposures_summary, compute_results_summary
+from experimentation.services import (
+    compute_exposures_summary,
+    compute_results_summary,
+    deliver_warehouse_events,
+)
 
 logger = structlog.get_logger("experimentation")
 
@@ -106,6 +122,38 @@ def remove_environment_ingestion_key(key: str) -> None:
     ingestion_sync_service.delete_ingestion_key(key)
 
 
+@register_recurring_task(run_every=DELIVERY_INTERVAL)
+def deliver_events_to_external_warehouses() -> None:
+    connection_ids = WarehouseConnection.objects.filter(
+        warehouse_type=WarehouseType.CLICKHOUSE,
+    ).values_list("id", flat=True)
+    for connection_id in connection_ids:
+        deliver_events_for_connection.delay(kwargs={"connection_id": connection_id})
+
+
+@register_task_handler(timeout=timedelta(minutes=9))
+def deliver_events_for_connection(connection_id: int) -> None:
+    connection = (
+        WarehouseConnection.objects.select_related(
+            "environment__project__organisation__ingestion_infrastructure",
+        )
+        .filter(id=connection_id)
+        .first()
+    )
+    if connection is None:
+        return
+
+    infrastructure = getattr(
+        connection.environment.project.organisation,
+        "ingestion_infrastructure",
+        None,
+    )
+    if infrastructure is None or not infrastructure.bucket_name:
+        return
+
+    deliver_warehouse_events(connection, bucket_name=infrastructure.bucket_name)
+
+
 @register_task_handler()
 def compute_experiment_exposures(experiment_id: int) -> None:
     experiment = (