From 3060b29d3e9cc523e38035ad063300650ad1b8bb Mon Sep 17 00:00:00 2001
From: Daesgar <daniel.e@posthog.com>
Date: Fri, 22 Nov 2024 16:52:35 +0100
Subject: [PATCH] chore: use posthog single-shard cluster for
 distributed_events_recent (#26339)

---
 docker/clickhouse/config.xml                         |  8 ++++++++
 .../migrations/0087_events_recent_single_shard.py    | 12 ++++++++++++
 posthog/clickhouse/table_engines.py                  |  7 +++++--
 posthog/models/event/sql.py                          |  6 +++++-
 posthog/settings/data_stores.py                      |  1 +
 5 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 posthog/clickhouse/migrations/0087_events_recent_single_shard.py

diff --git a/docker/clickhouse/config.xml b/docker/clickhouse/config.xml
index bc1c691f8b4..1eca858c32d 100644
--- a/docker/clickhouse/config.xml
+++ b/docker/clickhouse/config.xml
@@ -712,6 +712,14 @@
                 </replica>
             </shard>
         </posthog>
+        <posthog_single_shard>
+            <shard>
+                <replica>
+                    <host>localhost</host>
+                    <port>9000</port>
+                </replica>
+            </shard>
+        </posthog_single_shard>
     </remote_servers>
 
     <!-- The list of hosts allowed to use in URL-related storage engines and table functions.
diff --git a/posthog/clickhouse/migrations/0087_events_recent_single_shard.py b/posthog/clickhouse/migrations/0087_events_recent_single_shard.py
new file mode 100644
index 00000000000..6d5b445d6a5
--- /dev/null
+++ b/posthog/clickhouse/migrations/0087_events_recent_single_shard.py
@@ -0,0 +1,12 @@
+from posthog.clickhouse.client.migration_tools import run_sql_with_exceptions
+from posthog.models.event.sql import (
+    DISTRIBUTED_EVENTS_RECENT_TABLE_SQL,
+)
+
+from posthog.settings import CLICKHOUSE_CLUSTER
+
+
+operations = [
+    run_sql_with_exceptions(f"DROP TABLE IF EXISTS distributed_events_recent ON CLUSTER '{CLICKHOUSE_CLUSTER}'"),
+    run_sql_with_exceptions(DISTRIBUTED_EVENTS_RECENT_TABLE_SQL()),
+]
diff --git a/posthog/clickhouse/table_engines.py b/posthog/clickhouse/table_engines.py
index 9f959216fbd..109534c292c 100644
--- a/posthog/clickhouse/table_engines.py
+++ b/posthog/clickhouse/table_engines.py
@@ -75,9 +75,12 @@ class AggregatingMergeTree(MergeTreeEngine):
 
 
 class Distributed:
-    def __init__(self, data_table: str, sharding_key: str):
+    def __init__(self, data_table: str, sharding_key: str, cluster: str = settings.CLICKHOUSE_CLUSTER):
         self.data_table = data_table
         self.sharding_key = sharding_key
+        self.cluster = cluster
 
     def __str__(self):
-        return f"Distributed('{settings.CLICKHOUSE_CLUSTER}', '{settings.CLICKHOUSE_DATABASE}', '{self.data_table}', {self.sharding_key})"
+        return (
+            f"Distributed('{self.cluster}', '{settings.CLICKHOUSE_DATABASE}', '{self.data_table}', {self.sharding_key})"
+        )
diff --git a/posthog/models/event/sql.py b/posthog/models/event/sql.py
index 67f5ac97325..e6ea42067a6 100644
--- a/posthog/models/event/sql.py
+++ b/posthog/models/event/sql.py
@@ -261,7 +261,11 @@ TTL _timestamp + INTERVAL 7 DAY
 DISTRIBUTED_EVENTS_RECENT_TABLE_SQL = lambda: EVENTS_TABLE_BASE_SQL.format(
     table_name="distributed_events_recent",
     cluster=settings.CLICKHOUSE_CLUSTER,
-    engine=Distributed(data_table=EVENTS_RECENT_DATA_TABLE(), sharding_key="sipHash64(distinct_id)"),
+    engine=Distributed(
+        data_table=EVENTS_RECENT_DATA_TABLE(),
+        sharding_key="sipHash64(distinct_id)",
+        cluster=settings.CLICKHOUSE_SINGLE_SHARD_CLUSTER,
+    ),
     extra_fields=KAFKA_COLUMNS_WITH_PARTITION + INSERTED_AT_COLUMN + f", {KAFKA_TIMESTAMP_MS_COLUMN}",
     materialized_columns="",
     indexes="",
diff --git a/posthog/settings/data_stores.py b/posthog/settings/data_stores.py
index e64f38877f8..31eea8e08b7 100644
--- a/posthog/settings/data_stores.py
+++ b/posthog/settings/data_stores.py
@@ -154,6 +154,7 @@ CLICKHOUSE_CA: str | None = os.getenv("CLICKHOUSE_CA", None)
 CLICKHOUSE_SECURE: bool = get_from_env("CLICKHOUSE_SECURE", not TEST and not DEBUG, type_cast=str_to_bool)
 CLICKHOUSE_VERIFY: bool = get_from_env("CLICKHOUSE_VERIFY", True, type_cast=str_to_bool)
 CLICKHOUSE_ENABLE_STORAGE_POLICY: bool = get_from_env("CLICKHOUSE_ENABLE_STORAGE_POLICY", False, type_cast=str_to_bool)
+CLICKHOUSE_SINGLE_SHARD_CLUSTER: str = os.getenv("CLICKHOUSE_SINGLE_SHARD_CLUSTER", "posthog_single_shard")
 
 CLICKHOUSE_CONN_POOL_MIN: int = get_from_env("CLICKHOUSE_CONN_POOL_MIN", 20, type_cast=int)
 CLICKHOUSE_CONN_POOL_MAX: int = get_from_env("CLICKHOUSE_CONN_POOL_MAX", 1000, type_cast=int)