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
 from django.db.models.sql.where import WhereNode
@@ -11,13 +19,23 @@
 
 def aggregate(self, compiler, connection, operator=None, resolve_inner_expression=False):
     agg_expression, *_ = self.get_source_expressions()
-    if self.filter:
-        agg_expression = Case(
-            When(self.filter, then=agg_expression),
-            # Skip rows that don't meet the criteria.
-            default=Remove(),
-        )
-    lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
+    if self.filter is not None:
+        try:
+            lhs_mql = self.filter.as_mql(compiler, connection, as_expr=True)
+        except NotSupportedError:
+            # Generate a CASE statement for this aggregate.
+            agg_expression = Case(
+                When(self.filter.condition, then=agg_expression),
+                # Skip rows that don't meet the criteria.
+                default=Remove(),
+            )
+            lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
+        except FullResultSet:
+            lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
+        except EmptyResultSet:
+            lhs_mql = Value(None).as_mql(compiler, connection, as_expr=True)
+    else:
+        lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
     if resolve_inner_expression:
         return lhs_mql
     operator = operator or MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
@@ -34,16 +52,29 @@ def count(self, compiler, connection, resolve_inner_expression=False):
     if not self.distinct or resolve_inner_expression:
         conditions = [IsNull(agg_expression, False)]
         if self.filter:
-            conditions.append(self.filter)
-        inner_expression = Case(
-            When(WhereNode(conditions), then=agg_expression if self.distinct else Value(1)),
-            # Skip rows that don't meet the criteria.
-            default=Remove(),
-        )
-        inner_expression = inner_expression.as_mql(compiler, connection, as_expr=True)
+            try:
+                inner_expression = self.filter.as_mql(compiler, connection, as_expr=True)
+            except NotSupportedError:
+                conditions.append(self.filter.condition)
+                condition = When(
+                    WhereNode(conditions),
+                    then=agg_expression if self.distinct else Value(1),
+                )
+                inner_expression = Case(condition)
+            except FullResultSet:
+                inner_expression = agg_expression if self.distinct else Value(1)
+            except EmptyResultSet:
+                inner_expression = Remove() if self.distinct else Value(0)
+        else:
+            inner_expression = Case(
+                When(WhereNode(conditions), then=agg_expression if self.distinct else Value(1)),
+                # Skip rows that don't meet the criteria.
+                default=Remove(),
+            )
+        lhs_mql = inner_expression.as_mql(compiler, connection, as_expr=True)
         if resolve_inner_expression:
-            return inner_expression
-        return {"$sum": inner_expression}
+            return lhs_mql
+        return {"$sum": lhs_mql}
     # If distinct=True or resolve_inner_expression=False, sum the size of the
     # set.
     return {"$size": agg_expression.as_mql(compiler, connection, as_expr=True)}
@@ -57,8 +88,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