Added QuerySet that determines db with router help.

Added using(), create(), bulk_create(), all() methods to QuerySet

src/django_clickhouse/__init__.py

@@ -1 +0,0 @@
-from .configuration import config

src/django_clickhouse/clickhouse_models.py

@@ -14,8 +14,9 @@ from infi.clickhouse_orm.models import Model as InfiModel, ModelBase as InfiMode
 from six import with_metaclass
 from statsd.defaults.django import statsd
 
+from .query import QuerySet
 from .configuration import config
-from .database import connections, DEFAULT_DB_ALIAS
+from .database import connections
 from .models import ClickHouseSyncModel
 from .serializers import Django2ClickHouseModelSerializer
 from .utils import lazy_class_import
@@ -31,6 +32,8 @@ class ClickHouseModelMeta(InfiModelBase):
         if res.django_model and res.get_sync_delay():
             res.django_model.register_clickhouse_sync_model(res)
 
+        res.objects = QuerySet(res)
+
         return res
 
 
@@ -41,8 +44,9 @@ class ClickHouseModel(with_metaclass(ClickHouseModelMeta, InfiModel)):
     django_model = None
     django_model_serializer = Django2ClickHouseModelSerializer
 
-    read_db_aliases = (DEFAULT_DB_ALIAS,)
+    read_db_aliases = (config.DEFAULT_DB_ALIAS,)
-    write_db_aliases = (DEFAULT_DB_ALIAS,)
+    write_db_aliases = (config.DEFAULT_DB_ALIAS,)
+    migrate_db_aliases = write_db_aliases
 
     sync_enabled = False
     sync_batch_size = None
@@ -52,12 +56,16 @@ class ClickHouseModel(with_metaclass(ClickHouseModelMeta, InfiModel)):
     sync_lock_timeout = None
 
     @classmethod
-    def get_database(cls, for_write=False):
+    def objects_in(cls, database):  # type: (Database) -> QuerySet
-        # type: (bool) -> Database
+        return QuerySet(cls, database)
+
+    @classmethod
+    def get_database_alias(cls, for_write=False):
+        # type: (bool) -> str
         """
-        Gets database for read or write purposes
+        Gets database alias for read or write purposes
         :param for_write: Boolean flag if database is neede for read or for write
-        :return: infi.clickhouse_orm.Database instance
+        :return: Database alias to use
         """
         db_router = lazy_class_import(config.DATABASE_ROUTER)()
         if for_write:
@@ -65,6 +73,17 @@ class ClickHouseModel(with_metaclass(ClickHouseModelMeta, InfiModel)):
         else:
             return db_router.db_for_read(cls)
 
+    @classmethod
+    def get_database(cls, for_write=False):
+        # type: (bool) -> Database
+        """
+        Gets database alias for read or write purposes
+        :param for_write: Boolean flag if database is neede for read or for write
+        :return: infi.clickhouse_orm.Database instance
+        """
+        db_alias = cls.get_database_alias(for_write=for_write)
+        return connections[db_alias]
+
     @classmethod
     def get_django_model_serializer(cls, writable=False):  # type: (bool) -> Django2ClickHouseModelSerializer
         serializer_cls = lazy_class_import(cls.django_model_serializer)

src/django_clickhouse/configuration.py

@@ -21,7 +21,8 @@ DEFAULTS = {
     'STATSD_PREFIX': 'clickhouse',
     'MIGRATIONS_PACKAGE': 'clickhouse_migrations',
     'MIGRATE_WITH_DEFAULT_DB': True,
-    'CELERY_QUEUE': 'celery'
+    'CELERY_QUEUE': 'celery',
+    'DEFAULT_DB_ALIAS': 'default'
 }
 
 

src/django_clickhouse/database.py

@@ -3,8 +3,6 @@ from infi.clickhouse_orm.database import Database as InfiDatabase
 from .configuration import config
 from .exceptions import DBAliasError
 
-DEFAULT_DB_ALIAS = 'default'
-
 
 class Database(InfiDatabase):
     def __init__(self, **kwargs):
@@ -30,7 +28,7 @@ class ConnectionProxy:
 
     def get_connection(self, alias):
         if alias is None:
-            alias = DEFAULT_DB_ALIAS
+            alias = config.DEFAULT_DB_ALIAS
 
         if alias not in self._connections:
             if alias not in config.DATABASES:

