Refactor projection and expressions.

Author: WaVEV

django_mongodb/compiler.py

@@ -34,12 +34,19 @@ def __init__(self, *args, **kwargs):
     def _get_group_alias_column(self, expr, annotation_group_idx):
         """Generate a dummy field for use in the ids fields in $group."""
         replacement = None
-        if isinstance(expr, Col):
-            col = expr
+
+        # Unwrap Ref (in this part of the pipeline we can't do references over $projected fields).
+        expr_ = expr
+        while isinstance(expr_, Ref):
+            expr_ = expr_.source
+            replacement = expr_
+
+        if isinstance(expr_, Col):
+            col = expr_
         else:
             # If the column is a composite expression, create a field for it.
             alias = f"__annotation_group{next(annotation_group_idx)}"
-            col = self._get_column_from_expression(expr, alias)
+            col = self._get_column_from_expression(expr_, alias)
             replacement = col
         if self.collection_name == col.alias:
             return col.target.column, replacement
@@ -137,14 +144,21 @@ def _get_group_id_expressions(self, order_by):
             for expr, (_, _, is_ref) in order_by:
                 if not is_ref:
                     group_expressions |= set(expr.get_group_by_cols())
-        for expr, *_ in self.select:
-            group_expressions |= set(expr.get_group_by_cols())
         having_group_by = self.having.get_group_by_cols() if self.having else ()
         for expr in having_group_by:
             group_expressions.add(expr)
+
+        refs_viewed = set()
+        for expr, _, alias in self.select:
+            group_expressions |= set(expr.get_group_by_cols())
+            refs_viewed.add(alias)
         if isinstance(self.query.group_by, tuple | list):
-            group_expressions |= set(self.query.group_by)
-        elif self.query.group_by is None:
+            for expr in self.query.group_by:
+                if not isinstance(expr, Ref) or expr.refs not in refs_viewed:
+                    group_expressions.add(expr)
+                if isinstance(expr, Ref):
+                    refs_viewed.add(expr.refs)
+        if self.query.group_by is None:
             group_expressions = set()
         if not group_expressions:
             ids = None
@@ -428,32 +442,22 @@ def _get_aggregate_expressions(self, expr):
                 stack.extend(expr.get_source_expressions())
 
     def get_project_fields(self, columns=None, ordering=None):
-        fields = {}
+        fields = defaultdict(dict)
         for name, expr in columns or []:
+            collection = expr.alias if isinstance(expr, Col) else None
             try:
-                column = expr.target.column
-            except AttributeError:
-                # Generate the MQL for an annotation.
-                try:
-                    fields[name] = expr.as_mql(self, self.connection)
-                except EmptyResultSet:
-                    fields[name] = Value(False).as_mql(self, self.connection)
-                except FullResultSet:
-                    fields[name] = Value(True).as_mql(self, self.connection)
-            else:
-                # 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
+                fields[collection][name] = expr.as_mql(self, self.connection)
+            except EmptyResultSet:
+                fields[collection][name] = Value(False).as_mql(self, self.connection)
+            except FullResultSet:
+                fields[collection][name] = Value(True).as_mql(self, self.connection)
+        # Unwrap annotations.
+        fields.update(fields.pop(None, {}))
+        # Unwrap main collection's fields.
+        fields.update(fields.pop(self.collection_name, {}))
+        if fields and ordering:
             # Add order_by() fields.
-            for alias, expression in ordering or []:
-                nested_entity = alias.split(".", 1)[0] if "." in alias else None
-                if alias not in fields and nested_entity not in fields:
-                    fields[alias] = expression.as_mql(self, self.connection)
+            fields.update({alias: expr.as_mql(self, self.connection) for alias, expr in ordering})
         return fields
 
     def _get_ordering(self):

django_mongodb/expressions.py

@@ -77,7 +77,11 @@ def query(self, compiler, connection):  # noqa: ARG001
 
 
 def ref(self, compiler, connection):  # noqa: ARG001
-    return f"${self.refs}"
+    if isinstance(self.source, Col) and self.source.alias != compiler.collection_name:
+        prefix = f"{self.source.alias}."
+    else:
+        prefix = ""
+    return f"${prefix}{self.refs}"
 
 
 def star(self, compiler, connection):  # noqa: ARG001

django_mongodb/features.py

@@ -50,15 +50,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # Length of null considered zero rather than null.
         "db_functions.text.test_length.LengthTests.test_basic",
         # Wrong results in queries with multiple tables.
-        "annotations.tests.NonAggregateAnnotationTestCase.test_annotation_reverse_m2m",
-        "annotations.tests.NonAggregateAnnotationTestCase.test_annotation_with_m2m",
-        "annotations.tests.NonAggregateAnnotationTestCase.test_chaining_annotation_filter_with_m2m",
-        "annotations.tests.NonAggregateAnnotationTestCase.test_mti_annotations",
-        "expressions.test_queryset_values.ValuesExpressionsTests.test_values_list_expression",
-        "expressions.test_queryset_values.ValuesExpressionsTests.test_values_list_expression_flat",
         "expressions.tests.IterableLookupInnerExpressionsTests.test_expressions_in_lookups_join_choice",
         "queries.tests.Queries1Tests.test_order_by_tables",
-        "queries.tests.TestTicket24605.test_ticket_24605",
         "queries.tests.TestInvalidValuesRelation.test_invalid_values",
         # QuerySet.explain() not implemented:
         # https://github.com/mongodb-labs/django-mongodb/issues/28
@@ -254,6 +247,7 @@ def django_test_expected_failures(self):
             "queries.tests.Queries6Tests.test_tickets_8921_9188",
             "queries.tests.Queries6Tests.test_xor_subquery",
             "queries.tests.QuerySetBitwiseOperationTests.test_subquery_aliases",
+            "queries.tests.TestTicket24605.test_ticket_24605",
             "queries.tests.Ticket20101Tests.test_ticket_20101",
             "queries.tests.Ticket20788Tests.test_ticket_20788",
             "queries.tests.Ticket22429Tests.test_ticket_22429",