mirror of
https://github.com/carrotquest/django-clickhouse.git
synced 2025-10-29 06:57:31 +03:00
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from typing import Optional, Iterable, List, Type
|
|
|
|
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):
|
|
"""
|
|
Basic QuerySet to use
|
|
"""
|
|
|
|
def __init__(self, model_cls: Type[InfiModel], database: Optional[Database] = None) -> None:
|
|
super(QuerySet, self).__init__(model_cls, database)
|
|
self._db_alias = None
|
|
|
|
@property
|
|
def _database(self) -> 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: Database) -> None:
|
|
# HACK for correct work of all infi.clickhouse-orm methods
|
|
self._db = database
|
|
|
|
def get_database(self, for_write: bool = False) -> 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: 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) -> '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: Iterable[InfiModel], batch_size: int = 1000) -> List[InfiModel]:
|
|
self.get_database(for_write=True).insert(model_instances=model_instances, batch_size=batch_size)
|
|
return list(model_instances)
|
|
|
|
|
|
class AggregateQuerySet(QuerySet, InfiAggregateQuerySet):
|
|
pass
|