add aggregation function to keytransforms.

Author: WaVEV

django_mongodb/base.py

@@ -97,6 +97,7 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         "gte": lambda a, b: {"$gte": [a, b]},
         "lt": lambda a, b: {"$lt": [a, b]},
         "lte": lambda a, b: {"$lte": [a, b]},
+        "in": lambda a, b: {"$in": [a, b]},
     }
 
     display_name = "MongoDB"

django_mongodb/compiler.py

@@ -153,9 +153,18 @@ def get_columns(self):
         columns = (
             self.get_default_columns(select_mask) if self.query.default_cols else self.query.select
         )
-        return tuple((column.target.column, column) for column in columns) + tuple(
-            self.query.annotations.items()
-        )
+        annotations_cnt = 0
+        result = []
+        for column in columns:
+            if hasattr(column, "target"):
+                target = column.target.column
+            else:
+                annotations_cnt += 1
+                target = f"annotation_{annotations_cnt}"
+
+            result.append((target, column))
+
+        return tuple(result) + tuple(self.query.annotations.items())
 
     def _get_ordering(self):
         """

django_mongodb/fields.py

@@ -64,16 +64,9 @@ def from_db_value(self, value, expression, connection):
 
 def json_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 []
-    else:
-        result_value = []
-        for ind, elem in enumerate(node.rhs):
-            item = f"${value[ind]}" if isinstance(elem, KeyTransform) else value[ind]
-            result_value.append(item)
-        value = result_value
 
     return value
 
@@ -97,35 +90,31 @@ def contained_by(self, compiler, connection):  # noqa: ARG001
 
 
 def json_exact(self, compiler, connection):
-    lhs_mql = process_lhs(self, compiler, connection)
+    lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
     rhs_mql = json_process_rhs(self, compiler, connection)
     if rhs_mql == "null":
         return {"$or": [{lhs_mql: {"$eq": None}}, {lhs_mql: {"$exists": False}}]}
-        # return {lhs_mql: {"$eq": None, "$exists": True}}
-        # return key_transform_isnull(self, compiler, connection)
     return {lhs_mql: {"$eq": rhs_mql, "$exists": True}}
 
 
 def key_transform_isnull(self, compiler, connection):
-    lhs_mql = process_lhs(self, compiler, connection)
+    lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
     rhs_mql = json_process_rhs(self, compiler, connection)
-    # if rhs_mql is False:
-    #     return {lhs_mql: {"$neq": None}}
-    # return {"$or": [{lhs_mql: {"$eq": None}}, {lhs_mql: {"$exists": False}}]}
 
     # https://code.djangoproject.com/ticket/32252
     return {lhs_mql: {"$exists": not rhs_mql}}
 
 
 def key_transform_in(self, compiler, connection):
-    lhs_mql = process_lhs(self, compiler, connection)
+    lhs_mql = key_transform_agg(self.lhs, compiler, connection)
+    bare_lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
     value = json_process_rhs(self, compiler, connection)
-    rhs_mql = connection.operators[self.lookup_name](value)
-    return {"$expr": {lhs_mql: rhs_mql}}
+    expr = connection.mongo_aggregations[self.lookup_name](lhs_mql, value)
+    return {"$expr": expr, bare_lhs_mql: {"$exists": True}}
 
 
 def has_key_lookup(self, compiler, connection):
-    lhs = process_lhs(self, compiler, connection)
+    lhs = process_lhs(self, compiler, connection, bare_column_ref=True)
     rhs = self.rhs
     if not isinstance(rhs, (list | tuple)):
         rhs = [rhs]
@@ -148,10 +137,34 @@ def has_key_lookup(self, compiler, connection):
     return {self.mongo_operator: keys}
 
 
+def key_transform_agg(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)
+    result = f"{lhs_mql}"
+    for key in key_transforms:
+        get_field = {"$getField": {"input": result, "field": key}}
+        if key.isdigit():
+            result = {
+                "$cond": {
+                    "if": {"$isArray": result},
+                    "then": {"$arrayElemAt": [result, int(key)]},
+                    "else": get_field,
+                }
+            }
+        else:
+            result = get_field
+    return result
+
+
 def load_fields():
     JSONField.from_db_value = from_db_value
     DataContains.as_mql = data_contains
     KeyTransform.as_mql = key_transform
+    KeyTransform.as_mql_agg = key_transform_agg
     JSONExact.as_mql = json_exact
     ContainedBy.as_mql = contained_by
     HasKeyLookup.as_mql = has_key_lookup