1import json
  2from collections.abc import Callable, Sequence
  3from typing import TYPE_CHECKING, Any
  4
  5from plain.postgres import expressions, lookups
  6from plain.postgres.constants import LOOKUP_SEP
  7from plain.postgres.dialect import adapt_json_value
  8from plain.postgres.fields import TextField
  9from plain.postgres.fields.base import NOT_PROVIDED
 10from plain.postgres.lookups import (
 11    FieldGetDbPrepValueMixin,
 12    Lookup,
 13    OperatorLookup,
 14    Transform,
 15)
 16
 17from plain import exceptions
 18
 19from .base import DefaultableField
 20
 21if TYPE_CHECKING:
 22    from plain.postgres.connection import DatabaseConnection
 23    from plain.postgres.sql.compiler import SQLCompiler
 24
 25__all__ = ["JSONField"]
 26
 27
 28class JSONField(DefaultableField):
 29    db_type_sql = "jsonb"
 30    empty_strings_allowed = False
 31
 32    def __init__(
 33        self,
 34        *,
 35        encoder: type[json.JSONEncoder] | None = None,
 36        decoder: type[json.JSONDecoder] | None = None,
 37        required: bool = True,
 38        allow_null: bool = False,
 39        default: Any = NOT_PROVIDED,
 40        validators: Sequence[Callable[..., Any]] = (),
 41    ):
 42        if encoder and not callable(encoder):
 43            raise ValueError("The encoder parameter must be a callable object.")
 44        if decoder and not callable(decoder):
 45            raise ValueError("The decoder parameter must be a callable object.")
 46        self.encoder = encoder
 47        self.decoder = decoder
 48        super().__init__(
 49            required=required,
 50            allow_null=allow_null,
 51            default=default,
 52            validators=validators,
 53        )
 54
 55    def deconstruct(self) -> tuple[str, str, list[Any], dict[str, Any]]:
 56        name, path, args, kwargs = super().deconstruct()
 57        if self.encoder is not None:
 58            kwargs["encoder"] = self.encoder
 59        if self.decoder is not None:
 60            kwargs["decoder"] = self.decoder
 61        return name, path, args, kwargs
 62
 63    def from_db_value(
 64        self, value: Any, expression: Any, connection: DatabaseConnection
 65    ) -> Any:
 66        if value is None:
 67            return value
 68        # KeyTransform may extract non-string values directly.
 69        if isinstance(expression, KeyTransform) and not isinstance(value, str):
 70            return value
 71        try:
 72            return json.loads(value, cls=self.decoder)
 73        except json.JSONDecodeError:
 74            return value
 75
 76    def get_db_prep_value(
 77        self, value: Any, connection: DatabaseConnection, prepared: bool = False
 78    ) -> Any:
 79        if isinstance(value, expressions.Value) and isinstance(
 80            value.output_field, JSONField
 81        ):
 82            value = value.value
 83        elif hasattr(value, "as_sql"):
 84            return value
 85        return self.adapt_json_db_value(value)
 86
 87    def adapt_json_db_value(self, value: Any) -> Any:
 88        """Adapt a plain JSON value for this field's column — the one step
 89        subclasses with a different column type (EncryptedJSONField) replace.
 90
 91        May receive None (get_db_prep_save short-circuits it, but other
 92        get_db_prep_value callers don't): here it adapts to jsonb null;
 93        EncryptedJSONField maps it to SQL NULL."""
 94        return adapt_json_value(value, self.encoder)
 95
 96    def get_db_prep_save(self, value: Any, connection: DatabaseConnection) -> Any:
 97        if value is None:
 98            return value
 99        return self.get_db_prep_value(value, connection)
