feat(experimentation): publish warehouse connections to Redis and apply delivery status

api/core/fields.py

@@ -1,17 +1,9 @@
-import base64
-import hashlib
-import json
-from typing import Any, TypeVar
+from typing import TypeVar
 
-import structlog
-from cryptography.fernet import Fernet, InvalidToken
-from django.conf import settings
 from django.db import models
 
 from core.validators import validate_http_url_scheme, validate_no_internal_address
 
-logger = structlog.get_logger("core")
-
 _ST = TypeVar("_ST")
 _GT = TypeVar("_GT")
 
@@ -34,38 +26,3 @@ class NoSSRFURLField(models.URLField[_ST, _GT]):
         validate_http_url_scheme,
         validate_no_internal_address,
     ]
-
-
-def _get_fernet() -> Fernet:
-    secret: str = settings.WAREHOUSE_CREDENTIALS_SECRET
-    digest = hashlib.sha256(secret.encode()).digest()
-    return Fernet(base64.urlsafe_b64encode(digest))
-
-
-class EncryptedJSONField(models.TextField[Any, Any]):
-    def get_prep_value(self, value: Any) -> str | None:
-        if value is None:
-            return None
-        return _get_fernet().encrypt(json.dumps(value).encode()).decode()
-
-    def from_db_value(
-        self,
-        value: str | None,
-        expression: object,
-        connection: object,
-    ) -> Any:
-        if value is None:
-            return None
-        try:
-            plaintext = _get_fernet().decrypt(value.encode())
-        except InvalidToken:
-            logger.warning("encrypted_field.decrypt_failed", exc_info=True)
-            return None
-        return json.loads(plaintext)
-
-    def get_lookup(self, lookup_name: str) -> Any:
-        if lookup_name != "isnull":
-            raise NotImplementedError(
-                "EncryptedJSONField only supports isnull lookups."
-            )
-        return super().get_lookup(lookup_name)

api/experimentation/dataclasses.py

@@ -125,3 +125,10 @@ class ResultsSummary:
     # totals divide: a bucket's own conversions over its own new identities
     # compares different people and can exceed 100%.
     exposures_timeseries: ExposuresTimeseries
+
+
+@dataclass(frozen=True)
+class WarehouseDeliveryStatus:
+    connection_id: int
+    status: str
+    detail: str | None

api/experimentation/fields.py

@@ -0,0 +1,45 @@
+from typing import Any
+
+import structlog
+from cryptography.fernet import InvalidToken
+from django.db import models
+
+from experimentation.warehouse_credentials import (
+    decrypt_warehouse_credentials,
+    encrypt_warehouse_credentials,
+)
+
+logger = structlog.get_logger("experimentation")
+
+
+class EncryptedJSONField(models.TextField[Any, Any]):
+    """Stores a JSON value as Fernet ciphertext keyed on
+    WAREHOUSE_CREDENTIALS_SECRET. Used for a warehouse connection's
+    credentials, so the same ciphertext can be handed to the warehouse-delivery
+    service."""
+
+    def get_prep_value(self, value: Any) -> str | None:
+        if value is None:
+            return None
+        return encrypt_warehouse_credentials(value)
+
+    def from_db_value(
+        self,
+        value: str | None,
+        expression: object,
+        connection: object,
+    ) -> Any:
+        if value is None:
+            return None
+        try:
+            return decrypt_warehouse_credentials(value)
+        except InvalidToken:
+            logger.warning("encrypted_field.decrypt_failed", exc_info=True)
+            return None
+
+    def get_lookup(self, lookup_name: str) -> Any:
+        if lookup_name != "isnull":
+            raise NotImplementedError(
+                "EncryptedJSONField only supports isnull lookups."
+            )
+        return super().get_lookup(lookup_name)

api/experimentation/ingestion_redis.py

