add support for subqueries

Subquery, Exists, and QuerySet as a lookup value.

Author: WaVEV

README.md

@@ -146,9 +146,6 @@ Congratulations, your project is ready to go!
 - `QuerySet.delete()` and `update()` do not support queries that span multiple
   collections.
 
-- `Subquery`, `Exists`, and using a `QuerySet` in `QuerySet.annotate()` aren't
-  supported.
-
 - `DateTimeField` doesn't support microsecond precision, and correspondingly,
   `DurationField` stores milliseconds rather than microseconds.
 

django_mongodb/compiler.py

@@ -13,6 +13,7 @@
 from django.db.models.lookups import IsNull
 from django.db.models.sql import compiler
 from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, SINGLE
+from django.db.models.sql.datastructures import BaseTable
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
@@ -25,12 +26,16 @@ class SQLCompiler(compiler.SQLCompiler):
 
     query_class = MongoQuery
     GROUP_SEPARATOR = "___"
+    PARENT_FIELD_TEMPLATE = "parent__field__{}"
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.aggregation_pipeline = None
+        # Map columns to their subquery indices.
+        self.column_indices = {}
         # A list of OrderBy objects for this query.
         self.order_by_objs = None
+        self.subqueries = []
 
     def _unfold_column(self, col):
         """
@@ -154,23 +159,40 @@ def _prepare_annotations_for_aggregation_pipeline(self, order_by):
         group.update(having_group)
         return group, replacements
 
-    def _get_group_id_expressions(self, order_by):
-        """Generate group ID expressions for the aggregation pipeline."""
-        group_expressions = set()
-        replacements = {}
+    def _get_group_expressions(self, order_by):
+        if self.query.group_by is None:
+            return []
+        seen = set()
+        expressions = set()
+        if self.query.group_by is not True:
+            # If group_by isn't True, then it's a list of expressions.
+            for expr in self.query.group_by:
+                if not hasattr(expr, "as_sql"):
+                    expr = self.query.resolve_ref(expr)
+                if isinstance(expr, Ref):
+                    if expr.refs not in seen:
+                        seen.add(expr.refs)
+                        expressions.add(expr.source)
+                else:
+                    expressions.add(expr)
+        for expr, _, alias in self.select:
+            # Skip members that are already grouped.
+            if alias not in seen:
+                expressions |= set(expr.get_group_by_cols())
         if not self._meta_ordering:
             for expr, (_, _, is_ref) in order_by:
+                # Skip references.
                 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())
+                    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)
-        if isinstance(self.query.group_by, tuple | list):
-            group_expressions |= set(self.query.group_by)
-        elif self.query.group_by is None:
-            group_expressions = set()
+            expressions.add(expr)
+        return expressions
+
+    def _get_group_id_expressions(self, order_by):
+        """Generate group ID expressions for the aggregation pipeline."""
+        replacements = {}
+        group_expressions = self._get_group_expressions(order_by)
         if not group_expressions:
             ids = None
         else:
@@ -186,6 +208,8 @@ def _get_group_id_expressions(self, order_by):
                     ids[alias] = Value(True).as_mql(self, self.connection)
                 if replacement is not None:
                     replacements[col] = replacement
+                    if isinstance(col, Ref):
+                        replacements[col.source] = replacement
         return ids, replacements
 
     def _build_aggregation_pipeline(self, ids, group):
@@ -228,15 +252,15 @@ def pre_sql_setup(self, with_col_aliases=False):
             all_replacements.update(replacements)
             pipeline = self._build_aggregation_pipeline(ids, group)
             if self.having:
-                pipeline.append(
-                    {
-                        "$match": {
-                            "$expr": self.having.replace_expressions(all_replacements).as_mql(
-                                self, self.connection
-                            )
-                        }
-                    }
+                having = self.having.replace_expressions(all_replacements).as_mql(
+                    self, self.connection
                 )
+                # Add HAVING subqueries.
+                for query in self.subqueries or ():
+                    pipeline.extend(query.get_pipeline())
+                # Remove the added subqueries.
+                self.subqueries = []
+                pipeline.append({"$match": {"$expr": having}})
             self.aggregation_pipeline = pipeline
         self.annotations = {
             target: expr.replace_expressions(all_replacements)
@@ -388,6 +412,7 @@ def build_query(self, columns=None):
                 query.mongo_query = {"$expr": expr}
         if extra_fields:
             query.extra_fields = self.get_project_fields(extra_fields, force_expression=True)
+        query.subqueries = self.subqueries
         return query
 
     def get_columns(self):
@@ -431,7 +456,12 @@ def project_field(column):
 
     @cached_property
     def collection_name(self):
-        return self.query.get_meta().db_table
+        base_table = next(
+            v
+            for k, v in self.query.alias_map.items()
+            if isinstance(v, BaseTable) and self.query.alias_refcount[k]
+        )
+        return base_table.table_alias or base_table.table_name
 
     @cached_property
     def collection(self):
@@ -581,7 +611,7 @@ def _get_ordering(self):
         return tuple(fields), sort_ordering, tuple(extra_fields)
 
     def get_where(self):
-        return self.where
+        return getattr(self, "where", self.query.where)
 
     def explain_query(self):
         # Validate format (none supported) and options.
@@ -741,7 +771,7 @@ def build_query(self, columns=None):
             else None
         )
         subquery = compiler.build_query(columns)
-        query.subquery = subquery
+        query.subqueries = [subquery]
         return query
 
     def _make_result(self, result, columns=None):

django_mongodb/expressions.py

@@ -9,6 +9,7 @@
     Case,
     Col,
     CombinedExpression,
+    Exists,
     ExpressionWrapper,
     F,
     NegatedExpression,
@@ -50,6 +51,18 @@ def case(self, compiler, connection):
 
 
 def col(self, compiler, connection):  # noqa: ARG001
+    # If the column is part of a subquery and belongs to one of the parent
+    # queries, it will be stored for reference using $let in a $lookup stage.
+    if (
+        self.alias not in compiler.query.alias_refcount
+        or compiler.query.alias_refcount[self.alias] == 0
+    ):
+        try:
+            index = compiler.column_indices[self]
+        except KeyError:
+            index = len(compiler.column_indices)
+            compiler.column_indices[self] = index
+        return f"$${compiler.PARENT_FIELD_TEMPLATE.format(index)}"
     # Add the column's collection's alias for columns in joined collections.
     prefix = f"{self.alias}." if self.alias != compiler.collection_name else ""
     return f"${prefix}{self.target.column}"
@@ -79,8 +92,73 @@ def order_by(self, compiler, connection):
     return self.expression.as_mql(compiler, connection)
 
 
-def query(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError("Using a QuerySet in annotate() is not supported on MongoDB.")
+def query(self, compiler, connection, lookup_name=None):
+    subquery_compiler = self.get_compiler(connection=connection)
+    subquery_compiler.pre_sql_setup(with_col_aliases=False)
+    columns = subquery_compiler.get_columns()
+    field_name, expr = columns[0]
+    subquery = subquery_compiler.build_query(
+        columns
+        if subquery_compiler.query.annotations or not subquery_compiler.query.default_cols
+        else None
+    )
+    table_output = f"__subquery{len(compiler.subqueries)}"
+    from_table = next(
+        e.table_name for alias, e in self.alias_map.items() if self.alias_refcount[alias]
+    )
+    # To perform a subquery, a $lookup stage that escapsulates the entire
+    # subquery pipeline is added. The "let" clause defines the variables
+    # needed to bridge the main collection with the subquery.
+    subquery.subquery_lookup = {
+        "as": table_output,
+        "from": from_table,
+        "let": {
+            compiler.PARENT_FIELD_TEMPLATE.format(i): col.as_mql(compiler, connection)
+            for col, i in subquery_compiler.column_indices.items()
+        },
+    }
+    # The result must be a list of values. The output is compressed with an
+    # aggregation pipeline.
+    if lookup_name in ("in", "range"):
+        if subquery.aggregation_pipeline is None:
+            subquery.aggregation_pipeline = []
+        subquery.aggregation_pipeline.extend(
+            [
+                {
+                    "$facet": {
+                        "group": [
+                            {
+                                "$group": {
+                                    "_id": None,
+                                    "tmp_name": {
+                                        "$addToSet": expr.as_mql(subquery_compiler, connection)
+                                    },
+                                }
+                            }
+                        ]
+                    }
+                },
+                {
+                    "$project": {
+                        field_name: {
+                            "$ifNull": [
+                                {
+                                    "$getField": {
+                                        "input": {"$arrayElemAt": ["$group", 0]},
+                                        "field": "tmp_name",
+                                    }
+                                },
+                                [],
+                            ]
+                        }
+                    }
+                },
+            ]
+        )
+        # Erase project_fields since the required value is projected above.
+        subquery.project_fields = None
+    compiler.subqueries.append(subquery)
+    return f"${table_output}.{field_name}"
 
 
 def raw_sql(self, compiler, connection):  # noqa: ARG001
@@ -100,8 +178,16 @@ def star(self, compiler, connection):  # noqa: ARG001
     return {"$literal": True}
 
 
-def subquery(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError(f"{self.__class__.__name__} is not supported on MongoDB.")
+def subquery(self, compiler, connection, lookup_name=None):
+    return self.query.as_mql(compiler, connection, lookup_name=lookup_name)
+
+
+def exists(self, compiler, connection, lookup_name=None):
+    try:
+        lhs_mql = subquery(self, compiler, connection, lookup_name=lookup_name)
+    except EmptyResultSet:
+        return Value(False).as_mql(compiler, connection)
+    return connection.mongo_operators["isnull"](lhs_mql, False)
 
 
 def when(self, compiler, connection):
@@ -130,6 +216,7 @@ def register_expressions():
     Case.as_mql = case
     Col.as_mql = col
     CombinedExpression.as_mql = combined_expression
+    Exists.as_mql = exists
     ExpressionWrapper.as_mql = expression_wrapper
     F.as_mql = f
     NegatedExpression.as_mql = negated_expression