src/django_clickhouse/models.py

@@ -6,7 +6,7 @@ It saves all operations to storage in order to write them to ClickHouse later.
 from typing import Optional, Any, List, Type
 
 import six
-from django.db import transaction, DEFAULT_DB_ALIAS
+from django.db import transaction
 from django.db.models.signals import post_save, post_delete
 from django.dispatch import receiver
 from django.db.models import QuerySet as DjangoQuerySet, Manager as DjangoManager, Model as DjangoModel
@@ -171,7 +171,7 @@ class ClickHouseSyncModel(DjangoModel):
         :param using: Database alias registered instances are from
         :return: None
         """
-        model_pks = ['%s.%d' % (using or DEFAULT_DB_ALIAS, pk) for pk in model_pks]
+        model_pks = ['%s.%d' % (using or config.DEFAULT_DB_ALIAS, pk) for pk in model_pks]
 
         def _on_commit():
             for model_cls in cls.get_clickhouse_sync_models():

src/django_clickhouse/query.py

@@ -1,9 +1,79 @@
+from typing import Optional, Iterable
+
+from copy import copy
+from infi.clickhouse_orm.database import Database
+from infi.clickhouse_orm.models import Model as InfiModel
 from infi.clickhouse_orm.query import QuerySet as InfiQuerySet, AggregateQuerySet as InfiAggregateQuerySet
 
+from .database import connections
+
 
 class QuerySet(InfiQuerySet):
-    pass
+    """
+    Basic QuerySet to use
+    """
+
+    def __init__(self, model_cls, database=None):  # type: (InfiModel, Optional[Database]) -> None
+        super(QuerySet, self).__init__(model_cls, database)
+        self._db_alias = None
+
+    @property
+    def _database(self):  # type: () -> Database
+        # HACK for correct work of all infi.clickhouse-orm methods
+        # There are no write QuerySet methods now, so I use for_write=False by default
+        return self.get_database(for_write=False)
+
+    @_database.setter
+    def _database(self, database):  # type: (Database) -> None
+        # HACK for correct work of all infi.clickhouse-orm methods
+        self._db = database
+
+    def get_database(self, for_write=False):  # type: (bool) -> Database
+        """
+        Gets database to execute query on. Looks for constructor or using() method.
+        If nothing was set tries to get database from model class using router.
+        :param for_write: Return QuerySet for read or for write.
+        :return: Database instance
+        """
+        if not self._db:
+            if self._db_alias:
+                self._db = connections[self._db_alias] 
+            else:
+                self._db = self._model_cls.get_database(for_write=for_write)
+
+        return self._db
+
+    def using(self, db_alias):  # type: (str) -> QuerySet
+        """
+        Sets database alias to use for this query
+        :param db_alias: Database alias name from CLICKHOUSE_DATABASES config option
+        :return: None
+        """
+        qs = copy(self)
+        qs._db_alias = db_alias
+        qs._db = None  # Previous database should be forgotten
+        return qs
+
+    def all(self):  # type: () -> QuerySet
+        """
+        Returns all items of queryset
+        :return: QuerySet
+        """
+        return copy(self)
+
+    def create(self, **kwargs):
+        """
+        Create single item in database
+        :return: Created instance
+        """
+        instance = self._model_cls(**kwargs)
+        self.get_database(for_write=True).insert([instance])
+        return instance
+
+    def bulk_create(self, model_instances, batch_size=1000):  # type: (Iterable[InfiModel], int)
+        self.get_database(for_write=True).insert(model_instances=model_instances, batch_size=batch_size)
+        return model_instances
 
 
-class AggregateQuerySet(InfiAggregateQuerySet):
+class AggregateQuerySet(QuerySet, InfiAggregateQuerySet):
     pass

src/django_clickhouse/routers.py

@@ -50,4 +50,4 @@ class DefaultRouter:
                 if isinstance(hints['model'], six.string_types) else hints['model']
 
         model = lazy_class_import(model)
-        return db_alias in model.write_db_aliases
+        return db_alias in model.migrate_db_aliases

tests/clickhouse_models.py

@@ -14,6 +14,7 @@ class ClickHouseTestModel(ClickHouseModel):
     value = fields.Int32Field()
 
     engine = ReplacingMergeTree('created_date', ('id',))
+    migrate_db_aliases = ('default', 'secondary')
 
 
 class ClickHouseCollapseTestModel(ClickHouseModel):

tests/settings.py

@@ -69,6 +69,11 @@ CLICKHOUSE_DATABASES = {
         'db_name': 'test',
         'username': 'default',
         'password': ''
+    },
+    'secondary': {
+        'db_name': 'test_2',
+        'username': 'default',
+        'password': ''
     }
 }
 

tests/test_config.py

@@ -1,6 +1,6 @@
 from django.test import TestCase
 
-from django_clickhouse import config
+from django_clickhouse.configuration import config
 
 
 class ConfigTest(TestCase):

tests/test_query.py

@@ -0,0 +1,66 @@
+from unittest import TestCase
+
+import datetime
+
+from django_clickhouse.database import connections
+from django_clickhouse.migrations import migrate_app
+from django_clickhouse.query import QuerySet
+from tests.clickhouse_models import ClickHouseTestModel
+
+
+class TestQuerySet(TestCase):
+    def _recreate_db(self, db_alias):
+        db = connections[db_alias]
+        db.drop_database()
+        db.db_exists = False
+        db.create_database()
+
+        migrate_app('tests', db_alias)
+        return db
+
+    def setUp(self):
+        self.db = self._recreate_db('default')
+        self._recreate_db('secondary')
+
+    def test_all(self):
+        self.db.insert([ClickHouseTestModel(id=i, created_date=datetime.date.today(), value=i) for i in range(1, 4)])
+        qs = ClickHouseTestModel.objects.all()
+        print(qs.get_database(for_write=True).db_name)
+        self.assertIsInstance(qs, QuerySet)
+        self.assertEqual(3, qs.count())
+
+    def test_create(self):
+        ClickHouseTestModel.objects.create(id=1, created_date=datetime.date.today(), value=2)
+        res = list(self.db.select('SELECT * FROM $table', model_class=ClickHouseTestModel))
+        self.assertEqual(1, len(res))
+
+        self.assertEqual(1, res[0].id)
+        self.assertEqual(datetime.date.today(), res[0].created_date)
+        self.assertEqual(2, res[0].value)
+
+    def test_bulk_create(self):
+        ClickHouseTestModel.objects.bulk_create([
+            ClickHouseTestModel(id=i, created_date=datetime.date.today(), value=i) for i in range(1, 4)
+        ])
+        res = list(self.db.select('SELECT * FROM $table ORDER BY id', model_class=ClickHouseTestModel))
+        self.assertEqual(3, len(res))
+
+        for i in range(0, 3):
+            self.assertEqual(i + 1, res[i].id)
+            self.assertEqual(datetime.date.today(), res[0].created_date)
+            self.assertEqual(i + 1, res[i].value)
+
+    def test_using(self):
+        self.db.insert(
+            [ClickHouseTestModel(id=i, created_date=datetime.date.today(), value=i) for i in range(1, 4)]
+        )
+        connections['secondary'].insert([
+            ClickHouseTestModel(id=i, created_date=datetime.date.today(), value=i) for i in range(10, 12)
+        ])
+
+        self.assertEqual(3, ClickHouseTestModel.objects.count())
+        self.assertEqual(3, ClickHouseTestModel.objects_in(self.db).count())
+        self.assertEqual(2, ClickHouseTestModel.objects_in(self.db).using('secondary').count())
+
+        self.assertEqual(2, ClickHouseTestModel.objects.using('secondary').count())
+        self.assertEqual(3, ClickHouseTestModel.objects.using('default').count())

tests/test_sync.py

@@ -1,17 +1,14 @@
 import datetime
 import signal
-
-import os
 from multiprocessing import Process
 from time import sleep
 from unittest import skip, expectedFailure
 
-import random
+import os
 from django.db import connections as django_connections
 from django.db.models import F
-from django.test import TransactionTestCase, override_settings
+from django.test import TransactionTestCase
 
-from django_clickhouse import config
 from django_clickhouse.database import connections
 from django_clickhouse.migrations import migrate_app
 from tests.clickhouse_models import ClickHouseTestModel, ClickHouseCollapseTestModel, ClickHouseMultiTestModel