add support for QuerySet.aggregate()

Also add support for Count() in QuerySet.annotate().

Author: WaVEV

.github/workflows/test-python.yml

@@ -67,6 +67,7 @@ jobs:
       - name: Run tests
         run: >
           python3 django_repo/tests/runtests.py --settings mongodb_settings -v 2
+          aggregation
           annotations
           auth_tests.test_models.UserManagerTestCase
           backends.base.test_base.DatabaseWrapperTests

README.md

@@ -109,7 +109,6 @@ Migrations for 'admin':
 ## Known issues and limitations
 
 - The following `QuerySet` methods aren't supported:
-  - `aggregate()`
   - `bulk_update()`
   - `dates()`
   - `datetimes()`

django_mongodb/__init__.py

@@ -6,12 +6,14 @@
 
 check_django_compatability()
 
+from .aggregates import register_aggregates  # noqa: E402
 from .expressions import register_expressions  # noqa: E402
 from .fields import register_fields  # noqa: E402
 from .functions import register_functions  # noqa: E402
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
 
+register_aggregates()
 register_expressions()
 register_fields()
 register_functions()

django_mongodb/aggregates.py

@@ -0,0 +1,85 @@
+from copy import deepcopy
+
+from django.db.models.aggregates import Aggregate, Count, StdDev, Variance
+from django.db.models.expressions import Case, Value, When
+from django.db.models.lookups import Exact
+from django.db.models.sql.where import WhereNode
+
+from .query_utils import process_lhs
+
+# Aggregates whose MongoDB aggregation name differ from Aggregate.function.lower().
+MONGO_AGGREGATIONS = {Count: "sum"}
+
+
+def aggregate(
+    self,
+    compiler,
+    connection,
+    operator=None,
+    resolve_inner_expression=False,
+    **extra_context,  # noqa: ARG001
+):
+    if self.filter:
+        node = self.copy()
+        node.filter = None
+        source_expressions = node.get_source_expressions()
+        condition = When(self.filter, then=source_expressions[0])
+        node.set_source_expressions([Case(condition)] + source_expressions[1:])
+    else:
+        node = self
+    lhs_mql = process_lhs(node, compiler, connection)
+    if resolve_inner_expression:
+        return lhs_mql
+    operator = operator or MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
+    return {f"${operator}": lhs_mql}
+
+
+def count(self, compiler, connection, resolve_inner_expression=False, **extra_context):  # noqa: ARG001
+    """
+    When resolve_inner_expression=True, return the MQL that resolves as a
+    value. This is used to count different elements, so the inner values are
+    returned to be pushed into a set.
+    """
+    if not self.distinct or resolve_inner_expression:
+        if self.filter:
+            node = self.copy()
+            node.filter = None
+            source_expressions = node.get_source_expressions()
+            filter_ = deepcopy(self.filter)
+            filter_.add(
+                WhereNode([Exact(source_expressions[0], Value(None))], negated=True),
+                filter_.default,
+            )
+            condition = When(filter_, then=Value(1))
+            node.set_source_expressions([Case(condition)] + source_expressions[1:])
+            inner_expression = process_lhs(node, compiler, connection)
+        else:
+            lhs_mql = process_lhs(self, compiler, connection)
+            null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
+            inner_expression = {
+                "$cond": {"if": null_cond, "then": None, "else": lhs_mql if self.distinct else 1}
+            }
+        if resolve_inner_expression:
+            return inner_expression
+        return {"$sum": inner_expression}
+    # If distinct=True or resolve_inner_expression=False, sum the size of the
+    # set.
+    lhs_mql = process_lhs(self, compiler, connection)
+    # None shouldn't be counted, so subtract 1 if it's present.
+    exits_null = {"$cond": {"if": {"$in": [{"$literal": None}, lhs_mql]}, "then": -1, "else": 0}}
+    return {"$add": [{"$size": lhs_mql}, exits_null]}
+
+
+def stddev_variance(self, compiler, connection, **extra_context):
+    if self.function.endswith("_SAMP"):
+        operator = "stdDevSamp"
+    elif self.function.endswith("_POP"):
+        operator = "stdDevPop"
+    return aggregate(self, compiler, connection, operator=operator, **extra_context)
+
+
+def register_aggregates():
+    Aggregate.as_mql = aggregate
+    Count.as_mql = count
+    StdDev.as_mql = stddev_variance
+    Variance.as_mql = stddev_variance