Some fixes.

Author: WaVEV

django_mongodb/compiler.py

@@ -18,6 +18,29 @@ class SQLCompiler(compiler.SQLCompiler):
     query_class = MongoQuery
     _group_pipeline = None
 
+    def _prepare_expressions_for_pipeline(self, expression, target, aggregation_idx):
+        replacements = {}
+        group = {}
+        for sub_expr in self._get_aggregate_expressions(expression):
+            alias = f"__aggregation{aggregation_idx}" if sub_expr != expression else target
+
+            column_target = sub_expr.output_field.__class__()
+            column_target.db_column = alias
+            column_target.set_attributes_from_name(alias)
+            inner_column = Col(self.collection_name, column_target)
+            if sub_expr.distinct:
+                inner_expr = sub_expr.as_mql(self, self.connection, force_filters=True)
+                rhs = next(iter(inner_expr.values()))
+                group[alias] = {"$addToSet": rhs}
+                replacing_expr = sub_expr.copy()
+                replacing_expr.set_source_expressions([inner_column])
+            else:
+                group[alias] = sub_expr.as_mql(self, self.connection)
+                replacing_expr = inner_column
+
+            replacements[sub_expr] = replacing_expr
+        return replacements, group
+
     @staticmethod
     def _random_separtor():
         import random
@@ -32,39 +55,24 @@ def pre_sql_setup(self, with_col_aliases=False):
         self.annotations = {}
         group = {}
         group_expressions = set()
-        aggregation_idx = 1
         all_replacements = {}
-        for target, expr in self.query.annotation_select.items():
-            if not expr.contains_aggregate:
-                result_expr = expr
-            else:
-                replacements = {}
-                for sub_expr in self._get_aggregate_expressions(expr):
-                    if sub_expr != expr:
-                        alias = f"__aggregation{aggregation_idx}"
-                        aggregation_idx += 1
-                    else:
-                        alias = target
-
-                    column_target = expr.output_field.__class__()
-                    column_target.db_column = alias
-                    column_target.set_attributes_from_name(alias)
-                    inner_column = Col(self.collection_name, column_target)
-                    if sub_expr.distinct:
-                        inner_expr = sub_expr.as_mql(self, self.connection, force_filters=True)
-                        rhs = next(iter(inner_expr.values()))
-                        group[alias] = {"$addToSet": rhs}
-                        replacing_expr = sub_expr.copy()
-                        replacing_expr.set_source_expressions([inner_column])
-                    else:
-                        group[alias] = sub_expr.as_mql(self, self.connection)
-                        replacing_expr = inner_column
-
-                    replacements[sub_expr] = replacing_expr
+        for idx, (target, expr) in enumerate(self.query.annotation_select.items()):
+            if expr.contains_aggregate:
+                replacements, expr_group = self._prepare_expressions_for_pipeline(expr, target, idx)
                 result_expr = expr.replace_expressions(replacements)
                 all_replacements.update(replacements)
+                group.update(expr_group)
+            else:
+                result_expr = expr
             group_expressions |= set(expr.get_group_by_cols())
             self.annotations[target] = result_expr
+
+        having_replacements, having_group = self._prepare_expressions_for_pipeline(
+            self.having, None, len(self.query.annotation_select)
+        )
+        all_replacements.update(having_replacements)
+        group.update(having_group)
+
         if group or self.query.group_by:
             order_by = self.get_order_by()
             for expr, (_, _, is_ref) in order_by:
@@ -94,7 +102,14 @@ def pre_sql_setup(self, with_col_aliases=False):
                     break
             SEPARATOR = f"__{random_string}__"
 
+            annotation_group_idx = 0
+
             def _ccc(col):
+                nonlocal annotation_group_idx
+
+                if not isinstance(col, Col):
+                    annotation_group_idx += 1
+                    return "__annotation_group_1"
                 if self.collection_name == col.alias:
                     return col.target.column
                 return f"{col.alias}{SEPARATOR}{col.target.column}"
@@ -106,7 +121,6 @@ def _ccc(col):
                     _ccc(col): col.as_mql(self, self.connection)
                     # expression aren't needed in the group by clouse ()
                     for col in group_expressions
-                    if isinstance(col, Col)
                 }
             )
             pipeline = []
@@ -140,10 +154,7 @@ def _ccc(col):
                     else:
                         sets[key] = value
 
-                pipeline.append(
-                    # {"$addFields": {key: f"$_id.{value[1:]}" for key, value in ids.items()}}
-                    {"$addFields": sets}
-                )
+                pipeline.append({"$addFields": sets})
                 if "_id" not in sets:
                     pipeline.append({"$unset": "_id"})