mirror of
https://github.com/carrotquest/django-clickhouse.git
synced 2024-11-22 00:56:37 +03:00
This commit is contained in:
parent
f9e33a6845
commit
54923e03bf
|
@ -5,6 +5,7 @@ It saves all operations to storage in order to write them to ClickHouse later.
|
||||||
|
|
||||||
from typing import Optional, Any, Type, Set
|
from typing import Optional, Any, Type, Set
|
||||||
|
|
||||||
|
import functools
|
||||||
import six
|
import six
|
||||||
from django.db import transaction
|
from django.db import transaction
|
||||||
from django.db.models import QuerySet as DjangoQuerySet, Model as DjangoModel, Manager as DjangoManager
|
from django.db.models import QuerySet as DjangoQuerySet, Model as DjangoModel, Manager as DjangoManager
|
||||||
|
@ -71,19 +72,25 @@ class ClickHouseSyncBulkUpdateQuerySetMixin(ClickHouseSyncRegisterMixin, BulkUpd
|
||||||
|
|
||||||
return returning
|
return returning
|
||||||
|
|
||||||
def pg_bulk_update(self, *args, **kwargs):
|
def _decorate_method(self, name: str, operation: str, args, kwargs):
|
||||||
|
if not hasattr(super(), name):
|
||||||
|
raise AttributeError(name)
|
||||||
|
|
||||||
|
func = getattr(super(), name)
|
||||||
original_returning = kwargs.pop('returning', None)
|
original_returning = kwargs.pop('returning', None)
|
||||||
kwargs['returning'] = self._update_returning_param(original_returning)
|
kwargs['returning'] = self._update_returning_param(original_returning)
|
||||||
result = super().pg_bulk_update(*args, **kwargs)
|
result = func(*args, **kwargs)
|
||||||
self._register_ops('update', result)
|
self._register_ops(operation, result)
|
||||||
return result.count() if original_returning is None else result
|
return result.count() if original_returning is None else result
|
||||||
|
|
||||||
|
def pg_bulk_update(self, *args, **kwargs):
|
||||||
|
return self._decorate_method('pg_bulk_update', 'update', args, kwargs)
|
||||||
|
|
||||||
def pg_bulk_update_or_create(self, *args, **kwargs):
|
def pg_bulk_update_or_create(self, *args, **kwargs):
|
||||||
original_returning = kwargs.pop('returning', None)
|
return self._decorate_method('pg_bulk_update_or_create', 'update', args, kwargs)
|
||||||
kwargs['returning'] = self._update_returning_param(original_returning)
|
|
||||||
result = super().pg_bulk_update_or_create(*args, **kwargs)
|
def pg_bulk_create(self, *args, **kwargs):
|
||||||
self._register_ops('update', result)
|
return self._decorate_method('pg_bulk_create', 'insert', args, kwargs)
|
||||||
return result.count() if original_returning is None else result
|
|
||||||
|
|
||||||
|
|
||||||
class ClickHouseSyncQuerySetMixin(ClickHouseSyncRegisterMixin):
|
class ClickHouseSyncQuerySetMixin(ClickHouseSyncRegisterMixin):
|
||||||
|
|
|
@ -76,6 +76,24 @@ class TestOperations(TransactionTestCase):
|
||||||
self.assertSetEqual({('update', "%s.%d" % (self.db_alias, instance.pk)) for instance in items},
|
self.assertSetEqual({('update', "%s.%d" % (self.db_alias, instance.pk)) for instance in items},
|
||||||
set(self.storage.get_operations(self.clickhouse_model.get_import_key(), 10)))
|
set(self.storage.get_operations(self.clickhouse_model.get_import_key(), 10)))
|
||||||
|
|
||||||
|
def test_pg_bulk_create(self):
|
||||||
|
now_dt = now()
|
||||||
|
res = self.django_model.objects.pg_bulk_create([
|
||||||
|
{'value': i, 'created': now_dt, 'created_date': now_dt.date()}
|
||||||
|
for i in range(5)
|
||||||
|
])
|
||||||
|
self.assertEqual(5, res)
|
||||||
|
|
||||||
|
items = list(self.django_model.objects.filter(value__lt=100).order_by('value'))
|
||||||
|
self.assertEqual(5, len(items))
|
||||||
|
for i, instance in enumerate(items):
|
||||||
|
self.assertEqual(instance.created, now_dt)
|
||||||
|
self.assertEqual(instance.created_date, now_dt.date())
|
||||||
|
self.assertEqual(i, instance.value)
|
||||||
|
|
||||||
|
self.assertSetEqual({('insert', "%s.%d" % (self.db_alias, instance.pk)) for instance in items},
|
||||||
|
set(self.storage.get_operations(self.clickhouse_model.get_import_key(), 10)))
|
||||||
|
|
||||||
def test_pg_bulk_update(self):
|
def test_pg_bulk_update(self):
|
||||||
items = list(self.django_model.objects.filter(pk__in={1, 2}))
|
items = list(self.django_model.objects.filter(pk__in={1, 2}))
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue
Block a user