Allow migrating on a clickhouse cluster (#4197)

* Allow migrating clickhouse on cluster

Context under #4090

* Add option to print sql

ee/management/commands/migrate_clickhouse.py

@@ -1,11 +1,17 @@
 import datetime
+from textwrap import indent
 
 from django.core.management.base import BaseCommand
 from infi.clickhouse_orm import Database  # type: ignore
 from infi.clickhouse_orm.migrations import MigrationHistory  # type: ignore
 from infi.clickhouse_orm.utils import import_submodules  # type: ignore
 
-from posthog.settings import CLICKHOUSE_DATABASE, CLICKHOUSE_HTTP_URL, CLICKHOUSE_PASSWORD, CLICKHOUSE_USER
+from posthog.settings import (
+    CLICKHOUSE_DATABASE,
+    CLICKHOUSE_MIGRATION_HOSTS_HTTP_URLS,
+    CLICKHOUSE_PASSWORD,
+    CLICKHOUSE_USER,
+)
 
 MIGRATIONS_PACKAGE_NAME = "ee.clickhouse.migrations"
 
@@ -21,21 +27,36 @@ def add_arguments(self, parser):
         parser.add_argument(
             "--plan", action="store_true", help="Shows a list of the migration actions that will be performed."
         )
+        parser.add_argument(
+            "--print-sql",
+            action="store_true",
+            help="Only use with --plan. Also prints SQL for each migration to be applied.",
+        )
 
     def handle(self, *args, **options):
+        for index, host in enumerate(CLICKHOUSE_MIGRATION_HOSTS_HTTP_URLS):
+            print(f"Updating host {host} ({index + 1}/{len(CLICKHOUSE_MIGRATION_HOSTS_HTTP_URLS)})")
+            self.migrate(host, options)
+
+    def migrate(self, host, options):
         database = Database(
             CLICKHOUSE_DATABASE,
-            db_url=CLICKHOUSE_HTTP_URL,
+            db_url=host,
             username=CLICKHOUSE_USER,
             password=CLICKHOUSE_PASSWORD,
             verify_ssl_cert=False,
         )
 
         if options["plan"]:
             print("List of clickhouse migrations to be applied:")
-            for migration_name in self.get_migrations(database, options["upto"]):
+            migrations = list(self.get_migrations(database, options["upto"]))
+            for migration_name, operations in migrations:
                 print(f"Migration would get applied: {migration_name}")
-            else:
+                for op in operations:
+                    sql = getattr(op, "_sql")
+                    if options["print_sql"] and sql is not None:
+                        print(indent("\n\n".join(sql), "    "))
+            if len(migrations) == 0:
                 print("Clickhouse migrations up to date!")
         elif options["fake"]:
             for migration_name in self.get_migrations(database, options["upto"]):
@@ -60,7 +81,7 @@ def get_migrations(self, database, upto):
         unapplied_migrations = set(modules.keys()) - applied_migrations
 
         for migration_name in sorted(unapplied_migrations):
-            yield migration_name
+            yield migration_name, modules[migration_name].operations
 
             if int(migration_name[:4]) >= upto:
                 break

posthog/settings.py

@@ -157,7 +157,12 @@ def print_warning(warning_lines: Sequence[str]):
     _clickhouse_http_protocol = "https://"
     _clickhouse_http_port = "8443"
 
-CLICKHOUSE_HTTP_URL = _clickhouse_http_protocol + CLICKHOUSE_HOST + ":" + _clickhouse_http_port + "/"
+CLICKHOUSE_HTTP_URL = f"{_clickhouse_http_protocol}{CLICKHOUSE_HOST}:{_clickhouse_http_port}/"
+
+_clickhouse_hosts = get_from_env("CLICKHOUSE_MIGRATION_HOSTS", CLICKHOUSE_HOST).split(",")
+CLICKHOUSE_MIGRATION_HOSTS_HTTP_URLS = [
+    f"{_clickhouse_http_protocol}{host}:{_clickhouse_http_port}/" for host in _clickhouse_hosts
+]
 
 IS_HEROKU = get_from_env("IS_HEROKU", False, type_cast=strtobool)
 

task-definition.migration.json

@@ -106,6 +106,10 @@
                     "name": "CLICKHOUSE_PASSWORD",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_PASSWORD::"
                 },
+                {
+                    "name": "CLICKHOUSE_MIGRATION_HOSTS",
+                    "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_MIGRATION_HOSTS::"
+                },
                 {
                     "name": "DATABASE_URL",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:DATABASE_URL::"

task-definition.web.json

@@ -120,6 +120,10 @@
                     "name": "CLICKHOUSE_HOST",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_HOST::"
                 },
+                {
+                    "name": "CLICKHOUSE_MIGRATION_HOSTS",
+                    "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_MIGRATION_HOSTS::"
+                },
                 {
                     "name": "CLICKHOUSE_PASSWORD",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_PASSWORD::"

task-definition.worker.json

@@ -114,6 +114,10 @@
                     "name": "CLICKHOUSE_HOST",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_HOST::"
                 },
+                {
+                    "name": "CLICKHOUSE_MIGRATION_HOSTS",
+                    "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_MIGRATION_HOSTS::"
+                },
                 {
                     "name": "CLICKHOUSE_PASSWORD",
                     "valueFrom": "arn:aws:secretsmanager:us-east-1:795637471508:secret:Posthog_Production_Heroku-FQ2itU:CLICKHOUSE_PASSWORD::"