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