Add AggregateFilter, StringAgg.as_mql()

django/django@4b977a5

Author: timgraham

django_mongodb_backend/aggregates.py

@@ -1,4 +1,12 @@
-from django.db.models.aggregates import Aggregate, Count, StdDev, Variance
+from django.core.exceptions import EmptyResultSet, FullResultSet
+from django.db import NotSupportedError
+from django.db.models.aggregates import (
+    Aggregate,
+    Count,
+    StdDev,
+    StringAgg,
+    Variance,
+)
 from django.db.models.expressions import Case, Value, When
 from django.db.models.lookups import IsNull
 
@@ -9,15 +17,20 @@
 
 
 def aggregate(self, compiler, connection, operator=None, resolve_inner_expression=False):
-    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:]])
+    if self.filter is not None:
+        # Generate a CASE statement for this aggregate.
+        try:
+            lhs_mql = self.filter.as_mql(compiler, connection, as_expr=True)
+        except NotSupportedError:
+            source_expressions = self.get_source_expressions()
+            condition = Case(When(self.filter.condition, then=source_expressions[0]))
+            lhs_mql = condition.as_mql(compiler, connection, as_expr=True)
+        except FullResultSet:
+            lhs_mql = source_expressions[0].as_mql(compiler, connection, as_expr=True)
+        except EmptyResultSet:
+            lhs_mql = Value(None).as_mql(compiler, connection, as_expr=True)
     else:
-        node = self
-    lhs_mql = process_lhs(node, compiler, connection, as_expr=True)
+        lhs_mql = process_lhs(self, compiler, connection, as_expr=True)
     if resolve_inner_expression:
         return lhs_mql
     operator = operator or MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
@@ -32,14 +45,19 @@ def count(self, compiler, connection, resolve_inner_expression=False):
     """
     if not self.distinct or resolve_inner_expression:
         if self.filter:
-            node = self.copy()
-            node.filter = None
-            source_expressions = node.get_source_expressions()
-            condition = When(
-                self.filter, then=Case(When(IsNull(source_expressions[0], False), then=Value(1)))
-            )
-            node.set_source_expressions([Case(condition), *source_expressions[1:]])
-            inner_expression = process_lhs(node, compiler, connection, as_expr=True)
+            try:
+                inner_expression = self.filter.as_mql(compiler, connection, as_expr=True)
+            except NotSupportedError:
+                source_expressions = self.get_source_expressions()
+                condition = When(
+                    self.filter.condition,
+                    then=Case(When(IsNull(source_expressions[0], False), then=Value(1))),
+                )
+                inner_expression = Case(condition).as_mql(compiler, connection, as_expr=True)
+            except FullResultSet:
+                inner_expression = {"$sum": 1}
+            except EmptyResultSet:
+                inner_expression = {"$sum": 0}
         else:
             lhs_mql = process_lhs(self, compiler, connection, as_expr=True)
             null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
@@ -65,8 +83,13 @@ def stddev_variance(self, compiler, connection):
     return aggregate(self, compiler, connection, operator=operator)
 
 
+def string_agg(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError("StringAgg is not supported.")
+
+
 def register_aggregates():
     Aggregate.as_mql_expr = aggregate
     Count.as_mql_expr = count
     StdDev.as_mql_expr = stddev_variance
+    StringAgg.as_mql_expr = string_agg
     Variance.as_mql_expr = stddev_variance