zoo/zoo/zoo_data/lookups.py

116 lines
3.6 KiB
Python

# zoo - versatile objects management
# Copyright (C) 2016 Entr'ouvert
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import django
from django.db.models import Transform, TextField, DateField
from django.contrib.postgres.fields import jsonb
try:
from django.contrib.postgres.fields.jsonb import KeyTransform, KeyTransformTextLookupMixin
except ImportError:
# backport from Django 2.x
class KeyTransform(Transform):
operator = '->'
nested_operator = '#>'
def __init__(self, key_name, *args, **kwargs):
super(KeyTransform, self).__init__(*args, **kwargs)
self.key_name = key_name
def as_sql(self, compiler, connection):
key_transforms = [self.key_name]
previous = self.lhs
while isinstance(previous, KeyTransform):
key_transforms.insert(0, previous.key_name)
previous = previous.lhs
lhs, params = compiler.compile(previous)
if len(key_transforms) > 1:
return "(%s %s %%s)" % (lhs, self.nested_operator), [key_transforms] + params
try:
int(self.key_name)
except ValueError:
lookup = "'%s'" % self.key_name
else:
lookup = "%s" % self.key_name
return "(%s %s %s)" % (lhs, self.operator, lookup), params
jsonb.KeyTransform = KeyTransform
class KeyTextTransform(KeyTransform):
operator = '->>'
nested_operator = '#>>'
_output_field = TextField()
class KeyTransformTextLookupMixin(object):
"""
Mixin for combining with a lookup expecting a text lhs from a JSONField
key lookup. Make use of the ->> operator instead of casting key values to
text and performing the lookup on the resulting representation.
"""
def __init__(self, key_transform, *args, **kwargs):
assert isinstance(key_transform, KeyTransform)
key_text_transform = KeyTextTransform(
key_transform.key_name, *key_transform.source_expressions, **key_transform.extra
)
super(KeyTransformTextLookupMixin, self).__init__(key_text_transform, *args, **kwargs)
class Lower(Transform):
lookup_name = 'lower'
function = 'LOWER'
TextField.register_lookup(Lower)
class Unaccent(Transform):
lookup_name = 'unaccent'
function = 'immutable_unaccent'
TextField.register_lookup(Unaccent)
class Normalize(Transform):
lookup_name = 'normalize'
function = 'immutable_normalize'
TextField.register_lookup(Normalize)
class Date(Transform):
lookup_name = 'timestamp'
function = 'immutable_date'
_output_field = DateField()
TextField.register_lookup(Date)
class JSONUnaccent(KeyTransformTextLookupMixin, Unaccent):
pass
class JSONTimestamp(KeyTransformTextLookupMixin, Date):
pass
class JSONNormalize(KeyTransformTextLookupMixin, Normalize):
pass
KeyTransform.register_lookup(JSONUnaccent)
KeyTransform.register_lookup(JSONTimestamp)
KeyTransform.register_lookup(JSONNormalize)