mirror of
https://github.com/carrotquest/django-clickhouse.git
synced 2024-11-10 19:36:38 +03:00
Added QuerySet that determines db with router help.
Added using(), create(), bulk_create(), all() methods to QuerySet
This commit is contained in:
parent
ce1a206b04
commit
abe0ae45c5
|
@ -1 +0,0 @@
|
|||
from .configuration import config
|
|
@ -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,)
|
||||
write_db_aliases = (DEFAULT_DB_ALIAS,)
|
||||
read_db_aliases = (config.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):
|
||||
# type: (bool) -> Database
|
||||
def objects_in(cls, database): # type: (Database) -> QuerySet
|
||||
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)
|
||||
|
|
|
@ -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'
|
||||
}
|
||||
|
||||
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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():
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -14,6 +14,7 @@ class ClickHouseTestModel(ClickHouseModel):
|
|||
value = fields.Int32Field()
|
||||
|
||||
engine = ReplacingMergeTree('created_date', ('id',))
|
||||
migrate_db_aliases = ('default', 'secondary')
|
||||
|
||||
|
||||
class ClickHouseCollapseTestModel(ClickHouseModel):
|
||||
|
|
|
@ -69,6 +69,11 @@ CLICKHOUSE_DATABASES = {
|
|||
'db_name': 'test',
|
||||
'username': 'default',
|
||||
'password': ''
|
||||
},
|
||||
'secondary': {
|
||||
'db_name': 'test_2',
|
||||
'username': 'default',
|
||||
'password': ''
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
from django.test import TestCase
|
||||
|
||||
from django_clickhouse import config
|
||||
from django_clickhouse.configuration import config
|
||||
|
||||
|
||||
class ConfigTest(TestCase):
|
||||
|
|
66
tests/test_query.py
Normal file
66
tests/test_query.py
Normal file
|
@ -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())
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue
Block a user