Support union.

Author: WaVEV

django_mongodb/compiler.py

@@ -350,21 +350,19 @@ def build_query(self, columns=None):
         self.check_query()
         query = self.query_class(self)
         ordering_fields, sort_ordering, extra_fields = self._get_ordering()
-        query.project_fields = self.get_project_fields(columns, ordering_fields)
         query.ordering = sort_ordering
-        # If columns is None, then get_project_fields() won't add
-        # ordering_fields to $project. Use $addFields (extra_fields) instead.
-        if columns is None:
-            extra_fields += ordering_fields
-        if extra_fields:
-            query.extra_fields = self.get_project_fields(extra_fields, force_expression=True)
         if self.query.combinator:
             if not getattr(self.connection.features, f"supports_select_{self.query.combinator}"):
                 raise NotSupportedError(
                     f"{self.query.combinator} is not supported on this database backend."
                 )
             query.combinator_pipeline = self.get_combinator_queries()
         else:
+            query.project_fields = self.get_project_fields(columns, ordering_fields)
+            # If columns is None, then get_project_fields() won't add
+            # ordering_fields to $project. Use $addFields (extra_fields) instead.
+            if columns is None:
+                extra_fields += ordering_fields
             query.lookup_pipeline = self.get_lookup_pipeline()
             where = self.get_where()
             try:
@@ -373,6 +371,8 @@ def build_query(self, columns=None):
                 query.mongo_query = {}
             else:
                 query.mongo_query = {"$expr": expr}
+        if extra_fields:
+            query.extra_fields = self.get_project_fields(extra_fields, force_expression=True)
         return query
 
     def get_columns(self):
@@ -381,10 +381,8 @@ def get_columns(self):
         which should be loaded by the query.
         """
         select_mask = self.query.get_select_mask()
-        columns = filter(
-            # The extra order by columns are handled by order_by_objs variables.
-            lambda col: not isinstance(col, OrderBy),
-            self.get_default_columns(select_mask) if self.query.default_cols else self.query.select,
+        columns = (
+            self.get_default_columns(select_mask) if self.query.default_cols else self.query.select
         )
         # Populate QuerySet.select_related() data.
         related_columns = []
@@ -400,6 +398,9 @@ def project_field(column):
             if hasattr(column, "target"):
                 # column is a Col.
                 target = column.target.column
+            # Handle Order by columns as refs columns.
+            elif isinstance(column, OrderBy) and isinstance(column.expression, Ref):
+                target = column.expression.refs
             else:
                 # column is a Transform in values()/values_list() that needs a
                 # name for $proj.
@@ -427,6 +428,8 @@ def get_combinator_queries(self):
             query.get_compiler(self.using, self.connection, self.elide_empty)
             for query in self.query.combined_queries
         ]
+        main_query_columns = self.get_columns()
+        main_query_fields, _ = zip(*main_query_columns, strict=True)
         for compiler_ in compilers:
             try:
                 # If the columns list is limited, then all combined queries
@@ -443,8 +446,10 @@ def get_combinator_queries(self):
                     )
                 compiler_.pre_sql_setup()
                 # Standardize columns as main query required.
-                _, exprs = zip(*compiler_.get_columns(), strict=True)
-                columns = tuple(zip(self.query.values_select, exprs, strict=True))
+                columns = compiler_.get_columns()
+                if self.query.values_select:
+                    _, exprs = zip(*columns, strict=True)
+                    columns = tuple(zip(main_query_fields, exprs, strict=True))
                 parts.append((compiler_.build_query(columns), compiler_.collection_name))
 
             except EmptyResultSet:
@@ -464,12 +469,14 @@ def get_combinator_queries(self):
                 )
                 if not self.query.combinator_all:
                     ids = {}
-                    for alias, expr in self.get_columns():
-                        ids[alias] = (
-                            expr.as_mql(self, self.connection)
-                            if isinstance(expr, Col | Ref)
-                            else f"${alias}"
-                        )
+                    for alias, expr in main_query_columns:
+                        collection = expr.alias if isinstance(expr, Col) else None
+                        if collection and collection != self.collection_name:
+                            ids[
+                                f"{expr.alias}{self.GROUP_SEPARATOR}{expr.target.column}"
+                            ] = expr.as_mql(self, self.connection)
+                        else:
+                            ids[alias] = f"${alias}"
                     combinator_pipeline.append({"$group": {"_id": ids}})
                     projected_fields = defaultdict(dict)
                     for key in ids:

django_mongodb/query.py

@@ -3,7 +3,7 @@
 
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError
-from django.db.models.expressions import Case, When
+from django.db.models.expressions import Case, OrderBy, When
 from django.db.models.functions import Mod
 from django.db.models.lookups import Exact
 from django.db.models.sql.constants import INNER
@@ -80,6 +80,8 @@ def get_pipeline(self):
             pipeline.extend(self.aggregation_pipeline)
         if self.project_fields:
             pipeline.append({"$project": self.project_fields})
+        if self.combinator_pipeline:
+            pipeline.extend(self.combinator_pipeline)
         if self.extra_fields:
             pipeline.append({"$addFields": self.extra_fields})
         if self.ordering:
@@ -88,8 +90,6 @@ def get_pipeline(self):
             pipeline.append({"$skip": self.query.low_mark})
         if self.query.high_mark is not None:
             pipeline.append({"$limit": self.query.high_mark - self.query.low_mark})
-        if self.combinator_pipeline:
-            pipeline.extend(self.combinator_pipeline)
 
         return subquery_pipeline + pipeline
 
@@ -171,6 +171,10 @@ def join(self, compiler, connection):
     return lookup_pipeline
 
 
+def orderby(self, compiler, connection):
+    return self.expression.as_mql(compiler, connection)
+
+
 def where_node(self, compiler, connection):
     if self.connector == AND:
         full_needed, empty_needed = len(self.children), 1
@@ -236,4 +240,5 @@ def where_node(self, compiler, connection):
 def register_nodes():
     Join.as_mql = join
     NothingNode.as_mql = NothingNode.as_sql
+    OrderBy.as_mql = orderby
     WhereNode.as_mql = where_node