@@ -0,0 +1,17 @@
+from functools import lru_cache
+
+from django.conf import settings
+from redis.cluster import RedisCluster
+
+SOCKET_TIMEOUT = 1
+
+
+@lru_cache(maxsize=1)
+def get_client() -> RedisCluster:
+    """The Redis the ingestion server and the warehouse-delivery service read
+    their configuration from and, for the latter, write outcomes to."""
+    return RedisCluster.from_url(  # type: ignore[no-any-return]
+        settings.INGESTION_REDIS_URL,
+        socket_timeout=SOCKET_TIMEOUT,
+        socket_keepalive=True,
+    )

api/experimentation/ingestion_sync_service.py

@@ -1,10 +1,8 @@
 from __future__ import annotations
 
-from functools import lru_cache
 from typing import TYPE_CHECKING
 
-from django.conf import settings
-from redis.cluster import RedisCluster
+from experimentation.ingestion_redis import get_client
 
 if TYPE_CHECKING:
     from datetime import datetime
@@ -13,26 +11,14 @@
 INGESTION_ENVIRONMENT_DESTINATION_PREFIX = "experimentation:environment_destinations:"
 
 
-SOCKET_TIMEOUT = 1
-
-
-@lru_cache(maxsize=1)
-def _get_client() -> RedisCluster:
-    return RedisCluster.from_url(  # type: ignore[no-any-return]
-        settings.INGESTION_REDIS_URL,
-        socket_timeout=SOCKET_TIMEOUT,
-        socket_keepalive=True,
-    )
-
-
 def set_ingestion_key(
     key: str,
     *,
     environment_key: str,
     expires_at: datetime | None = None,
 ) -> None:
     redis_key = f"{INGESTION_ENVIRONMENT_KEY_PREFIX}{key}"
