Implementing aggregate compiler.

Author: WaVEV

django_mongodb/compiler.py

@@ -17,9 +17,10 @@ class SQLCompiler(compiler.SQLCompiler):
     """Base class for all Mongo compilers."""
 
     query_class = MongoQuery
+    _group_pipeline = None
 
-    def pre_sql_setup(self, *args, **kargs):
-        pre_setup = super().pre_sql_setup(*args, **kargs)
+    def pre_sql_setup(self, with_col_aliases=False):
+        pre_setup = super().pre_sql_setup(with_col_aliases=with_col_aliases)
         self.annotations = {}
         group = {}
         group_expressions = set()
@@ -30,17 +31,20 @@ def pre_sql_setup(self, *args, **kargs):
             else:
                 replacements = {}
                 for sub_expr in self._get_aggregate_expressions(expr):
-                    alias = f"__aggregation{aggregation_idx}"
+                    if sub_expr != expr:
+                        alias = f"__aggregation{aggregation_idx}"
+                        aggregation_idx += 1
+                    else:
+                        alias = target
+                    group_expressions |= set(sub_expr.get_group_by_cols())
                     group[alias] = sub_expr.as_mql(self, self.connection)
-                    aggregation_idx += 1
                     column_target = expr.output_field.__class__()
                     column_target.set_attributes_from_name(alias)
                     replacements[sub_expr] = Col(self.collection_name, column_target)
                 result_expr = expr.replace_expressions(replacements)
 
             self.annotations[target] = result_expr
         if group:
-            """
             order_by = self.get_order_by()
             for expr, (_, _, is_ref) in order_by:
                 # Skip references to the SELECT clause, as all expressions in
@@ -50,7 +54,8 @@ def pre_sql_setup(self, *args, **kargs):
             having_group_by = self.having.get_group_by_cols() if self.having else ()
             for expr in having_group_by:
                 group_expressions.add(expr)
-            """
+            if isinstance(self.query.group_by, tuple | list):
+                group_expressions |= set(self.query.group_by)
 
             ids = (
                 None
@@ -61,7 +66,6 @@ def pre_sql_setup(self, *args, **kargs):
                 }
             )
             group["_id"] = ids
-
             pipeline = [{"$group": group}]
             if ids:
                 pipeline.append(
@@ -79,8 +83,8 @@ def pre_sql_setup(self, *args, **kargs):
     def execute_sql(
         self, result_type=MULTI, chunked_fetch=False, chunk_size=GET_ITERATOR_CHUNK_SIZE
     ):
-        self.pre_sql_setup()
         # QuerySet.count()
+        self.pre_sql_setup()
         if self.query.annotations == {"__count": Count("*")}:
             return [self.get_count()]
 
@@ -337,17 +341,6 @@ def get_lookup_pipeline(self):
             result += self.query.alias_map[alias].as_mql(self, self.connection)
         return result
 
-    def _get_aggregate_expressions2(self, expr):
-        stack = [(None, expr)]
-        while stack:
-            parent, expr = stack.pop()
-            if isinstance(expr, Aggregate):
-                yield parent
-            elif hasattr(expr, "get_source_expressions"):
-                stack.extend(
-                    [((expr, idx), se) for idx, se in enumerate(expr.get_source_expressions())]
-                )
-
     def _get_aggregate_expressions(self, expr):
         stack = [expr]
         while stack:
@@ -494,4 +487,14 @@ def check_query(self):
 
 
 class SQLAggregateCompiler(SQLCompiler):
-    pass
+    def build_query(self, columns=None):
+        query = self.query_class(self)
+        query.project_fields = self.get_project_fields(tuple(self.query.annotation_select.items()))
+
+        compiler = self.query.inner_query.get_compiler(
+            self.using,
+            elide_empty=self.elide_empty,
+        )
+        compiler.pre_sql_setup(with_col_aliases=False)
+        query.sub_query = compiler.build_query()
+        return query

django_mongodb/expressions.py

@@ -76,7 +76,7 @@ def query(self, compiler, connection):  # noqa: ARG001
 
 
 def ref(self, compiler, connection):  # noqa: ARG001
-    return self.refs
+    return f"${self.refs}"
 
 
 def subquery(self, compiler, connection):  # noqa: ARG001

django_mongodb/query.py

@@ -47,6 +47,7 @@ def __init__(self, compiler):
         self.collection = self.compiler.get_collection()
         self.collection_name = self.compiler.collection_name
         self.mongo_query = getattr(compiler.query, "raw_query", {})
+        self.sub_query = None
         self.lookup_pipeline = None
         self.annotation_stage = None
         self.project_fields = None
@@ -88,19 +89,8 @@ def delete(self):
         options = self.connection.operation_flags.get("delete", {})
         return self.collection.delete_many(self.mongo_query, **options).deleted_count
 
-    @wrap_database_errors
-    def get_cursor(self, count=False, limit=None, skip=None):
-        """
-        Return a pymongo CommandCursor that can be iterated on to give the
-        results of the query.
-
-        If `count` is True, return a single document with the number of
-        documents that match the query.
-
-        Use `limit` or `skip` to override those options of the query.
-        """
-        # Construct the query pipeline.
-        pipeline = []
+    def get_pipeline(self, count=False, limit=None, skip=None):
+        pipeline = [] if self.sub_query is None else self.sub_query.get_pipeline()
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
         if self.mongo_query:
@@ -111,16 +101,34 @@ def get_cursor(self, count=False, limit=None, skip=None):
             pipeline.append({"$project": self.project_fields})
         if self.ordering:
             pipeline.append({"$sort": dict(self.ordering)})
+
         if skip is not None:
             pipeline.append({"$skip": skip})
         elif self.query.low_mark > 0:
             pipeline.append({"$skip": self.query.low_mark})
+
         if limit is not None:
             pipeline.append({"$limit": limit})
         elif self.query.high_mark is not None:
             pipeline.append({"$limit": self.query.high_mark - self.query.low_mark})
-        if count:
-            pipeline.append({"$group": {"_id": None, "__count": {"$sum": 1}}})
+
+        return pipeline
+
+    @wrap_database_errors
+    def get_cursor(self, count=False, limit=None, skip=None):
+        """
+        Return a pymongo CommandCursor that can be iterated on to give the
+        results of the query.
+
+        If `count` is True, return a single document with the number of
+        documents that match the query.
+
+        Use `limit` or `skip` to override those options of the query.
+        """
+        if self.query.low_mark == self.query.high_mark:
+            return []
+
+        pipeline = self.get_pipeline()
         return self.collection.aggregate(pipeline)