edits

timgraham · timgraham · commit 554dacb4ac3f · 2024-08-05T13:43:33.000-04:00
diff --git a/README.md b/README.md
@@ -116,12 +116,9 @@ Migrations for 'admin':
   - `extra()`
   - `prefetch_related()`
 
-
 - `QuerySet.delete()` and `update()` do not support queries that span multiple
   collections.
 
-- `aggregate()` Sort with aggregate may not work well.
-
 - `Subquery`, `Exists`, and using a `QuerySet` in `QuerySet.annotate()` aren't
   supported.
 
diff --git a/django_mongodb/compiler.py b/django_mongodb/compiler.py
@@ -45,7 +45,7 @@ def _get_column_from_expression(self, expr, alias):
         return Col(self.collection_name, column_target)
 
     def _prepare_expressions_for_pipeline(self, expression, target, count):
-        """Prepare expressions for the MongoDB aggregation pipeline."""
+        """Prepare expressions for the aggregation pipeline."""
         replacements = {}
         group = {}
         for sub_expr in self._get_aggregate_expressions(expression):
@@ -62,12 +62,12 @@ def _prepare_expressions_for_pipeline(self, expression, target, count):
             else:
                 group[alias] = sub_expr.as_mql(self, self.connection)
                 replacing_expr = inner_column
-
+            # Count must return 0 rather than null.
             if isinstance(sub_expr, Count):
                 replacing_expr = Coalesce(replacing_expr, 0)
+            # Variance = StdDev^2
             if isinstance(sub_expr, Variance):
                 replacing_expr = Power(replacing_expr, 2)
-
             replacements[sub_expr] = replacing_expr
         return replacements, group
 
@@ -170,10 +170,10 @@ def _build_group_pipeline(self, ids, group):
     def pre_sql_setup(self, with_col_aliases=False):
         extra_select, order_by, group_by = super().pre_sql_setup(with_col_aliases=with_col_aliases)
         group, all_replacements = self._prepare_annotations_for_group_pipeline()
-
-        # The query.group_by is either None (no GROUP BY at all), True
-        # (group by select fields), or a list of expressions to be added
-        # to the group by.
+        # query.group_by is either:
+        # - None: no GROUP BY
+        # - True: group by select fields
+        # - a list of expressions to group by.
         if group or self.query.group_by:
             ids, replacements = self._get_group_id_expressions(order_by)
             all_replacements.update(replacements)
@@ -194,7 +194,6 @@ def pre_sql_setup(self, with_col_aliases=False):
             target: expr.replace_expressions(all_replacements)
             for target, expr in self.query.annotation_select.items()
         }
-
         return extra_select, order_by, group_by
 
     def execute_sql(
@@ -455,18 +454,16 @@ def get_project_fields(self, columns=None, ordering=None):
                 # 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
-
+            # Add order_by() fields.
             for column, _ in ordering or []:
                 foreign_table = column.split(".", 1)[0] if "." in column else None
                 if column not in fields and foreign_table not in fields:
                     fields[column] = 1
-
         return fields
 
 
diff --git a/django_mongodb/features.py b/django_mongodb/features.py
@@ -72,7 +72,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "db_functions.math.test_round.RoundTests.test_integer_with_negative_precision",
         # Truncating in another timezone doesn't work becauase MongoDB converts
         # the result back to UTC.
-        "aggregation.tests.AggregateTestCase.test_aggregation_default_using_date_from_database",
         "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_func_with_timezone",
         "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_timezone_applied_before_truncation",
         # Length of null considered zero rather than null.
@@ -104,7 +103,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # QuerySet.explain() not implemented:
         # https://github.com/mongodb-labs/django-mongodb/issues/28
         "queries.test_explain.ExplainUnsupportedTests.test_message",
-        # Sum returns 0 instead of None in mongodb.
+        # The $sum aggregation returns 0 instead of None for null.
         "aggregation.test_filter_argument.FilteredAggregateTests.test_plain_annotate",
         "aggregation.tests.AggregateTestCase.test_aggregation_default_passed_another_aggregate",
         "aggregation.tests.AggregateTestCase.test_annotation_expressions",
@@ -504,7 +503,9 @@ def django_test_expected_failures(self):
             "timezones.tests.NewDatabaseTests.test_cursor_explicit_time_zone",
             "timezones.tests.NewDatabaseTests.test_raw_sql",
         },
