diff --git a/src/django_clickhouse/models.py b/src/django_clickhouse/models.py index 698eda2..98d37bb 100644 --- a/src/django_clickhouse/models.py +++ b/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 diff --git a/src/django_clickhouse/storages.py b/src/django_clickhouse/storages.py index 50eaba4..2c7babf 100644 --- a/src/django_clickhouse/storages.py +++ b/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) diff --git a/tests/clickhouse_models.py b/tests/clickhouse_models.py index 071ee6b..8b9cfb5 100644 --- a/tests/clickhouse_models.py +++ b/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',)) diff --git a/tests/migrations/0001_initial.py b/tests/migrations/0001_initial.py index ec3f5f0..ead88c1 100644 --- a/tests/migrations/0001_initial.py +++ b/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, + }, + ), ] diff --git a/tests/models.py b/tests/models.py index eefda05..701a7d0 100644 --- a/tests/models.py +++ b/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() diff --git a/tests/test_models.py b/tests/test_models.py index c517b39..1eb3257 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2,8 +2,9 @@ import datetime from django.test import TransactionTestCase -from tests.clickhouse_models import ClickHouseTestModel -from tests.models import TestModel +from tests.clickhouse_models import ClickHouseTestModel, ClickHouseSecondTestModel, ClickHouseCollapseTestModel, \ + 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())