Initial commit.

Author: WaVEV

django_mongodb/compiler.py

@@ -114,7 +114,8 @@ def check_query(self):
             isinstance(a, Aggregate) and not isinstance(a, Count)
             for a in self.query.annotations.values()
         ):
-            raise NotSupportedError("QuerySet.aggregate() isn't supported on MongoDB.")
+            # raise NotSupportedError("QuerySet.aggregate() isn't supported on MongoDB.")
+            pass
 
     def get_count(self, check_exists=False):
         """
@@ -141,8 +142,10 @@ def get_count(self, check_exists=False):
     def build_query(self, columns=None):
         """Check if the query is supported and prepare a MongoQuery."""
         self.check_query()
-        query = self.query_class(self, columns)
+        query = self.query_class(self)
+        query.project_fields = self._get_project_fields(columns)
         query.lookup_pipeline = self.get_lookup_pipeline()
+        query.annotation_stage = self._get_group_pipeline()
         try:
             query.mongo_query = {"$expr": self.query.where.as_mql(self, self.connection)}
         except FullResultSet:
@@ -247,6 +250,73 @@ def get_lookup_pipeline(self):
             result += self.query.alias_map[alias].as_mql(self, self.connection)
         return result
 
+    def _get_group_pipeline(self):
+        pipeline = None
+        if any(isinstance(a, Aggregate) for a in self.query.annotations.values()):
+            result = {}
+            # self.get_group_by(self.select, [])
+            for alias, annotation in self.query.annotation_select.items():
+                value = annotation.as_mql(self, self.connection)
+                if isinstance(value, list):
+                    value = value[0]
+                result[alias] = value
+
+            expressions = set()
+            for expr, *_ in self.select:
+                expressions |= set(expr.get_group_by_cols())
+            order_by = self.get_order_by()
+            for expr, (_, _, is_ref) in order_by:
+                # Skip references to the SELECT clause, as all expressions in
+                # the SELECT clause are already part of the GROUP BY.
+                if not is_ref:
+                    expressions |= set(expr.get_group_by_cols())
+            having_group_by = self.having.get_group_by_cols() if self.having else ()
+            for expr in having_group_by:
+                expressions.add(expr)
+
+            ids = (
+                None
+                if not expressions
+                else {col.target.column: col.as_mql(self, self.connection) for col in expressions}
+            )
+            result["_id"] = ids
+
+            pipeline = [{"$group": result}]
+            if ids:
+                pipeline.append(
+                    {"$addFields": {key: f"$_id.{value[1:]}" for key, value in ids.items()}}
+                )
+                if "_id" not in ids:
+                    pipeline.append({"$unSet": "$_id"})
+
+        return pipeline
+
+    def _get_project_fields(self, columns=None):
+        fields = {}
+        for name, expr in columns or []:
+            try:
+                column = name if isinstance(expr, Aggregate) else expr.target.column
+            except AttributeError:
+                # Generate the MQL for an annotation.
+                try:
+                    fields[name] = expr.as_mql(self, self.connection)
+                except EmptyResultSet:
+                    fields[name] = Value(False).as_mql(self, self.connection)
+                except FullResultSet:
+                    fields[name] = Value(True).as_mql(self, self.connection)
+            else:
+                # If name != column, then this is an annotatation referencing
+                # another column.
+                fields[name] = 1 if name == column else f"${column}"
+
+        if fields:
+            # Add related fields.
+            for alias in self.query.alias_map:
+                if self.query.alias_refcount[alias] and self.collection_name != alias:
+                    fields[alias] = 1
+
+        return fields
+
 
 class SQLInsertCompiler(SQLCompiler):
     def execute_sql(self, returning_fields=None):

django_mongodb/functions.py

@@ -1,4 +1,5 @@
 from django.db import NotSupportedError
+from django.db.models.aggregates import Aggregate, Avg, Count, Max, Min, StdDev, Sum, Variance
 from django.db.models.expressions import Func
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
 from django.db.models.functions.datetime import (
@@ -33,6 +34,17 @@
 
 from .query_utils import process_lhs
 
+MONGO_AGGREGATION = {
+    Avg: "avg",
+    Count: "sum",
+    Max: "max",
+    Min: "min",
+    StdDev: "stddev",
+    Sum: "sum",
+    Variance: "variance",
+}
+
+
 MONGO_OPERATORS = {
     Ceil: "ceil",
     Coalesce: "ifNull",
@@ -57,6 +69,10 @@
 }
 
 
+def aggregate(self, compiler, connection):  # noqa: ARG001
+    pass
+
+
 def cast(self, compiler, connection):
     output_type = connection.data_types[self.output_field.get_internal_type()]
     lhs_mql = process_lhs(self, compiler, connection)[0]
@@ -187,6 +203,7 @@ def trunc(self, compiler, connection):
 
 
 def register_functions():
+    Aggregate.as_mql = aggregate
     Cast.as_mql = cast
     Concat.as_mql = concat
     ConcatPair.as_mql = concat_pair

django_mongodb/query.py

@@ -3,7 +3,7 @@
 
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError
-from django.db.models.expressions import Case, Value, When
+from django.db.models.expressions import Case, When
 from django.db.models.functions import Mod
 from django.db.models.lookups import Exact
 from django.db.models.sql.constants import INNER
@@ -37,18 +37,19 @@ class MongoQuery:
     built by Django to a "representation" more suitable for MongoDB.
     """
 