-    _get_client().set(
+    get_client().set(
         redis_key,
         environment_key,
         exat=int(expires_at.timestamp()) if expires_at is not None else None,
@@ -41,14 +27,14 @@ def set_ingestion_key(
 
 def delete_ingestion_key(key: str) -> None:
     redis_key = f"{INGESTION_ENVIRONMENT_KEY_PREFIX}{key}"
-    _get_client().delete(redis_key)
+    get_client().delete(redis_key)
 
 
 def set_ingestion_destination(client_api_key: str, *, topic: str) -> None:
     redis_key = f"{INGESTION_ENVIRONMENT_DESTINATION_PREFIX}{client_api_key}"
-    _get_client().set(redis_key, topic)
+    get_client().set(redis_key, topic)
 
 
 def delete_ingestion_destination(client_api_key: str) -> None:
     redis_key = f"{INGESTION_ENVIRONMENT_DESTINATION_PREFIX}{client_api_key}"
-    _get_client().delete(redis_key)
+    get_client().delete(redis_key)

api/experimentation/migrations/0010_warehouse_connection_credentials_and_status_detail.py

@@ -1,6 +1,6 @@
 from django.db import migrations, models
 
-import core.fields
+import experimentation.fields
 
 
 class Migration(migrations.Migration):
@@ -12,7 +12,7 @@ class Migration(migrations.Migration):
         migrations.AddField(
             model_name="warehouseconnection",
             name="credentials",
-            field=core.fields.EncryptedJSONField(blank=True, null=True),
+            field=experimentation.fields.EncryptedJSONField(blank=True, null=True),
         ),
         migrations.AddField(
             model_name="warehouseconnection",

api/experimentation/models.py

@@ -13,14 +13,14 @@
     hook,
 )
 
-from core.fields import EncryptedJSONField
 from core.models import SoftDeleteExportableModel
 from environments.models import Environment
 from experimentation.dataclasses import (
     ExposuresSummary,
     ResultsSummary,
     WarehouseEventStats,
 )
+from experimentation.fields import EncryptedJSONField
 from experimentation.types import MetricDefinition
 
 # A computation's payload is the serialised form of its summary dataclass; the
@@ -78,7 +78,11 @@ class Meta:
         ]
 
     @hook(AFTER_CREATE)  # type: ignore[misc]
-    @hook(AFTER_UPDATE, when="warehouse_type", has_changed=True)  # type: ignore[misc]
+    @hook(  # type: ignore[misc]
+        AFTER_UPDATE,
+        when_any=["warehouse_type", "config", "credentials"],
+        has_changed=True,
+    )
     @hook(AFTER_DELETE)  # type: ignore[misc]
     def sync_to_ingestion(self) -> None:
         from experimentation.tasks import sync_environment_ingestion

api/experimentation/services.py

@@ -22,7 +22,10 @@
 from cohorts.models import Cohort
 from core.dataclasses import AuthorData
 from environments.tasks import rebuild_environment_document
-from experimentation import warehouse_delivery_service
+from experimentation import (
+    warehouse_delivery_sync_service,
+    warehouse_verification_service,
+)
 from experimentation.constants import (
     CONTROL_VARIANT_KEY,
     EXPERIMENT_FLAG,
@@ -1413,15 +1416,15 @@ def verify_clickhouse_connection(
     log = logger.bind(environment__id=connection.environment_id)
     try:
         log = log.bind(organisation__id=connection.environment.project.organisation_id)
-        with warehouse_delivery_service.delivery_client(
+        with warehouse_verification_service.delivery_client(
             connection,
             send_receive_timeout=CLICKHOUSE_VERIFY_TIMEOUT_SECONDS,
         ) as client:
-            warehouse_delivery_service.check_events_table_exists(client)
+            warehouse_verification_service.check_events_table_exists(client)
     except Exception as error:
         connection.status = WarehouseConnectionStatus.ERRORED
-        connection.status_detail = warehouse_delivery_service.describe_warehouse_error(
-            error
+        connection.status_detail = (
+            warehouse_verification_service.describe_warehouse_error(error)
         )
         if persist:
             connection.save(update_fields=["status", "status_detail"])
@@ -1461,6 +1464,29 @@ def refresh_warehouse_connection_status(
     return connection
 
 
+def annotate_warehouse_delivery_statuses(
+    connections: Sequence[WarehouseConnection],
+) -> None:
+    """For external connections that passed verification, show what the
+    warehouse-delivery service last saw. Read-only: nothing is saved."""
+    verified = [
+        connection
+        for connection in connections
+        if connection.warehouse_type != WarehouseType.FLAGSMITH
+        and connection.status == WarehouseConnectionStatus.CONNECTED
+    ]
+    if not verified:
+        return
+    statuses = warehouse_delivery_sync_service.get_warehouse_delivery_statuses(
+        [connection.id for connection in verified]
+    )
+    for connection in verified:
+        outcome = statuses.get(connection.id)
+        if outcome is not None and outcome.status == WarehouseConnectionStatus.ERRORED:
+            connection.status = WarehouseConnectionStatus.ERRORED
+            connection.status_detail = outcome.detail
+
+
 def annotate_warehouse_event_stats(
     connection: WarehouseConnection,
     environment_key: str,
@@ -1501,7 +1527,7 @@ def _get_customer_warehouse_event_stats_cached(
     if cached == _CUSTOMER_EVENT_UNAVAILABLE:
         return None
     try:
-        with warehouse_delivery_service.delivery_client(
+        with warehouse_verification_service.delivery_client(
             connection,
             send_receive_timeout=CLICKHOUSE_VERIFY_TIMEOUT_SECONDS,
         ) as client:
@@ -1539,7 +1565,7 @@ def _get_customer_clickhouse_event_names(
     if cached == _CUSTOMER_EVENT_UNAVAILABLE:
         return None
     try:
-        with warehouse_delivery_service.delivery_client(
+        with warehouse_verification_service.delivery_client(
             connection,
             send_receive_timeout=CLICKHOUSE_EVENT_NAMES_TIMEOUT_SECONDS,
         ) as client:

api/experimentation/tasks.py

@@ -6,12 +6,13 @@
 from task_processor.exceptions import TaskBackoffError
 
 from environments.models import Environment, EnvironmentAPIKey
-from experimentation import ingestion_sync_service
+from experimentation import ingestion_sync_service, warehouse_delivery_sync_service
 from experimentation.constants import EXTERNAL_WAREHOUSE_EVENTS_TOPIC
 from experimentation.models import (
     Experiment,
     ExperimentExposures,
     ExperimentResults,
+    WarehouseConnection,
     WarehouseType,
 )
 from experimentation.services import (
@@ -47,14 +48,40 @@ def sync_environment_ingestion(environment_id: int) -> None:
         for api_key in environment.api_keys.all():
             ingestion_sync_service.delete_ingestion_key(api_key.key)
         ingestion_sync_service.delete_ingestion_destination(environment.api_key)
+        # The connection is gone, so include deleted ones to find its id.
+        connection_ids = (
+            WarehouseConnection.objects.all_with_deleted()
+            .filter(environment_id=environment.id)
+            .values_list("id", flat=True)
+        )
+        warehouse_delivery_sync_service.remove_warehouse_connection(environment.api_key)
+        warehouse_delivery_sync_service.delete_warehouse_delivery_statuses(
+            list(connection_ids)
+        )
         return
 
-    # Destination first, then keys. As soon as a key is in Redis the ingestion
-    # server accepts events for it, and if no destination is stored yet those
-    # events go to Flagsmith's own topic instead of the external warehouse one.
+    # Connection details, then destination, then keys. Each step makes the next
+    # one safe: the warehouse-delivery service drops events for an environment
+    # whose connection it cannot find in Redis, and the ingestion server sends
+    # events for an environment with no destination to Flagsmith's own topic.
     if connection.warehouse_type == WarehouseType.FLAGSMITH:
         ingestion_sync_service.delete_ingestion_destination(environment.api_key)
+        warehouse_delivery_sync_service.remove_warehouse_connection(environment.api_key)
+        warehouse_delivery_sync_service.delete_warehouse_delivery_statuses(
+            [connection.id]
+        )
     else:
+        warehouse_delivery_sync_service.publish_warehouse_connection(
+            environment.api_key,
+            connection_id=connection.id,
+            warehouse_type=connection.warehouse_type,
+            config=connection.config or {},
+            credentials=connection.credentials,
+        )
+        # Its last outcome was about the details just replaced.
+        warehouse_delivery_sync_service.delete_warehouse_delivery_statuses(
+            [connection.id]
+        )
         ingestion_sync_service.set_ingestion_destination(
             environment.api_key,
             topic=EXTERNAL_WAREHOUSE_EVENTS_TOPIC,

api/experimentation/views.py

@@ -61,6 +61,7 @@
 )
 from experimentation.services import (
     EVENT_NAMES_SUPPORTED_WAREHOUSE_TYPES,
+    annotate_warehouse_delivery_statuses,
     annotate_warehouse_event_stats,
     apply_experiment_rollout,
     create_experiment_audit_log,
@@ -150,12 +151,14 @@ def list(self, request: Request, *args: object, **kwargs: object) -> Response:
         if not exclude_event_stats:
             for connection in connections:
                 annotate_warehouse_event_stats(connection, environment_api_key)
+        annotate_warehouse_delivery_statuses(connections)
         serializer = self.get_serializer(connections, many=True)
         return Response(serializer.data)
 
     def retrieve(self, request: Request, *args: object, **kwargs: object) -> Response:
         connection = self.get_object()
         annotate_warehouse_event_stats(connection, self.kwargs["environment_api_key"])
+        annotate_warehouse_delivery_statuses([connection])
         serializer = self.get_serializer(connection)
         return Response(serializer.data)