491 lines
15 KiB
Python
491 lines
15 KiB
Python
import threading
|
|
import functools
|
|
from django.db import models, transaction
|
|
from django.db.models import Q, F, Count, Sum, Avg, Min, Max, Subquery, OuterRef, Case, When, Value
|
|
from gvsds.model_base import FieldExpression, snake_to_pascal, pascal_to_snake
|
|
|
|
|
|
_local = threading.local()
|
|
|
|
|
|
def _get_session():
|
|
if not hasattr(_local, 'db_session'):
|
|
_local.db_session = Session()
|
|
return _local.db_session
|
|
|
|
|
|
def _get_cache():
|
|
if not hasattr(_local, 'query_cache'):
|
|
_local.query_cache = {}
|
|
return _local.query_cache
|
|
|
|
|
|
def _invalidate_cache():
|
|
if hasattr(_local, 'query_cache'):
|
|
_local.query_cache.clear()
|
|
|
|
|
|
class RelationAccessor:
|
|
def __init__(self, model_class, prefix=''):
|
|
self._model = model_class
|
|
self._prefix = prefix
|
|
|
|
def __getattr__(self, name):
|
|
if name.startswith('_'):
|
|
return object.__getattribute__(self, name)
|
|
try:
|
|
self._model._meta.get_field(pascal_to_snake(name))
|
|
full_name = f'{self._prefix}__{pascal_to_snake(name)}' if self._prefix else pascal_to_snake(name)
|
|
return FieldExpression(self._model, full_name)
|
|
except Exception:
|
|
for rel in self._model._meta.get_fields(include_hidden=True):
|
|
if hasattr(rel, 'related_model') and rel.related_model:
|
|
rel_name = ''
|
|
if hasattr(rel, 'get_accessor_name'):
|
|
rel_name = rel.get_accessor_name()
|
|
elif hasattr(rel, 'name'):
|
|
rel_name = rel.name
|
|
if rel_name and (rel_name == pascal_to_snake(name) or snake_to_pascal(rel_name) == name):
|
|
full_prefix = f'{self._prefix}__{rel_name}' if self._prefix else rel_name
|
|
return RelationAccessor(rel.related_model, full_prefix)
|
|
return FieldExpression(self._model, name)
|
|
|
|
|
|
class _ValuesResult:
|
|
def __init__(self, qs, fields):
|
|
self._qs = qs
|
|
self._fields = fields
|
|
|
|
def filter(self, *args, **kwargs):
|
|
self._qs = self._qs.filter(*args, **kwargs)
|
|
return self
|
|
|
|
def first(self):
|
|
obj = self._qs.first()
|
|
if obj is None:
|
|
return None
|
|
if len(self._fields) == 1:
|
|
return getattr(obj, 'val0', None)
|
|
return tuple(getattr(obj, f'val{i}', None) for i in range(len(self._fields)))
|
|
|
|
def all(self):
|
|
results = []
|
|
for obj in self._qs:
|
|
if len(self._fields) == 1:
|
|
results.append(getattr(obj, 'val0', None))
|
|
else:
|
|
results.append(tuple(getattr(obj, f'val{i}', None) for i in range(len(self._fields))))
|
|
return results
|
|
|
|
def scalar(self):
|
|
obj = self._qs.first()
|
|
if obj is None:
|
|
return None
|
|
return getattr(obj, 'val0', None)
|
|
|
|
|
|
class FluentQuery:
|
|
def __init__(self, model_class):
|
|
self._model = model_class
|
|
self._qs = model_class.objects.all()
|
|
self._joined_models = []
|
|
self._cache_enabled = True
|
|
|
|
def _clone(self):
|
|
q = FluentQuery(self._model)
|
|
q._qs = self._qs
|
|
q._joined_models = list(self._joined_models)
|
|
q._cache_enabled = self._cache_enabled
|
|
return q
|
|
|
|
def _cache_key(self, method):
|
|
try:
|
|
query_str = str(self._qs.query)
|
|
except Exception:
|
|
query_str = id(self._qs)
|
|
return (self._model, method, query_str)
|
|
|
|
def _convert_kwargs(self, kwargs):
|
|
django_kwargs = {}
|
|
for key, val in kwargs.items():
|
|
parts = key.split('__')
|
|
if not parts[0].endswith('_id') and parts[0] != '':
|
|
try:
|
|
self._model._meta.get_field(parts[0])
|
|
except Exception:
|
|
parts[0] = pascal_to_snake(parts[0])
|
|
django_kwargs['__'.join(parts)] = val
|
|
return django_kwargs
|
|
|
|
def _remap_q(self, q):
|
|
new_children = []
|
|
for child in q.children:
|
|
if isinstance(child, Q):
|
|
new_children.append(self._remap_q(child))
|
|
elif isinstance(child, tuple):
|
|
key, val = child
|
|
new_key = self._remap_key(key)
|
|
new_children.append((new_key, val))
|
|
else:
|
|
new_children.append(child)
|
|
if new_children:
|
|
new_q = Q(*new_children)
|
|
new_q.connector = q.connector
|
|
new_q.negated = q.negated
|
|
return new_q
|
|
return q
|
|
|
|
def _remap_key(self, key):
|
|
parts = key.split('__')
|
|
field_name = parts[0]
|
|
try:
|
|
self._model._meta.get_field(field_name)
|
|
return key
|
|
except Exception:
|
|
pass
|
|
for joined in self._joined_models:
|
|
try:
|
|
joined._meta.get_field(field_name)
|
|
for rel in self._model._meta.get_fields(include_hidden=True):
|
|
if hasattr(rel, 'related_model') and rel.related_model == joined:
|
|
acc = ''
|
|
if hasattr(rel, 'get_accessor_name'):
|
|
acc = rel.get_accessor_name()
|
|
elif hasattr(rel, 'name'):
|
|
acc = rel.name
|
|
if acc and acc != '+':
|
|
return acc + '__' + key
|
|
return joined._meta.model_name + '__' + key
|
|
except Exception:
|
|
pass
|
|
return key
|
|
|
|
def filter(self, *args, **kwargs):
|
|
converted_args = []
|
|
for arg in args:
|
|
if isinstance(arg, Q) and self._joined_models:
|
|
converted_args.append(self._remap_q(arg))
|
|
else:
|
|
converted_args.append(arg)
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
self._qs = self._qs.filter(*converted_args, **django_kwargs)
|
|
return self
|
|
|
|
def exclude(self, *args, **kwargs):
|
|
converted_args = []
|
|
for arg in args:
|
|
if isinstance(arg, Q) and self._joined_models:
|
|
converted_args.append(self._remap_q(arg))
|
|
else:
|
|
converted_args.append(arg)
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
self._qs = self._qs.exclude(*converted_args, **django_kwargs)
|
|
return self
|
|
|
|
def order_by(self, *args):
|
|
converted = []
|
|
for a in args:
|
|
if isinstance(a, str):
|
|
if a.startswith('-'):
|
|
converted.append('-' + pascal_to_snake(a[1:]))
|
|
else:
|
|
converted.append(pascal_to_snake(a))
|
|
elif isinstance(a, FieldExpression):
|
|
converted.append(str(a))
|
|
else:
|
|
converted.append(str(a))
|
|
self._qs = self._qs.order_by(*converted)
|
|
return self
|
|
|
|
def first(self):
|
|
if self._cache_enabled:
|
|
cache = _get_cache()
|
|
key = self._cache_key('first')
|
|
if key in cache:
|
|
return cache[key]
|
|
obj = self._qs.first()
|
|
if obj is not None:
|
|
_get_session().track(obj)
|
|
if self._cache_enabled:
|
|
_get_cache()[self._cache_key('first')] = obj
|
|
return obj
|
|
|
|
def all(self):
|
|
if self._cache_enabled:
|
|
cache = _get_cache()
|
|
key = self._cache_key('all')
|
|
if key in cache:
|
|
return cache[key]
|
|
results = list(self._qs)
|
|
session = _get_session()
|
|
for obj in results:
|
|
session.track(obj)
|
|
if self._cache_enabled:
|
|
_get_cache()[self._cache_key('all')] = results
|
|
return results
|
|
|
|
def count(self):
|
|
if self._cache_enabled:
|
|
cache = _get_cache()
|
|
key = self._cache_key('count')
|
|
if key in cache:
|
|
return cache[key]
|
|
result = self._qs.count()
|
|
if self._cache_enabled:
|
|
_get_cache()[self._cache_key('count')] = result
|
|
return result
|
|
|
|
def get(self, *args, **kwargs):
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
if args:
|
|
obj = self._qs.filter(*args).get(**django_kwargs)
|
|
else:
|
|
obj = self._qs.get(**django_kwargs)
|
|
_get_session().track(obj)
|
|
return obj
|
|
|
|
def create(self, **kwargs):
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
obj = self._model.objects.create(**django_kwargs)
|
|
_invalidate_cache()
|
|
return obj
|
|
|
|
def get_or_create(self, defaults=None, **kwargs):
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
if defaults:
|
|
defaults = self._convert_kwargs(defaults)
|
|
obj, created = self._model.objects.get_or_create(defaults=defaults, **django_kwargs)
|
|
if created:
|
|
_invalidate_cache()
|
|
_get_session().track(obj)
|
|
return obj, created
|
|
|
|
def limit(self, n):
|
|
self._qs = self._qs[:n]
|
|
return self
|
|
|
|
def offset(self, n):
|
|
self._qs = self._qs[n:]
|
|
return self
|
|
|
|
def join(self, *args):
|
|
for arg in args:
|
|
if isinstance(arg, type) and issubclass(arg, models.Model):
|
|
if arg not in self._joined_models:
|
|
self._joined_models.append(arg)
|
|
return self
|
|
|
|
def select_related(self, *fields):
|
|
converted = [pascal_to_snake(f) for f in fields]
|
|
self._qs = self._qs.select_related(*converted)
|
|
return self
|
|
|
|
def prefetch_related(self, *fields):
|
|
converted = [pascal_to_snake(f) for f in fields]
|
|
self._qs = self._qs.prefetch_related(*converted)
|
|
return self
|
|
|
|
def annotate(self, **kwargs):
|
|
self._qs = self._qs.annotate(**kwargs)
|
|
return self
|
|
|
|
def aggregate(self, **kwargs):
|
|
return self._qs.aggregate(**kwargs)
|
|
|
|
def values(self, *fields):
|
|
from gvsds.fsr import _AggregateExpr
|
|
annotations = {}
|
|
converted = []
|
|
for i, f in enumerate(fields):
|
|
if isinstance(f, _AggregateExpr):
|
|
annotations[f'val{i}'] = f.aggregate
|
|
converted.append(f'val{i}')
|
|
elif isinstance(f, (models.Aggregate, models.Expression)):
|
|
annotations[f'val{i}'] = f
|
|
converted.append(f'val{i}')
|
|
elif isinstance(f, str):
|
|
converted.append(pascal_to_snake(f) if not f.startswith('-') else f)
|
|
else:
|
|
converted.append(str(f))
|
|
qs = self._qs
|
|
if annotations:
|
|
qs = qs.annotate(**annotations)
|
|
return _ValuesResult(qs, fields)
|
|
|
|
def with_entities(self, *fields):
|
|
return self.values(*fields)
|
|
|
|
def params(self, **kwargs):
|
|
return self
|
|
|
|
def values_list(self, *fields, flat=False):
|
|
converted = [pascal_to_snake(f) for f in fields]
|
|
return self._qs.values_list(*converted, flat=flat)
|
|
|
|
def update(self, **kwargs):
|
|
django_kwargs = self._convert_kwargs(kwargs)
|
|
result = self._qs.update(**django_kwargs)
|
|
_invalidate_cache()
|
|
return result
|
|
|
|
def delete(self):
|
|
result = self._qs.delete()
|
|
_invalidate_cache()
|
|
return result
|
|
|
|
def exists(self):
|
|
return self._qs.exists()
|
|
|
|
def distinct(self):
|
|
self._qs = self._qs.distinct()
|
|
return self
|
|
|
|
def paginate(self, page=1, per_page=20, error_out=False):
|
|
from django.core.paginator import Paginator
|
|
if not self._qs.ordered:
|
|
meta = getattr(self._model, '_meta', None)
|
|
if meta and meta.pk:
|
|
self._qs = self._qs.order_by(meta.pk.name)
|
|
paginator = Paginator(self._qs, per_page)
|
|
try:
|
|
page_obj = paginator.page(page)
|
|
except Exception:
|
|
page_obj = paginator.page(1)
|
|
page_obj.total = paginator.count
|
|
page_obj.items = page_obj.object_list
|
|
page_obj.pages = paginator.num_pages
|
|
page_obj.page = page
|
|
page_obj.per_page = per_page
|
|
page_obj.has_next = page_obj.has_next()
|
|
page_obj.has_prev = page_obj.has_previous()
|
|
page_obj.next_num = page_obj.next_page_number() if page_obj.has_next else None
|
|
page_obj.prev_num = page_obj.previous_page_number() if page_obj.has_prev else None
|
|
return page_obj
|
|
|
|
def no_cache(self):
|
|
self._cache_enabled = False
|
|
return self
|
|
|
|
def using(self, alias):
|
|
self._qs = self._qs.using(alias)
|
|
return self
|
|
|
|
def __iter__(self):
|
|
return iter(self.all())
|
|
|
|
def __len__(self):
|
|
return self.count()
|
|
|
|
def __bool__(self):
|
|
return self.exists()
|
|
|
|
def __repr__(self):
|
|
return f'FluentQuery({self._model.__name__})'
|
|
|
|
|
|
class Session:
|
|
def __init__(self):
|
|
self._pending_adds = []
|
|
self._pending_deletes = []
|
|
self._tracked = set()
|
|
|
|
def add(self, obj):
|
|
self._pending_adds.append(obj)
|
|
|
|
def delete(self, obj):
|
|
self._pending_deletes.append(obj)
|
|
self._tracked.discard(obj)
|
|
|
|
def track(self, obj):
|
|
if isinstance(obj, models.Model):
|
|
self._tracked.add(obj)
|
|
|
|
def commit(self):
|
|
delete_pks = {(obj.__class__, obj.pk) for obj in self._pending_deletes if obj.pk is not None}
|
|
with transaction.atomic():
|
|
for obj in self._pending_deletes:
|
|
obj.delete()
|
|
for obj in self._pending_adds:
|
|
obj.save()
|
|
for obj in self._tracked:
|
|
if obj.pk is not None and (obj.__class__, obj.pk) in delete_pks:
|
|
continue
|
|
try:
|
|
obj.save()
|
|
except Exception:
|
|
pass
|
|
self._pending_adds.clear()
|
|
self._pending_deletes.clear()
|
|
self._tracked.clear()
|
|
_invalidate_cache()
|
|
|
|
def rollback(self):
|
|
self._pending_adds.clear()
|
|
self._pending_deletes.clear()
|
|
self._tracked.clear()
|
|
_invalidate_cache()
|
|
|
|
def query(self, *args):
|
|
from gvsds.fsr import _AggregateExpr
|
|
model = None
|
|
aggregates = []
|
|
for a in args:
|
|
if isinstance(a, type) and issubclass(a, models.Model):
|
|
model = a
|
|
elif isinstance(a, _AggregateExpr):
|
|
if a.model:
|
|
model = a.model
|
|
aggregates.append(a.aggregate)
|
|
elif isinstance(a, (models.Aggregate, models.Expression)):
|
|
aggregates.append(a)
|
|
if not aggregates and model:
|
|
return FluentQuery(model)
|
|
if aggregates and model:
|
|
fq = FluentQuery(model)
|
|
return fq.values(*args)
|
|
return None
|
|
|
|
def flush(self):
|
|
self.commit()
|
|
|
|
def execute(self, stmt):
|
|
from PDAPI_django import _DeleteStmt
|
|
if isinstance(stmt, _DeleteStmt):
|
|
stmt.execute()
|
|
elif isinstance(stmt, str):
|
|
from django.db import connection
|
|
with connection.cursor() as cursor:
|
|
cursor.execute(stmt)
|
|
|
|
|
|
class DB:
|
|
@property
|
|
def session(self):
|
|
return _get_session()
|
|
|
|
@staticmethod
|
|
def or_(*args):
|
|
result = None
|
|
for q in args:
|
|
if isinstance(q, Q):
|
|
if result is None:
|
|
result = q
|
|
else:
|
|
result |= q
|
|
return result if result is not None else Q()
|
|
|
|
@staticmethod
|
|
def and_(*args):
|
|
result = Q()
|
|
for q in args:
|
|
if isinstance(q, Q):
|
|
result &= q
|
|
return result
|
|
|
|
|
|
db = DB()
|
|
|
|
|
|
def FQ(model_class):
|
|
return FluentQuery(model_class)
|