100
101    def get_transform(self, name: str) -> Callable[..., Transform] | None:
102        # Always returns a transform (never None in practice)
103        transform = super().get_transform(name)
104        if transform:
105            return transform
106        return KeyTransformFactory(name)
107
108    def validate(self, value: Any, model_instance: Any) -> None:
109        super().validate(value, model_instance)
110        try:
111            json.dumps(value, cls=self.encoder)
112        except TypeError:
113            raise exceptions.ValidationError(
114                "Value must be valid JSON.",
115                code="invalid",
116                params={"value": value},
117            )
118
119
120class DataContains(FieldGetDbPrepValueMixin, OperatorLookup):
121    lookup_name = "contains"
122    # PostgreSQL @> operator checks if left JSON contains right JSON.
123    operator = "@>"
124
125
126class ContainedBy(FieldGetDbPrepValueMixin, OperatorLookup):
127    lookup_name = "contained_by"
128    # PostgreSQL <@ operator checks if left JSON is contained by right JSON.
129    operator = "<@"
130
131
132class HasKeyLookup(OperatorLookup):
133    """Lookup for checking if a JSON field has a key."""
134
135    logical_operator: str | None = None
136
137    def as_sql(
138        self, compiler: SQLCompiler, connection: DatabaseConnection
139    ) -> tuple[str, tuple[Any, ...]]:
140        # Handle KeyTransform on RHS by expanding it into LHS chain.
141        if isinstance(self.rhs, KeyTransform):
142            *_, rhs_key_transforms = self.rhs.preprocess_lhs(compiler, connection)
143            for key in rhs_key_transforms[:-1]:
144                self.lhs = KeyTransform(key, self.lhs)
145            self.rhs = rhs_key_transforms[-1]
146        return super().as_sql(compiler, connection)
147
148
149class HasKey(HasKeyLookup):
150    lookup_name = "has_key"
151    # PostgreSQL ? operator checks if key exists.
152    operator = "?"
153    prepare_rhs = False
154
155
156class HasKeys(HasKeyLookup):
157    lookup_name = "has_keys"
158    # PostgreSQL ?& operator checks if all keys exist.
159    operator = "?&"
160    logical_operator = " AND "
161
162    def get_prep_lookup(self) -> list[str]:
163        return [str(item) for item in self.rhs]
164
165
166class HasAnyKeys(HasKeys):
167    lookup_name = "has_any_keys"
168    # PostgreSQL ?| operator checks if any key exists.
169    operator = "?|"
170    logical_operator = " OR "
171
172
173class JSONExact(lookups.Exact):
174    can_use_none_as_rhs = True
175
176    def process_rhs(
177        self, compiler: SQLCompiler, connection: DatabaseConnection
178    ) -> tuple[str, list[Any]] | tuple[list[str], list[Any]]:
179        rhs, rhs_params = super().process_rhs(compiler, connection)
180        if isinstance(rhs, str):
181            # Treat None lookup values as null.
182            if rhs == "%s" and rhs_params == [None]:
183                rhs_params = ["null"]
184            return rhs, rhs_params
185        else:
186            return rhs, rhs_params
187
188
189class JSONIContains(lookups.IContains):
190    pass
191
192
193JSONField.register_lookup(DataContains)
194JSONField.register_lookup(ContainedBy)
195JSONField.register_lookup(HasKey)
196JSONField.register_lookup(HasKeys)
197JSONField.register_lookup(HasAnyKeys)
198JSONField.register_lookup(JSONExact)
199JSONField.register_lookup(JSONIContains)
200
201
202class KeyTransform(Transform):
203    # PostgreSQL -> operator extracts JSON object field as JSON.
204    operator = "->"
205    # PostgreSQL #> operator extracts nested JSON path as JSON.
206    nested_operator = "#>"
207
208    def __init__(self, key_name: str, *args: Any, **kwargs: Any):
209        super().__init__(*args, **kwargs)
210        self.key_name = str(key_name)
211
212    def preprocess_lhs(
213        self, compiler: SQLCompiler, connection: DatabaseConnection
214    ) -> tuple[str, tuple[Any, ...], list[str]]:
215        key_transforms = [self.key_name]
216        previous = self.lhs
217        while isinstance(previous, KeyTransform):
218            key_transforms.insert(0, previous.key_name)
219            previous = previous.lhs
220        lhs, params = compiler.compile(previous)
221        return lhs, params, key_transforms
222
223    def as_sql(
224        self,
225        compiler: SQLCompiler,
226        connection: DatabaseConnection,
227        function: str | None = None,
228        template: str | None = None,
229        arg_joiner: str | None = None,
230        **extra_context: Any,
231    ) -> tuple[str, list[Any]]:
232        lhs, params, key_transforms = self.preprocess_lhs(compiler, connection)
233        if len(key_transforms) > 1:
234            sql = f"({lhs} {self.nested_operator} %s)"
235            return sql, list(params) + [key_transforms]
236        try:
237            lookup = int(self.key_name)
238        except ValueError:
239            lookup = self.key_name
240        return f"({lhs} {self.operator} %s)", list(params) + [lookup]
241
242
243class KeyTextTransform(KeyTransform):
244    # PostgreSQL ->> operator extracts JSON object field as text.
245    operator = "->>"
246    # PostgreSQL #>> operator extracts nested JSON path as text.
247    nested_operator = "#>>"
248    output_field = TextField()
249
250    @classmethod
251    def from_lookup(cls, lookup: str) -> Any:
252        transform, *keys = lookup.split(LOOKUP_SEP)
253        if not keys:
254            raise ValueError("Lookup must contain key or index transforms.")
255        for key in keys:
256            transform = cls(key, transform)
257        return transform
258
259
260KT = KeyTextTransform.from_lookup
261
262
263class KeyTransformTextLookupMixin(Lookup):
264    """
265    Mixin for lookups expecting text LHS from a JSONField key lookup.
266    Uses the ->> operator to extract JSON values as text.
267    """
268
269    def __init__(self, key_transform: Any, *args: Any, **kwargs: Any):
270        if not isinstance(key_transform, KeyTransform):
271            raise TypeError(
272                "Transform should be an instance of KeyTransform in order to "
273                "use this lookup."
274            )
275        key_text_transform = KeyTextTransform(
276            key_transform.key_name,
277            *key_transform.source_expressions,
278            **key_transform.extra,
279        )
280        super().__init__(key_text_transform, *args, **kwargs)
281
282
283class KeyTransformIsNull(lookups.IsNull):
284    # key__isnull=False is the same as has_key='key'
285    pass
286
287
288class KeyTransformIn(lookups.In):
289    def resolve_expression_parameter(
290        self,
291        compiler: SQLCompiler,
292        connection: DatabaseConnection,
293        sql: str,
294        param: Any,
295    ) -> tuple[str, list[Any]]:
296        sql, params = super().resolve_expression_parameter(
297            compiler,
298            connection,
299            sql,
300            param,
301        )
302        return sql, list(params)
303
304
305class KeyTransformExact(JSONExact):
306    def process_rhs(
307        self, compiler: SQLCompiler, connection: DatabaseConnection
308    ) -> tuple[str, list[Any]] | tuple[list[str], list[Any]]:
309        if isinstance(self.rhs, KeyTransform):
310            return super(lookups.Exact, self).process_rhs(compiler, connection)
311        return super().process_rhs(compiler, connection)
312
313
314class KeyTransformIExact(KeyTransformTextLookupMixin, lookups.IExact):
315    pass
316
317
318class KeyTransformIContains(KeyTransformTextLookupMixin, lookups.IContains):
319    pass
320
321
322class KeyTransformStartsWith(KeyTransformTextLookupMixin, lookups.StartsWith):
323    pass
324
325
326class KeyTransformIStartsWith(KeyTransformTextLookupMixin, lookups.IStartsWith):
327    pass
328
329
330class KeyTransformEndsWith(KeyTransformTextLookupMixin, lookups.EndsWith):
331    pass
332
333
334class KeyTransformIEndsWith(KeyTransformTextLookupMixin, lookups.IEndsWith):
335    pass
336
337
338class KeyTransformRegex(KeyTransformTextLookupMixin, lookups.Regex):
339    pass
340
341
342class KeyTransformIRegex(KeyTransformTextLookupMixin, lookups.IRegex):
343    pass
344
345
346class KeyTransformLt(lookups.LessThan):
347    pass
348
349
350class KeyTransformLte(lookups.LessThanOrEqual):
351    pass
352
353
354class KeyTransformGt(lookups.GreaterThan):
355    pass
356
357
358class KeyTransformGte(lookups.GreaterThanOrEqual):
359    pass
360
361
362KeyTransform.register_lookup(KeyTransformIn)
363KeyTransform.register_lookup(KeyTransformExact)
364KeyTransform.register_lookup(KeyTransformIExact)
365KeyTransform.register_lookup(KeyTransformIsNull)
366KeyTransform.register_lookup(KeyTransformIContains)
367KeyTransform.register_lookup(KeyTransformStartsWith)
368KeyTransform.register_lookup(KeyTransformIStartsWith)
369KeyTransform.register_lookup(KeyTransformEndsWith)
370KeyTransform.register_lookup(KeyTransformIEndsWith)
371KeyTransform.register_lookup(KeyTransformRegex)
372KeyTransform.register_lookup(KeyTransformIRegex)
373
374KeyTransform.register_lookup(KeyTransformLt)
375KeyTransform.register_lookup(KeyTransformLte)
376KeyTransform.register_lookup(KeyTransformGt)
377KeyTransform.register_lookup(KeyTransformGte)
378
379
380class KeyTransformFactory:
381    def __init__(self, key_name: str):
382        self.key_name = key_name
383
384    def __call__(self, *args: Any, **kwargs: Any) -> KeyTransform:
385        return KeyTransform(self.key_name, *args, **kwargs)