Merge pull request #5 from PostHog/fix-migrations-by-not-replicating

Separate way to migrate without replication

setup.py

@@ -1,7 +1,7 @@
 
 SETUP_INFO = dict(
     name = 'infi.clickhouse_orm',
-    version = '2.1.0.post11',
+    version = '2.1.0.post15',
     author = 'James Greenhill',
     author_email = '[email protected]',
 

src/infi/clickhouse_orm/database.py

@@ -330,7 +330,7 @@ def paginate(self, model_class, order_by, page_num=1, page_size=100, conditions=
             page_size=page_size
         )
 
-    def migrate(self, migrations_package_name, up_to=9999):
+    def migrate(self, migrations_package_name, up_to=9999, replicated=False):
         '''
         Executes schema migrations.
 
@@ -340,7 +340,7 @@ def migrate(self, migrations_package_name, up_to=9999):
         '''
         from .migrations import MigrationHistory
         logger = logging.getLogger('migrations')
-        applied_migrations = self._get_applied_migrations(migrations_package_name)
+        applied_migrations = self._get_applied_migrations(migrations_package_name, replicated=replicated)
         modules = import_submodules(migrations_package_name)
         unapplied_migrations = set(modules.keys()) - applied_migrations
         for name in sorted(unapplied_migrations):
@@ -351,11 +351,13 @@ def migrate(self, migrations_package_name, up_to=9999):
             if int(name[:4]) >= up_to:
                 break
 
-    def _get_applied_migrations(self, migrations_package_name):
-        from .migrations import MigrationHistory
-        self.create_table(MigrationHistory)
+    def _get_applied_migrations(self, migrations_package_name, replicated):
+        from .migrations import MigrationHistory, MigrationHistoryReplicated
+
+        Model = MigrationHistoryReplicated if replicated else MigrationHistory
+        self.create_table(Model)
         query = "SELECT module_name from $table WHERE package_name = '%s'" % migrations_package_name
-        query = self._substitute(query, MigrationHistory)
+        query = self._substitute(query, Model)
         return set(obj.module_name for obj in self.select(query))
 
     def _send(self, data, settings=None, stream=False):

src/infi/clickhouse_orm/migrations.py

@@ -280,17 +280,20 @@ class MigrationHistory(Model):
     module_name = StringField()
     applied = DateField()
 
+    engine = MergeTree('applied', ('package_name', 'module_name'))
+
+    @classmethod
+    def table_name(cls):
+        return "infi_clickhouse_orm_migrations"
+
+class MigrationHistoryReplicated(Model):
     engine = MergeTree(
         "applied",
         ("package_name", "module_name"),
         replica_table_path="/clickhouse/prod/tables/noshard/posthog.infi_clickhouse_orm_migrations",
         replica_name="{replica}-{shard}",
     )
 
-    @classmethod
-    def table_name(cls):
-        return "infi_clickhouse_orm_migrations"
-
 
 # Expose only relevant classes in import *
 __all__ = get_subclass_names(locals(), Operation)