From 2d4c8554cd7aed5031979f78809f7bc5abfbff7f Mon Sep 17 00:00:00 2001
From: M1ha <work_shvein_mihail@mail.ru>
Date: Tue, 18 Dec 2018 15:45:10 +0500
Subject: [PATCH] Changed migrations routing, dividing it into replicated and
 not replicated operations

---
 src/django_clickhouse/clickhouse_models.py |  3 ++-
 src/django_clickhouse/migrations.py        |  2 +-
 src/django_clickhouse/routers.py           | 12 +++++++++---
 tests/clickhouse_models.py                 |  3 ++-
 tests/test_migrations.py                   |  2 +-
 5 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/django_clickhouse/clickhouse_models.py b/src/django_clickhouse/clickhouse_models.py
index 91e93f8..5f7d578 100644
--- a/src/django_clickhouse/clickhouse_models.py
+++ b/src/django_clickhouse/clickhouse_models.py
@@ -47,7 +47,8 @@ class ClickHouseModel(with_metaclass(ClickHouseModelMeta, InfiModel)):
 
     read_db_aliases = (config.DEFAULT_DB_ALIAS,)
     write_db_aliases = (config.DEFAULT_DB_ALIAS,)
-    migrate_db_aliases = write_db_aliases
+    migrate_replicated_db_aliases = (config.DEFAULT_DB_ALIAS,)
+    migrate_non_replicated_db_aliases = (config.DEFAULT_DB_ALIAS,)
 
     sync_enabled = False
     sync_batch_size = None
diff --git a/src/django_clickhouse/migrations.py b/src/django_clickhouse/migrations.py
index 564da68..e185c47 100644
--- a/src/django_clickhouse/migrations.py
+++ b/src/django_clickhouse/migrations.py
@@ -34,7 +34,7 @@ class Migration:
             model_class = getattr(op, 'model_class', None)
             hints = getattr(op, 'hints', {})
 
-            if db_router.allow_migrate(db_alias, self.__module__, model=model_class, **hints):
+            if db_router.allow_migrate(db_alias, self.__module__, op, model=model_class, **hints):
                 database = database or connections[db_alias]
                 op.apply(database)
 
diff --git a/src/django_clickhouse/routers.py b/src/django_clickhouse/routers.py
index 008c2e0..f56de90 100644
--- a/src/django_clickhouse/routers.py
+++ b/src/django_clickhouse/routers.py
@@ -5,6 +5,7 @@ from typing import Optional
 
 import random
 import six
+from infi.clickhouse_orm.migrations import Operation, DropTable, CreateTable
 
 from .clickhouse_models import ClickHouseModel
 from .configuration import config
@@ -32,12 +33,13 @@ class DefaultRouter:
         """
         return random.choice(model.write_db_aliases)
 
-    def allow_migrate(self, db_alias, app_label, model=None, **hints):
-        # type: (str, str, Optional[ClickHouseModel], **dict) -> bool
+    def allow_migrate(self, db_alias, app_label, operation, model=None, **hints):
+        # type: (str, str, Operation, Optional[ClickHouseModel], **dict) -> bool
         """
         Checks if migration can be applied to given database
         :param db_alias: Database alias to check
         :param app_label: App from which migration is got
+        :param operation: Operation object to perform
         :param model: Model migration is applied to
         :param hints: Hints to make correct decision
         :return: boolean
@@ -50,4 +52,8 @@ class DefaultRouter:
                 if isinstance(hints['model'], six.string_types) else hints['model']
 
         model = lazy_class_import(model)
-        return db_alias in model.migrate_db_aliases
+
+        if operation.__class__ not in {CreateTable, DropTable}:
+            return db_alias in model.migrate_replicated_db_aliases
+        else:
+            return db_alias in model.migrate_non_replicated_db_aliases
diff --git a/tests/clickhouse_models.py b/tests/clickhouse_models.py
index 8b347c0..071ee6b 100644
--- a/tests/clickhouse_models.py
+++ b/tests/clickhouse_models.py
@@ -15,7 +15,8 @@ class ClickHouseTestModel(ClickHouseModel):
     value = fields.Int32Field()
 
     engine = ReplacingMergeTree('created_date', ('id',))
-    migrate_db_aliases = ('default', 'secondary')
+    migrate_replicated_db_aliases = ('default', 'secondary')
+    migrate_non_replicated_db_aliases = ('default', 'secondary')
 
 
 class ClickHouseCollapseTestModel(ClickHouseModel):
diff --git a/tests/test_migrations.py b/tests/test_migrations.py
index 3e5693b..333e7ee 100644
--- a/tests/test_migrations.py
+++ b/tests/test_migrations.py
@@ -8,7 +8,7 @@ from tests.clickhouse_models import ClickHouseTestModel
 
 
 class NoMigrateRouter(DefaultRouter):
-    def allow_migrate(self, db_alias, app_label, model=None, **hints):
+    def allow_migrate(self, db_alias, app_label, operation, model=None, **hints):
         return False