Support Union query, first approach.

Author: WaVEV

django_mongodb/compiler.py

@@ -349,7 +349,6 @@ 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
@@ -359,13 +358,21 @@ def build_query(self, columns=None):
             extra_fields += ordering_fields
         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 = {}
+        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.mongo_query = {"$expr": expr}
+            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}
         return query
 
     def get_columns(self):
@@ -412,6 +419,61 @@ 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
+        ]
+        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(with_col_aliases=False)
+                # Avoid $project (columns=None) if unneeded.
+                columns = (
+                    compiler_.get_columns()
+                    if compiler_.query.annotations or not compiler_.query.default_cols
+                    else None
+                )
+                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
+
+        combinator_pipeline = parts.pop(0)[0].get_pipeline() if parts else None
+        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 = {}
+                    annotation_group_idx = itertools.count(start=1)
+                    for _, expr in self.get_columns():
+                        alias, replacement = self._get_group_alias_column(
+                            expr, annotation_group_idx
+                        )
+                        ids[alias] = expr.as_mql(self, self.connection)
+                    combinator_pipeline.append({"$group": {"_id": ids}})
+        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

@@ -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}>"
@@ -69,7 +70,8 @@ def get_cursor(self):
         return self.collection.aggregate(self.get_pipeline())
 
     def get_pipeline(self):
-        pipeline = self.subquery.get_pipeline() if self.subquery else []
+        subquery_pipeline = self.subquery.get_pipeline() if self.subquery else []
+        pipeline = []
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
         if self.mongo_query:
@@ -86,7 +88,10 @@ 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})
-        return pipeline
+        if self.combinator_pipeline:
+            pipeline.extend(self.combinator_pipeline)
+
+        return subquery_pipeline + pipeline
 
 
 def join(self, compiler, connection):