Support Union.

Author: WaVEV

django_mongodb/compiler.py

@@ -7,7 +7,7 @@
 from django.db import IntegrityError, NotSupportedError
 from django.db.models import Count
 from django.db.models.aggregates import Aggregate, Variance
-from django.db.models.expressions import Case, Col, Ref, Value, When
+from django.db.models.expressions import Case, Col, OrderBy, Ref, Value, When
 from django.db.models.functions.comparison import Coalesce
 from django.db.models.functions.math import Power
 from django.db.models.lookups import IsNull
@@ -349,23 +349,30 @@ def build_query(self, columns=None):
         """Check if the query is supported and prepare a MongoQuery."""
         self.check_query()
         query = self.query_class(self)
-        query.lookup_pipeline = self.get_lookup_pipeline()
         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 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:
+                expr = where.as_mql(self, self.connection) if where else {}
+            except FullResultSet:
+                query.mongo_query = {}
+            else:
+                query.mongo_query = {"$expr": expr}
         if extra_fields:
             query.extra_fields = self.get_project_fields(extra_fields, force_expression=True)
-        where = self.get_where()
-        try:
-            expr = where.as_mql(self, self.connection) if where else {}
-        except FullResultSet:
-            query.mongo_query = {}
-        else:
-            query.mongo_query = {"$expr": expr}
         return query
 
     def get_columns(self):
@@ -391,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.
@@ -412,6 +422,79 @@ def collection_name(self):
     def collection(self):
         return self.connection.get_collection(self.collection_name)
 
+    def get_combinator_queries(self):
+        parts = []
+        compilers = [
+            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
+                # must have the same columns list. Set the selects defined on
+                # the query on all combined queries, if not already set.
+                if not compiler_.query.values_select and self.query.values_select:
+                    compiler_.query = compiler_.query.clone()
+                    compiler_.query.set_values(
+                        (
+                            *self.query.extra_select,
+                            *self.query.values_select,
+                            *self.query.annotation_select,
+                        )
+                    )
+                compiler_.pre_sql_setup()
+                # Standardize columns as main query required.
+                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:
+                # Omit the empty queryset with UNION and with DIFFERENCE if the
+                # first queryset is nonempty.
+                if self.query.combinator == "union":
+                    continue
+                raise
+        # Raise EmptyResultSet if all the combinator queries are empty.
+        if not parts:
+            raise EmptyResultSet
+        combinator_pipeline = parts.pop(0)[0].get_pipeline()
+        if self.query.combinator == "union":
+            for part, collection in parts:
+                combinator_pipeline.append(
+                    {"$unionWith": {"coll": collection, "pipeline": part.get_pipeline()}}
+                )
+                if not self.query.combinator_all:
+                    ids = {}
+                    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:
+                        value = f"$_id.{key}"
+                        if self.GROUP_SEPARATOR in key:
+                            table, field = key.split(self.GROUP_SEPARATOR)
+                            projected_fields[table][field] = value
+                        else:
+                            projected_fields[key] = value
+                    # Convert defaultdict to dict so it doesn't appear as
+                    # "defaultdict(<CLASS 'dict'>, ..." in query logging.
+                    combinator_pipeline.append({"$addFields": dict(projected_fields)})
+                    if "_id" not in projected_fields:
+                        combinator_pipeline.append({"$unset": "_id"})
+        else:
+            raise NotSupportedError(f"Combinator {self.query.combinator} isn't supported.")
+        return combinator_pipeline
+
     def get_lookup_pipeline(self):
         result = []
         for alias in tuple(self.query.alias_map):

django_mongodb/features.py

@@ -23,7 +23,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     supports_select_difference = False
     supports_select_intersection = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/72
-    supports_select_union = False
+    supports_select_union = True
     supports_sequence_reset = False
     supports_table_check_constraints = False
     supports_temporal_subtraction = True

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
@@ -51,6 +51,7 @@ def __init__(self, compiler):
         self.project_fields = None
         self.aggregation_pipeline = compiler.aggregation_pipeline
         self.extra_fields = None
+        self.combinator_pipeline = None
 
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"
@@ -78,6 +79,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:
@@ -166,6 +169,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
@@ -231,4 +238,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