feat(Segment membership inspection PoC): Daily ClickHouse-backed per-env segment counts (#7464)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: khvn26

.github/workflows/update-flagsmith-environment.yml

@@ -9,6 +9,9 @@ defaults:
   run:
     working-directory: api
 
+permissions:
+  contents: read
+
 jobs:
   update_server_defaults:
     runs-on: depot-ubuntu-latest

api/app/routers.py

@@ -3,6 +3,7 @@
 from django.db.models import Model
 
 AnalyticsDatabaseName = Literal["analytics"]
+ClickHouseDatabaseName = Literal["clickhouse"]
 
 
 class AnalyticsRouter:
@@ -38,3 +39,38 @@ def allow_migrate(self, db: str, app_label: str, **hints: Any) -> bool | None:
         if db == "analytics":
             return app_label in self.route_app_labels
         return None
+
+
+class ClickHouseRouter:
+    route_app_labels = ["clickhouse"]
+
+    def db_for_read(
+        self, model: type[Model], **hints: Any
+    ) -> ClickHouseDatabaseName | None:
+        if model._meta.app_label in self.route_app_labels:
+            return "clickhouse"
+        return None
+
+    def db_for_write(
+        self, model: type[Model], **hints: Any
+    ) -> ClickHouseDatabaseName | None:
+        if model._meta.app_label in self.route_app_labels:
+            return "clickhouse"
+        return None
+
+    def allow_relation(self, obj1: Model, obj2: Model, **hints: Any) -> bool | None:
+        # ClickHouse has no FKs and we don't expose CH-app models, so any
+        # relation involving this app is forbidden.
+        if (
+            obj1._meta.app_label in self.route_app_labels
+            or obj2._meta.app_label in self.route_app_labels
+        ):
+            return False
+        return None
+
+    def allow_migrate(self, db: str, app_label: str, **hints: Any) -> bool | None:
+        if db == "clickhouse":
+            return app_label in self.route_app_labels
+        if app_label in self.route_app_labels:
+            return False
+        return None

api/app/settings/common.py

@@ -118,6 +118,8 @@
     "features.workflows.core",
     "features.release_pipelines.core",
     "segments",
+    "segment_membership",
+    "clickhouse",
     "app",
     "e2etests",
     "simple_history",
@@ -1446,3 +1448,39 @@
 PYLON_IDENTITY_VERIFICATION_SECRET = env.str("PYLON_IDENTITY_VERIFICATION_SECRET", None)
 
 OSIC_UPDATE_BATCH_SIZE = env.int("OSIC_UPDATE_BATCH_SIZE", default=500)
+
+# ClickHouse backs the segment_membership backfill and refresh tasks. Set
+# CLICKHOUSE_URL (DSN form) or any CLICKHOUSE_HOST + discrete fields to enable.
+# Discrete settings override the matching field parsed from the URL.
+CLICKHOUSE_URL = env.str("CLICKHOUSE_URL", default=None)
+CLICKHOUSE_HOST = env.str("CLICKHOUSE_HOST", default=None)
+CLICKHOUSE_PORT = env.int("CLICKHOUSE_PORT", default=None)
+CLICKHOUSE_USER = env.str("CLICKHOUSE_USER", default=None)
+CLICKHOUSE_PASSWORD = env.str("CLICKHOUSE_PASSWORD", default=None)
+CLICKHOUSE_DATABASE = env.str("CLICKHOUSE_DATABASE", default=None)
+CLICKHOUSE_SECURE = env.bool("CLICKHOUSE_SECURE", default=None)
+
+CLICKHOUSE_ENABLED = bool(CLICKHOUSE_URL or CLICKHOUSE_HOST)
+
+# Always installed: the router fences the `clickhouse` app's migrations off
+# the default Postgres database whether or not a CH alias is configured.
+DATABASE_ROUTERS.append("app.routers.ClickHouseRouter")
+
+if CLICKHOUSE_ENABLED:
+    _clickhouse_db: dict[str, Any] = {
+        "ENGINE": "clickhouse_backend.backend",
+        "HOST": CLICKHOUSE_HOST,
+        "PORT": CLICKHOUSE_PORT,
+        "USER": CLICKHOUSE_USER,
+        "PASSWORD": CLICKHOUSE_PASSWORD,
+        "NAME": CLICKHOUSE_DATABASE,
+        "OPTIONS": {
+            "dsn": CLICKHOUSE_URL,
+            "secure": CLICKHOUSE_SECURE,
+            "settings": {
+                # ClickHouse Cloud 25.12 requires this for `JSON`-column DDL.
+                "allow_experimental_json_type": 1,
+            },
+        },
+    }
+    DATABASES["clickhouse"] = _clickhouse_db  # type: ignore[assignment]

api/clickhouse/__init__.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class ClickHouseConfig(AppConfig):
+    name = "clickhouse"
+    label = "clickhouse"

api/clickhouse/apps.py

@@ -0,0 +1,28 @@
+from django.db import migrations
+
+
+_SCHEMA_DDL = """\
+CREATE TABLE IF NOT EXISTS IDENTITIES (
+    environment_id String,
+    -- (environment_id, identifier) is the natural unique key in Flagsmith's
+    -- identity model — dedupes ReplacingMergeTree without a synthetic id.
+    identifier String,
+    identity_key String,
+    -- Stored per top-level key as typed subcolumns; SQL NULL for empty traits.
+    traits JSON,
+    -- ReplacingMergeTree version column; most-recent insert wins per PK.
+    inserted_at DateTime DEFAULT now()
+)
+ENGINE = ReplacingMergeTree(inserted_at)
+ORDER BY (environment_id, identifier)
+"""
+
+
+class Migration(migrations.Migration):
+    # ClickHouse has no transactional DDL.
+    atomic = False
+    initial = True
+
+    operations = [
+        migrations.RunSQL(_SCHEMA_DDL, reverse_sql="DROP TABLE IF EXISTS IDENTITIES"),
+    ]

api/clickhouse/migrations/0001_create_identities.py

@@ -44,6 +44,8 @@ dependencies = [
     "drf-writable-nested>=0.6.2,<0.7.0",
     "django-filter>=2.4.0,<2.5.0",
     "flagsmith-flag-engine>=10.1.0,<11.0.0",
+    "flagsmith-sql-flag-engine>=0.1.0,<0.2.0",
+    "django-clickhouse-backend>=1.4,<2.0",
     "boto3>=1.35.95,<1.36.0",
     "slack-sdk>=3.9.0,<3.10.0",
     "asgiref>=3.8.1,<3.9.0",
@@ -71,7 +73,7 @@ dependencies = [
     "hubspot-api-client>=12.0.0,<13.0.0",
     "djangorestframework-dataclasses>=1.3.1,<2.0.0",
     "pyotp>=2.9.0,<3.0.0",
-    "flagsmith-common[common-core,flagsmith-schemas,task-processor]>=3.9.0,<4",
+    "flagsmith-common[common-core,flagsmith-schemas,task-processor]>=3.9.1,<4",
     "django-stubs>=5.1.3,<6.0.0",
     "tzdata>=2024.1,<2025.0.0",
     "djangorestframework-simplejwt>=5.5.1,<6.0.0",
@@ -541,6 +543,10 @@ ignore_missing_imports = true
 module = ["openfeature_flagsmith.*"]
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = ["clickhouse_backend.*", "clickhouse_driver.*"]
+ignore_missing_imports = true
+
 [[tool.mypy.overrides]]
 module = ["scim.*"]
 ignore_missing_imports = true

api/clickhouse/migrations/__init__.py

@@ -45,6 +45,9 @@ run_task_processor() {
     if [ -n "$TASK_PROCESSOR_DATABASE_URL" ] || [ -n "$TASK_PROCESSOR_DATABASE_NAME" ]; then
         waitfordb --waitfor 30 --migrations --database task_processor
     fi
+    if [ -n "$CLICKHOUSE_URL" ] || [ -n "$CLICKHOUSE_HOST" ]; then
+        waitfordb --waitfor 30 --migrations --database clickhouse
+    fi
     exec flagsmith start \
              --bind 0.0.0.0:8000 \
              --access-logfile $ACCESS_LOG_LOCATION \
@@ -68,6 +71,12 @@ migrate_task_processor_db(){
     fi
     python manage.py migrate --database task_processor
 }
+migrate_clickhouse_db(){
+    if [ -z "$CLICKHOUSE_URL" ] && [ -z "$CLICKHOUSE_HOST" ]; then
+        return 0
+    fi
+    python manage.py migrate --database clickhouse
+}
 bootstrap(){
     python manage.py bootstrap
 }
@@ -81,17 +90,20 @@ if [ "$1" = "migrate" ]; then
     migrate
     migrate_analytics_db
     migrate_task_processor_db
+    migrate_clickhouse_db
 elif [ "$1" = "serve" ]; then
     serve
 elif [ "$1" = "run-task-processor" ]; then
     migrate
     migrate_analytics_db
     migrate_task_processor_db
+    migrate_clickhouse_db
     run_task_processor
 elif [ "$1" = "migrate-and-serve" ]; then
     migrate
     migrate_analytics_db
     migrate_task_processor_db
+    migrate_clickhouse_db
     bootstrap
     serve
 else