JSONField: Haskey lookups.

Author: WaVEV

django_mongodb/features.py

@@ -7,6 +7,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     supports_ignore_conflicts = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/8
     supports_json_field = True
+    has_native_json_field = True
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/7
     supports_transactions = False
     uses_savepoints = False

django_mongodb/fields.py

@@ -2,11 +2,20 @@
 from django.core import exceptions
 from django.db.models import JSONField
 from django.db.models.fields import AutoField, Field
-from django.db.models.fields.json import DataContains, JSONExact, KeyTransform
+from django.db.models.fields.json import (
+    ContainedBy,
+    DataContains,
+    HasAnyKeys,
+    HasKey,
+    HasKeyLookup,
+    HasKeys,
+    JSONExact,
+    KeyTransform,
+)
 from django.utils.translation import gettext_lazy as _
 
 from .base import DatabaseWrapper
-from .query_utils import process_lhs, process_rhs
+from .query_utils import process_lhs
 
 
 class MongoAutoField(AutoField):
@@ -50,19 +59,35 @@ def from_db_value(self, value, expression, connection):
     )
 
 
+def process_rhs(node, compiler, connection):
+    _, value = node.process_rhs(compiler, connection)
+    lookup_name = node.lookup_name
+    if lookup_name not in ("in", "range"):
+        value = value[0] if len(value) > 0 else []
+
+    return value
+
+
 def key_transform(self, compiler, connection):
-    _, _, key_transforms = self.preprocess_lhs(compiler, connection)
-    lhs_mql = process_lhs(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_mql = previous.as_mql(compiler, connection)
     return ".".join([lhs_mql, *key_transforms])
 
 
 def data_contains(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)
     value = process_rhs(self, compiler, connection)
-    # rhs_mql = connection.operators[self.lookup_name](value)
     if not value:
-        return {lhs_mql: {"$ne": None}}
-    return {lhs_mql: {"$all": value}}
+        return {lhs_mql: {"$type": "array"}}
+    if isinstance(value, list):
+        return {"$and": [{lhs_mql: v} for v in value]}
+    if isinstance(value, dict):
+        return {"$and": [{f"{lhs_mql}.{k}": v} for k, v in value.items()]}
+    return {lhs_mql: value}
 
 
 def json_exact(self, compiler, connection):
@@ -71,8 +96,58 @@ def json_exact(self, compiler, connection):
     return {lhs_mql: {"$eq": rhs_mql}}
 
 
+def contained_by(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    if isinstance(value, list):
+        if not value:
+            return {lhs_mql: []}
+        return {"$and": [{"$or": [{lhs_mql: v}, {lhs_mql: {"$exists": False}}]} for v in value]}
+    if isinstance(value, dict):
+        if not value:
+            return {lhs_mql: {}}
+        return {
+            "$and": [
+                {"$or": [{f"{lhs_mql}.{k}": v}, {f"{lhs_mql}.{k}": {"$exists": False}}]}
+                for k, v in value.items()
+            ]
+        }
+    return {lhs_mql: value}
+
+
+def has_key_lookup(self, compiler, connection):
+    lhs = process_lhs(self, compiler, connection)
+    rhs = self.rhs
+    if not isinstance(rhs, (list | tuple)):
+        rhs = [rhs]
+    paths = []
+    for key in rhs:
+        if isinstance(key, KeyTransform):
+            *_, rhs_key_transforms = key.preprocess_lhs(compiler, connection)
+            rhs_key_transforms = ".".join(
+                rhs_key_transforms
+            )  # process_lhs(key, compiler, connection)
+        else:
+            rhs_key_transforms = str(key)
+        rhs_json_path = f"{lhs}.{rhs_key_transforms}"
+        paths.append(rhs_json_path)
+
+    keys = []
+    for path in paths:
+        keys.append({path: {"$exists": True}})
+    if self.mongo_operator is None:
+        assert len(keys) == 1
+        return keys[0]
+    return {self.mongo_operator: keys}
+
+
 def load_fields():
     JSONField.from_db_value = from_db_value
     DataContains.as_mql = data_contains
     KeyTransform.as_mql = key_transform
     JSONExact.as_mql = json_exact
+    ContainedBy.as_mql = contained_by
+    HasKeyLookup.as_mql = has_key_lookup
+    HasAnyKeys.mongo_operator = "$or"
+    HasKey.mongo_operator = None
+    HasKeys.mongo_operator = "$and"