From 369320b390d132799804728844673215540a04b2 Mon Sep 17 00:00:00 2001
From: Karl-Aksel Puulmann <oxymaccy@gmail.com>
Date: Fri, 16 Jul 2021 14:41:06 +0300
Subject: [PATCH] Separate way to migrate without replication

Will be used locally/in tests
---
 setup.py                              |  2 +-
 src/infi/clickhouse_orm/database.py   | 14 ++++++++------
 src/infi/clickhouse_orm/migrations.py | 11 +++++++----
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/setup.py b/setup.py
index 7d78984..085a3f5 100644
--- a/setup.py
+++ b/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 = 'fuziontech@gmail.com',
 
diff --git a/src/infi/clickhouse_orm/database.py b/src/infi/clickhouse_orm/database.py
index d42e224..64de0cb 100644
--- a/src/infi/clickhouse_orm/database.py
+++ b/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):
diff --git a/src/infi/clickhouse_orm/migrations.py b/src/infi/clickhouse_orm/migrations.py
index 8ed2d00..17dba5d 100644
--- a/src/infi/clickhouse_orm/migrations.py
+++ b/src/infi/clickhouse_orm/migrations.py
@@ -280,6 +280,13 @@ 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"),
@@ -287,10 +294,6 @@ class MigrationHistory(Model):
         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)