132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
from django.contrib.postgres.fields import ArrayField
|
|
from django.contrib.postgres.fields.ranges import RangeField
|
|
from django.db.models import CharField, JSONField, Lookup
|
|
from django.db.models.fields.json import KeyTextTransform
|
|
|
|
from .fields import CachedValueField, ChoiceSetField
|
|
|
|
__all__ = (
|
|
'ChoiceValueLookup',
|
|
'Empty',
|
|
'JSONEmpty',
|
|
'NetContainsOrEquals',
|
|
'NetHost',
|
|
'RangeContains',
|
|
)
|
|
|
|
|
|
class RangeContains(Lookup):
|
|
"""
|
|
Filter ArrayField(RangeField) columns where ANY element-range contains the scalar RHS.
|
|
|
|
Usage (ORM):
|
|
Model.objects.filter(<range_array_field>__range_contains=<scalar>)
|
|
|
|
Works with int4range[], int8range[], daterange[], tstzrange[], etc.
|
|
"""
|
|
|
|
lookup_name = 'range_contains'
|
|
|
|
def as_sql(self, compiler, connection):
|
|
# Compile LHS (the array-of-ranges column/expression) and RHS (scalar)
|
|
lhs, lhs_params = self.process_lhs(compiler, connection)
|
|
rhs, rhs_params = self.process_rhs(compiler, connection)
|
|
|
|
# Guard: only allow ArrayField whose base_field is a PostgreSQL RangeField
|
|
field = getattr(self.lhs, 'output_field', None)
|
|
if not (isinstance(field, ArrayField) and isinstance(field.base_field, RangeField)):
|
|
raise TypeError('range_contains is only valid for ArrayField(RangeField) columns')
|
|
|
|
# Range-contains-element using EXISTS + UNNEST keeps the range on the LHS: r @> value
|
|
sql = f"EXISTS (SELECT 1 FROM unnest({lhs}) AS r WHERE r @> {rhs})"
|
|
params = lhs_params + rhs_params
|
|
return sql, params
|
|
|
|
|
|
class ChoiceValueLookup(Lookup):
|
|
"""
|
|
Match rows where any [value, label] pair in a ChoiceSetField has the given value.
|
|
|
|
Compares the RHS against the first element (the value) of each pair.
|
|
"""
|
|
lookup_name = 'choice_value'
|
|
prepare_rhs = False
|
|
|
|
def as_sql(self, compiler, connection):
|
|
lhs, lhs_params = self.process_lhs(compiler, connection)
|
|
rhs, rhs_params = self.process_rhs(compiler, connection)
|
|
# Slice the value column of the two-dimensional array and match any element
|
|
return f'{rhs} = ANY({lhs}[:][1:1])', [*rhs_params, *lhs_params]
|
|
|
|
|
|
class Empty(Lookup):
|
|
"""
|
|
Filter on whether a string is empty.
|
|
"""
|
|
lookup_name = 'empty'
|
|
prepare_rhs = False
|
|
|
|
def as_sql(self, compiler, connection):
|
|
sql, params = compiler.compile(self.lhs)
|
|
if self.rhs:
|
|
return f"CAST(LENGTH({sql}) AS BOOLEAN) IS NOT TRUE", params
|
|
return f"CAST(LENGTH({sql}) AS BOOLEAN) IS TRUE", params
|
|
|
|
|
|
class JSONEmpty(Lookup):
|
|
"""
|
|
Support "empty" lookups for JSONField keys.
|
|
|
|
A key is considered empty if it is "", null, or does not exist.
|
|
"""
|
|
lookup_name = 'empty'
|
|
|
|
def as_sql(self, compiler, connection):
|
|
# self.lhs.lhs is the parent expression (could be a JSONField or another KeyTransform)
|
|
# Rebuild the expression using KeyTextTransform to guarantee ->> (text)
|
|
text_expr = KeyTextTransform(self.lhs.key_name, self.lhs.lhs)
|
|
lhs_sql, lhs_params = compiler.compile(text_expr)
|
|
|
|
value = self.rhs
|
|
if value not in (True, False):
|
|
raise ValueError("The 'empty' lookup only accepts True or False.")
|
|
|
|
condition = '' if value else 'NOT '
|
|
sql = f"(NULLIF({lhs_sql}, '') IS {condition}NULL)"
|
|
|
|
return sql, lhs_params
|
|
|
|
|
|
class NetHost(Lookup):
|
|
"""
|
|
Similar to ipam.lookups.NetHost, but casts the field to INET.
|
|
"""
|
|
lookup_name = 'net_host'
|
|
|
|
def as_sql(self, qn, connection):
|
|
lhs, lhs_params = self.process_lhs(qn, connection)
|
|
rhs, rhs_params = self.process_rhs(qn, connection)
|
|
params = lhs_params + rhs_params
|
|
return f'HOST(CAST({lhs} AS INET)) = HOST({rhs})', params
|
|
|
|
|
|
class NetContainsOrEquals(Lookup):
|
|
"""
|
|
Similar to ipam.lookups.NetContainsOrEquals, but casts the field to INET.
|
|
"""
|
|
lookup_name = 'net_contains_or_equals'
|
|
|
|
def as_sql(self, qn, connection):
|
|
lhs, lhs_params = self.process_lhs(qn, connection)
|
|
rhs, rhs_params = self.process_rhs(qn, connection)
|
|
params = lhs_params + rhs_params
|
|
return f'CAST({lhs} AS INET) >>= {rhs}', params
|
|
|
|
|
|
ArrayField.register_lookup(RangeContains)
|
|
ChoiceSetField.register_lookup(ChoiceValueLookup)
|
|
CharField.register_lookup(Empty)
|
|
JSONField.register_lookup(JSONEmpty)
|
|
CachedValueField.register_lookup(NetHost)
|
|
CachedValueField.register_lookup(NetContainsOrEquals)
|