Bugfix in getting sync models

src/django_clickhouse/models.py

@@ -3,7 +3,7 @@ This file contains base django model to be synced with ClickHouse.
 It saves all operations to storage in order to write them to ClickHouse later.
 """
 
-from typing import Optional, Any, List, Type
+from typing import Optional, Any, List, Type, Set
 
 import six
 from django.db import transaction
@@ -128,7 +128,6 @@ class ClickHouseSyncModel(DjangoModel):
     """
     Base model for syncing data. Each django model synced with data must inherit this
     """
-    _clickhouse_sync_models = []
     objects = ClickHouseSyncManager()
 
     class Meta:
@@ -151,13 +150,16 @@ class ClickHouseSyncModel(DjangoModel):
         :param model_cls: Model class to register
         :return: None
         """
-        cls._clickhouse_sync_models.append(model_cls)
+        if not hasattr(cls, '_clickhouse_sync_models'):
+            cls._clickhouse_sync_models = set()
+
+        cls._clickhouse_sync_models.add(model_cls)
 
     @classmethod
-    def get_clickhouse_sync_models(cls):  # type: () -> List['django_clickhouse.clickhouse_models.ClickHouseModel']
+    def get_clickhouse_sync_models(cls):  # type: () -> Set['django_clickhouse.clickhouse_models.ClickHouseModel']
         """
         Returns all clickhouse models, listening to this class
-        :return:
+        :return: A set of model classes to sync
         """
         return cls._clickhouse_sync_models
 

src/django_clickhouse/storages.py

@@ -8,15 +8,15 @@ Storage should be able to restore current importing batch, if something goes wro
 """
 import datetime
 import logging
-import os
 from typing import Any, Optional, List, Tuple
 
+import os
 from statsd.defaults.django import statsd
 
-from django_clickhouse.redis import redis_zadd
-from django_clickhouse.utils import check_pid
 from .configuration import config
 from .exceptions import ConfigurationError, RedisLockTimeoutError
+from .redis import redis_zadd
+from .utils import check_pid, get_subclasses
 
 logger = logging.getLogger('django-clickhouse')
 
@@ -254,6 +254,12 @@ class RedisStorage(Storage):
             if keys:
                 self._redis.delete(*keys)
 
+        from .clickhouse_models import ClickHouseModel
+        for model in get_subclasses(ClickHouseModel):
+            if isinstance(model.get_storage(), self.__class__):
+                key = "%s.sync.%s.queue" % (config.STATSD_PREFIX, model.get_import_key())
+                statsd.gauge(key, 0)
+
     def get_last_sync_time(self, import_key):
         sync_ts_key = self.REDIS_KEY_LAST_SYNC_TS.format(import_key=import_key)
         res = self._redis.get(sync_ts_key)

tests/clickhouse_models.py

@@ -2,7 +2,7 @@ from infi.clickhouse_orm import fields
 
 from django_clickhouse.clickhouse_models import ClickHouseModel, ClickHouseMultiModel
 from django_clickhouse.engines import ReplacingMergeTree, CollapsingMergeTree
-from tests.models import TestModel
+from tests.models import TestModel, SecondTestModel
 
 
 class ClickHouseTestModel(ClickHouseModel):
@@ -37,3 +37,15 @@ class ClickHouseMultiTestModel(ClickHouseMultiModel):
     sub_models = [ClickHouseTestModel, ClickHouseCollapseTestModel]
     sync_delay = 2
     sync_enabled = True
+
+
+class ClickHouseSecondTestModel(ClickHouseModel):
+    django_model = SecondTestModel
+    sync_delay = 2
+    sync_enabled = True
+
+    id = fields.Int32Field()
+    created_date = fields.DateField()
+    value = fields.Int32Field()
+
+    engine = ReplacingMergeTree('created_date', ('id',))

tests/migrations/0001_initial.py

@@ -23,4 +23,15 @@ class Migration(migrations.Migration):
                 'abstract': False,
             },
         ),
+        migrations.CreateModel(
+            name='SecondTestModel',
+            fields=[
+                ('id', models.AutoField()),
+                ('value', models.IntegerField()),
+                ('created_date', models.DateField()),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
     ]

tests/models.py

@@ -9,3 +9,8 @@ from django_clickhouse.models import ClickHouseSyncModel
 class TestModel(ClickHouseSyncModel):
     value = models.IntegerField()
     created_date = models.DateField()
+
+
+class SecondTestModel(ClickHouseSyncModel):
+    value = models.IntegerField()
+    created_date = models.DateField()

tests/test_models.py

@@ -2,8 +2,9 @@ import datetime
 
 from django.test import TransactionTestCase
 
-from tests.clickhouse_models import ClickHouseTestModel
+from tests.clickhouse_models import ClickHouseTestModel, ClickHouseSecondTestModel, ClickHouseCollapseTestModel, \
-from tests.models import TestModel
+    ClickHouseMultiTestModel
+from tests.models import TestModel, SecondTestModel
 
 
 # TestCase can't be used here:
@@ -110,3 +111,8 @@ class ClickHouseDjangoModelTest(TransactionTestCase):
     def test_qs_delete(self):
         TestModel.objects.filter(pk=1).delete()
         self.assertListEqual([('delete', 'default.1')], self.storage.get_operations(ClickHouseTestModel.get_import_key(), 10))
+
+    def test_clickhouse_sync_models(self):
+        self.assertSetEqual({ClickHouseSecondTestModel}, SecondTestModel.get_clickhouse_sync_models())
+        self.assertSetEqual({ClickHouseTestModel, ClickHouseCollapseTestModel, ClickHouseMultiTestModel},
+                            TestModel.get_clickhouse_sync_models())