-    def __init__(self, compiler, columns):
+    def __init__(self, compiler):
         self.compiler = compiler
         self.connection = compiler.connection
         self.ops = compiler.connection.ops
         self.query = compiler.query
-        self.columns = columns
         self._negated = False
         self.ordering = []
         self.collection = self.compiler.get_collection()
         self.collection_name = self.compiler.collection_name
         self.mongo_query = getattr(compiler.query, "raw_query", {})
         self.lookup_pipeline = None
+        self.annotation_stage = None
+        self.project_fields = None
 
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"
@@ -103,35 +104,16 @@ def get_cursor(self, count=False, limit=None, skip=None):
         """
         if self.query.low_mark == self.query.high_mark:
             return []
-        fields = {}
-        for name, expr in self.columns or []:
-            try:
-                column = expr.target.column
-            except AttributeError:
-                # Generate the MQL for an annotation.
-                try:
-                    fields[name] = expr.as_mql(self.compiler, self.connection)
-                except EmptyResultSet:
-                    fields[name] = Value(False).as_mql(self.compiler, self.connection)
-                except FullResultSet:
-                    fields[name] = Value(True).as_mql(self.compiler, self.connection)
-            else:
-                # If name != column, then this is an annotatation referencing
-                # another column.
-                fields[name] = 1 if name == column else f"${column}"
-        if fields:
-            # Add related fields.
-            for alias in self.query.alias_map:
-                if self.query.alias_refcount[alias] and self.collection_name != alias:
-                    fields[alias] = 1
         # Construct the query pipeline.
         pipeline = []
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
         if self.mongo_query:
             pipeline.append({"$match": self.mongo_query})
-        if fields:
-            pipeline.append({"$project": fields})
+        if self.annotation_stage:
+            pipeline.extend(self.annotation_stage)
+        if self.project_fields:
+            pipeline.append({"$project": self.project_fields})
         if self.ordering:
             pipeline.append({"$sort": dict(self.ordering)})
         if skip is not None:

django_mongodb/query_utils.py

@@ -1,4 +1,5 @@
 from django.core.exceptions import FullResultSet
+from django.db.models.aggregates import Aggregate
 from django.db.models.expressions import Value
 
 
@@ -15,6 +16,8 @@ def process_lhs(node, compiler, connection):
                 result.append(expr.as_mql(compiler, connection))
             except FullResultSet:
                 result.append(Value(True).as_mql(compiler, connection))
+        if isinstance(node, Aggregate):
+            return result[0]
         return result
     # node is a Transform with just one source expression, aliased as "lhs".
     if is_direct_value(node.lhs):