|
| 1 | +from django.db import NotSupportedError |
| 2 | +from django.db.models.fields.json import ( |
| 3 | + ContainedBy, |
| 4 | + DataContains, |
| 5 | + HasAnyKeys, |
| 6 | + HasKey, |
| 7 | + HasKeyLookup, |
| 8 | + HasKeys, |
| 9 | + JSONExact, |
| 10 | + KeyTransform, |
| 11 | + KeyTransformIn, |
| 12 | + KeyTransformIsNull, |
| 13 | + KeyTransformNumericLookupMixin, |
| 14 | +) |
| 15 | + |
| 16 | +from ..lookups import builtin_lookup |
| 17 | +from ..query_utils import process_lhs, process_rhs |
| 18 | + |
| 19 | + |
| 20 | +def contained_by(self, compiler, connection): # noqa: ARG001 |
| 21 | + raise NotSupportedError("contained_by lookup is not supported on this database backend.") |
| 22 | + |
| 23 | + |
| 24 | +def data_contains(self, compiler, connection): # noqa: ARG001 |
| 25 | + raise NotSupportedError("contains lookup is not supported on this database backend.") |
| 26 | + |
| 27 | + |
| 28 | +def _has_key_predicate(path, root_column, negated=False): |
| 29 | + """Return MQL to check for the existence of `path`.""" |
| 30 | + result = { |
| 31 | + "$and": [ |
| 32 | + # The path must exist (i.e. not be "missing"). |
| 33 | + {"$ne": [{"$type": path}, "missing"]}, |
| 34 | + # If the JSONField value is None, an additional check for not null |
| 35 | + # is needed since $type returns null instead of "missing". |
| 36 | + {"$ne": [root_column, None]}, |
| 37 | + ] |
| 38 | + } |
| 39 | + if negated: |
| 40 | + result = {"$not": result} |
| 41 | + return result |
| 42 | + |
| 43 | + |
| 44 | +def has_key_lookup(self, compiler, connection): |
| 45 | + """Return MQL to check for the existence of a key.""" |
| 46 | + rhs = self.rhs |
| 47 | + lhs = process_lhs(self, compiler, connection) |
| 48 | + if not isinstance(rhs, list | tuple): |
| 49 | + rhs = [rhs] |
| 50 | + paths = [] |
| 51 | + # Transform any "raw" keys into KeyTransforms to allow consistent handling |
| 52 | + # in the code that follows. |
| 53 | + for key in rhs: |
| 54 | + rhs_json_path = key if isinstance(key, KeyTransform) else KeyTransform(key, self.lhs) |
| 55 | + paths.append(rhs_json_path.as_mql(compiler, connection)) |
| 56 | + keys = [] |
| 57 | + for path in paths: |
| 58 | + keys.append(_has_key_predicate(path, lhs)) |
| 59 | + if self.mongo_operator is None: |
| 60 | + return keys[0] |
| 61 | + return {self.mongo_operator: keys} |
| 62 | + |
| 63 | + |
| 64 | +_process_rhs = JSONExact.process_rhs |
| 65 | + |
| 66 | + |
| 67 | +def json_exact_process_rhs(self, compiler, connection): |
| 68 | + """Skip JSONExact.process_rhs()'s conversion of None to "null".""" |
| 69 | + return ( |
| 70 | + super(JSONExact, self).process_rhs(compiler, connection) |
| 71 | + if connection.vendor == "mongodb" |
| 72 | + else _process_rhs(self, compiler, connection) |
| 73 | + ) |
| 74 | + |
| 75 | + |
| 76 | +def key_transform(self, compiler, connection): |
| 77 | + """ |
| 78 | + Return MQL for this KeyTransform (JSON path). |
| 79 | +
|
| 80 | + JSON paths cannot always be represented simply as $var.key1.key2.key3 due |
| 81 | + to possible array types. Therefore, indexing arrays requires the use of |
| 82 | + `arrayElemAt`. Additionally, $cond is necessary to verify the type before |
| 83 | + performing the operation. |
| 84 | + """ |
| 85 | + key_transforms = [self.key_name] |
| 86 | + previous = self.lhs |
| 87 | + # Collect all key transforms in order. |
| 88 | + while isinstance(previous, KeyTransform): |
| 89 | + key_transforms.insert(0, previous.key_name) |
| 90 | + previous = previous.lhs |
| 91 | + lhs_mql = previous.as_mql(compiler, connection) |
| 92 | + result = lhs_mql |
| 93 | + # Build the MQL path using the collected key transforms. |
| 94 | + for key in key_transforms: |
| 95 | + get_field = {"$getField": {"input": result, "field": key}} |
| 96 | + # Handle array indexing if the key is a digit. If key is something |
| 97 | + # like '001', it's not an array index despite isdigit() returning True. |
| 98 | + if key.isdigit() and str(int(key)) == key: |
| 99 | + result = { |
| 100 | + "$cond": { |
| 101 | + "if": {"$isArray": result}, |
| 102 | + "then": {"$arrayElemAt": [result, int(key)]}, |
| 103 | + "else": get_field, |
| 104 | + } |
| 105 | + } |
| 106 | + else: |
| 107 | + result = get_field |
| 108 | + return result |
| 109 | + |
| 110 | + |
| 111 | +def key_transform_in(self, compiler, connection): |
| 112 | + """ |
| 113 | + Return MQL to check if a JSON path exists and that its values are in the |
| 114 | + set of specified values (rhs). |
| 115 | + """ |
| 116 | + lhs_mql = process_lhs(self, compiler, connection) |
| 117 | + # Traverse to the root column. |
| 118 | + previous = self.lhs |
| 119 | + while isinstance(previous, KeyTransform): |
| 120 | + previous = previous.lhs |
| 121 | + root_column = previous.as_mql(compiler, connection) |
| 122 | + value = process_rhs(self, compiler, connection) |
| 123 | + # Construct the expression to check if lhs_mql values are in rhs values. |
| 124 | + expr = connection.mongo_operators[self.lookup_name](lhs_mql, value) |
| 125 | + return {"$and": [_has_key_predicate(lhs_mql, root_column), expr]} |
| 126 | + |
| 127 | + |
| 128 | +def key_transform_is_null(self, compiler, connection): |
| 129 | + """ |
| 130 | + Return MQL to check the nullability of a key. |
| 131 | +
|
| 132 | + If `isnull=True`, the query matches objects where the key is missing or the |
| 133 | + root column is null. If `isnull=False`, the query negates the result to |
| 134 | + match objects where the key exists. |
| 135 | +
|
| 136 | + Reference: https://code.djangoproject.com/ticket/32252 |
| 137 | + """ |
| 138 | + lhs_mql = process_lhs(self, compiler, connection) |
| 139 | + rhs_mql = process_rhs(self, compiler, connection) |
| 140 | + # Get the root column. |
| 141 | + previous = self.lhs |
| 142 | + while isinstance(previous, KeyTransform): |
| 143 | + previous = previous.lhs |
| 144 | + root_column = previous.as_mql(compiler, connection) |
| 145 | + return _has_key_predicate(lhs_mql, root_column, negated=rhs_mql) |
| 146 | + |
| 147 | + |
| 148 | +def key_transform_numeric_lookup_mixin(self, compiler, connection): |
| 149 | + """ |
| 150 | + Return MQL to check if the field exists (i.e., is not "missing" or "null") |
| 151 | + and that the field matches the given numeric lookup expression. |
| 152 | + """ |
| 153 | + expr = builtin_lookup(self, compiler, connection) |
| 154 | + lhs = process_lhs(self, compiler, connection) |
| 155 | + # Check if the type of lhs is not "missing" or "null". |
| 156 | + not_missing_or_null = {"$not": {"$in": [{"$type": lhs}, ["missing", "null"]]}} |
| 157 | + return {"$and": [expr, not_missing_or_null]} |
| 158 | + |
| 159 | + |
| 160 | +def register_json_field(): |
| 161 | + ContainedBy.as_mql = contained_by |
| 162 | + DataContains.as_mql = data_contains |
| 163 | + HasAnyKeys.mongo_operator = "$or" |
| 164 | + HasKey.mongo_operator = None |
| 165 | + HasKeyLookup.as_mql = has_key_lookup |
| 166 | + HasKeys.mongo_operator = "$and" |
| 167 | + JSONExact.process_rhs = json_exact_process_rhs |
| 168 | + KeyTransform.as_mql = key_transform |
| 169 | + KeyTransformIn.as_mql = key_transform_in |
| 170 | + KeyTransformIsNull.as_mql = key_transform_is_null |
| 171 | + KeyTransformNumericLookupMixin.as_mql = key_transform_numeric_lookup_mixin |
0 commit comments