-        "Custom functions with SQL don't work on MongoDB.": {
+        "Custom aggregations/functions with SQL don't work on MongoDB.": {
+            "aggregation.tests.AggregateTestCase.test_add_implementation",
+            "aggregation.tests.AggregateTestCase.test_multi_arg_aggregate",
             "annotations.tests.NonAggregateAnnotationTestCase.test_custom_functions",
             "annotations.tests.NonAggregateAnnotationTestCase.test_custom_functions_can_ref_other_functions",
         },
@@ -533,6 +534,7 @@ def django_test_expected_failures(self):
             "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_quarter_func_boundaries",
         },
         "TruncDate database function not supported.": {
+            "aggregation.tests.AggregateTestCase.test_aggregation_default_using_date_from_database",
             "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_date_func",
             "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_date_none",
             "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_lookup_name_sql_injection",
@@ -588,9 +590,6 @@ def django_test_expected_failures(self):
         },
         "Test not applicable for MongoDB's SQLCompiler.": {
             "queries.test_iterator.QuerySetIteratorTests",
-            # Custom aggregations:
-            "aggregation.tests.AggregateTestCase.test_add_implementation",
-            "aggregation.tests.AggregateTestCase.test_multi_arg_aggregate",
         },
     }
 
diff --git a/django_mongodb/functions.py b/django_mongodb/functions.py
@@ -41,8 +41,8 @@
 
 MONGO_AGGREGATIONS = {
     Count: "sum",
-    StdDev: "stdDev",
-    Variance: "stdDev",
+    StdDev: "stdDev",  # Samp or Pop suffix added in aggregate().
+    Variance: "stdDev",  # Likewise.
 }
 MONGO_OPERATORS = {
     Ceil: "ceil",
@@ -79,11 +79,11 @@ def aggregate(self, compiler, connection, **extra_context):  # noqa: ARG001
         node = self
     lhs_mql = process_lhs(node, compiler, connection)
     operator = MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
-    if "_SAMP" in self.function:
+    # Add suffixes to StdDev/Variance.
+    if self.function.endswith("_SAMP"):
         operator += "Samp"
-    elif "_POP" in self.function:
+    elif self.function.endswith("_POP"):
         operator += "Pop"
-
     return {f"${operator}": lhs_mql}
 
 
@@ -119,26 +119,24 @@ def cot(self, compiler, connection):
 
 def count(self, compiler, connection, resolve_inner_expression=False, **extra_context):  # noqa: ARG001
     """
-    When resolve_inner_expression is True, it returns the argument as 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.
+    When resolve_inner_expression is True, return the argument as 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:
-            copy = self.copy()
-            copy.filter = None
-            source_expressions = copy.get_source_expressions()
+            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))
-            copy.set_source_expressions([Case(condition)] + source_expressions[1:])
-            node = copy
+            node.set_source_expressions([Case(condition)] + source_expressions[1:])
             inner_expression = process_lhs(node, compiler, connection)
         else:
-            node = self
             lhs_mql = process_lhs(self, compiler, connection)
             null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
             inner_expression = {
@@ -147,10 +145,10 @@ def count(self, compiler, connection, resolve_inner_expression=False, **extra_co
         if resolve_inner_expression:
             return inner_expression
         return {"$sum": inner_expression}
-
-    # When count is called with distinct without the flag, we sum the size of the set.
+    # If distinct=True or resolve_inner_expression=False, sum the size
+    # of the set.
     lhs_mql = process_lhs(self, compiler, connection)
-    # And subtract 1 if None is in the set (it shouldn't have been counted).
+    # Subtract 1 if None is in the set (it shouldn't have been counted).
     exits_null = {"$cond": {"if": {"$in": [{"$literal": None}, lhs_mql]}, "then": -1, "else": 0}}
     return {"$add": [{"$size": lhs_mql}, exits_null]}
 
diff --git a/django_mongodb/query.py b/django_mongodb/query.py
@@ -87,12 +87,10 @@ def get_cursor(self):
         """
         if self.query.low_mark == self.query.high_mark:
             return []
-
-        pipeline = self.get_pipeline()
-        return self.collection.aggregate(pipeline)
+        return self.collection.aggregate(self.get_pipeline())
 
     def get_pipeline(self):
-        pipeline = [] if self.subquery is None else self.subquery.get_pipeline()
+        pipeline = self.subquery.get_pipeline() if self.subquery else []
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
         if self.mongo_query:
