add support for QuerySet.annotate()

Author: timgraham

README.md

@@ -47,7 +47,6 @@ DATABASES = {
 ## Known issues and limitations
 
 - The following `QuerySet` methods aren't supported:
-  - `annotate()`
   - `aggregate()`
   - `dates()`
   - `datetimes()`

django_mongodb/base.py

@@ -54,7 +54,25 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         "TimeField": "date",
         "UUIDField": "string",
     }
+    # Django uses these operators to generate SQL queries before it generates
+    # MQL queries.
     operators = {
+        "exact": "= %s",
+        "iexact": "= UPPER(%s)",
+        "contains": "LIKE %s",
+        "icontains": "LIKE UPPER(%s)",
+        "regex": "~ %s",
+        "iregex": "~* %s",
+        "gt": "> %s",
+        "gte": ">= %s",
+        "lt": "< %s",
+        "lte": "<= %s",
+        "startswith": "LIKE %s",
+        "endswith": "LIKE %s",
+        "istartswith": "LIKE UPPER(%s)",
+        "iendswith": "LIKE UPPER(%s)",
+    }
+    mongo_operators = {
         "exact": lambda val: val,
         "gt": lambda val: {"$gt": val},
         "gte": lambda val: {"$gte": val},
@@ -73,6 +91,13 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         "regex": lambda val: re.compile(val),
         "iregex": lambda val: re.compile(val, re.IGNORECASE),
     }
+    mongo_aggregates = {
+        "exact": lambda a, b: {"$eq": [a, b]},
+        "gt": lambda a, b: {"$gt": [a, b]},
+        "gte": lambda a, b: {"$gte": [a, b]},
+        "lt": lambda a, b: {"$lt": [a, b]},
+        "lte": lambda a, b: {"$lte": [a, b]},
+    }
 
     display_name = "MongoDB"
     vendor = "mongodb"

django_mongodb/compiler.py

@@ -1,6 +1,6 @@
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError, NotSupportedError
-from django.db.models import NOT_PROVIDED, Count, Expression, Value
+from django.db.models import NOT_PROVIDED, Count, Expression
 from django.db.models.aggregates import Aggregate
 from django.db.models.constants import LOOKUP_SEP
 from django.db.models.sql import compiler
@@ -57,11 +57,16 @@ def has_results(self):
 
     def get_converters(self, columns):
         converters = {}
-        for column in columns:
+        for name_column in columns:
+            try:
+                name, column = name_column
+            except TypeError:
+                # e.g., Count("*")
+                continue
             backend_converters = self.connection.ops.get_db_converters(column)
             field_converters = column.field.get_db_converters(self.connection)
             if backend_converters or field_converters:
-                converters[column.target.column] = backend_converters + field_converters
+                converters[name] = backend_converters + field_converters
         return converters
 
     def _make_result(self, entity, columns, converters, tuple_expected=False):
@@ -72,15 +77,14 @@ def _make_result(self, entity, columns, converters, tuple_expected=False):
         names as keys.
         """
         result = []
-        for col in columns:
+        for name, col in columns:
             field = col.field
-            column = col.target.column
-            value = entity.get(column, NOT_PROVIDED)
+            value = entity.get(name, NOT_PROVIDED)
             if value is NOT_PROVIDED:
                 value = field.get_default()
             elif converters:
                 # Decode values using Django's database converters API.
-                for converter in converters.get(column, ()):
+                for converter in converters.get(name, ()):
                     value = converter(value, col, self.connection)
             result.append(value)
         if tuple_expected:
@@ -91,12 +95,6 @@ def check_query(self):
         """Check if the current query is supported by the database."""
         if self.query.is_empty():
             raise EmptyResultSet()
-        # Supported annotations are Exists() and Count().
-        if self.query.annotations and self.query.annotations not in (
-            {"a": Value(1)},
-            {"__count": Count("*")},
-        ):
-            raise NotSupportedError("QuerySet.annotate() is not supported on MongoDB.")
         if self.query.distinct:
             # This is a heuristic to detect QuerySet.datetimes() and dates().
             # "datetimefield" and "datefield" are the names of the annotations
@@ -144,11 +142,17 @@ def build_query(self, columns=None):
         return query
 
     def get_columns(self):
-        """Return columns which should be loaded by the query."""
+        """
+        Return a tuple of (name, expression) with the columns and annotations
+        which should be loaded by the query.
+        """
         select_mask = self.query.get_select_mask()
-        return (
+        columns = (
             self.get_default_columns(select_mask) if self.query.default_cols else self.query.select
         )
+        return tuple((column.target.column, column) for column in columns) + tuple(
+            self.query.annotations.items()
+        )
 
     def _get_ordering(self):
         """

django_mongodb/features.py

@@ -133,27 +133,30 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "timezones.tests.LegacyDatabaseTests.test_query_aggregation",
             "timezones.tests.NewDatabaseTests.test_query_aggregation",
         },
-        "QuerySet.annotate() not supported.": {
-            "lookup.test_decimalfield.DecimalFieldLookupTests",
+        "QuerySet.annotate() has some limitations.": {
+            # Exists not supported.
             "lookup.tests.LookupTests.test_exact_exists",
             "lookup.tests.LookupTests.test_nested_outerref_lhs",
+            "lookup.tests.LookupQueryingTests.test_filter_exists_lhs",
+            # QuerySet.alias() doesn't work.
             "lookup.tests.LookupQueryingTests.test_alias",
-            "lookup.tests.LookupQueryingTests.test_annotate",
+            # Comparing two fields doesn't work.
             "lookup.tests.LookupQueryingTests.test_annotate_field_greater_than_field",
-            "lookup.tests.LookupQueryingTests.test_annotate_field_greater_than_literal",
-            "lookup.tests.LookupQueryingTests.test_annotate_field_greater_than_value",
-            "lookup.tests.LookupQueryingTests.test_annotate_greater_than_or_equal",
-            "lookup.tests.LookupQueryingTests.test_annotate_greater_than_or_equal_float",
-            "lookup.tests.LookupQueryingTests.test_annotate_less_than_float",
+            # Value() not supported.
             "lookup.tests.LookupQueryingTests.test_annotate_literal_greater_than_field",
             "lookup.tests.LookupQueryingTests.test_annotate_value_greater_than_value",
+            # annotate() with combined expressions doesn't work:
+            # 'WhereNode' object has no attribute 'field'
             "lookup.tests.LookupQueryingTests.test_combined_annotated_lookups_in_filter",
             "lookup.tests.LookupQueryingTests.test_combined_annotated_lookups_in_filter_false",
             "lookup.tests.LookupQueryingTests.test_combined_lookups",
+            # Case not supported.
             "lookup.tests.LookupQueryingTests.test_conditional_expression",
-            "lookup.tests.LookupQueryingTests.test_filter_exists_lhs",
+            # Using expression in filter() doesn't work.
             "lookup.tests.LookupQueryingTests.test_filter_lookup_lhs",
+            # Subquery not supported.
             "lookup.tests.LookupQueryingTests.test_filter_subquery_lhs",
+            # ExpressionWrapper not supported.
             "lookup.tests.LookupQueryingTests.test_filter_wrapped_lookup_lhs",
         },
         "QuerySet.dates() is not supported on MongoDB.": {
@@ -190,6 +193,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "defer.tests.DeferTests.test_only_baseclass_when_subclass_has_no_added_fields",
             "defer.tests.TestDefer2.test_defer_inheritance_pk_chaining",
             "defer_regress.tests.DeferRegressionTest.test_ticket_16409",
+            "lookup.test_decimalfield.DecimalFieldLookupTests",
             "lookup.tests.LookupQueryingTests.test_multivalued_join_reuse",
             "lookup.tests.LookupTests.test_filter_by_reverse_related_field_transform",
             "lookup.tests.LookupTests.test_lookup_collision",

django_mongodb/lookups.py

@@ -8,10 +8,16 @@
 def builtin_lookup(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
     value = process_rhs(self, compiler, connection)
-    rhs_mql = connection.operators[self.lookup_name](value)
+    rhs_mql = connection.mongo_operators[self.lookup_name](value)
     return {lhs_mql: rhs_mql}
 
 
+def builtin_lookup_agg(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    return connection.mongo_aggregates[self.lookup_name](lhs_mql, value)
+
+
 def exact(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)
     value = process_rhs(self, compiler, connection)
@@ -28,7 +34,7 @@ def is_null(self, compiler, connection):
     if not isinstance(self.rhs, bool):
         raise ValueError("The QuerySet value for an isnull lookup must be True or False.")
     lhs_mql = process_lhs(self, compiler, connection, bare_column_ref=True)
-    rhs_mql = connection.operators["isnull"](self.rhs)
+    rhs_mql = connection.mongo_operators["isnull"](self.rhs)
     return {lhs_mql: rhs_mql}
 
 
@@ -38,6 +44,7 @@ def uuid_text_mixin(self, compiler, connection):  # noqa: ARG001
 
 def register_lookups():
     BuiltinLookup.as_mql = builtin_lookup
+    BuiltinLookup.as_mql_agg = builtin_lookup_agg
     Exact.as_mql = exact
     In.as_mql = RelatedIn.as_mql = in_
     IsNull.as_mql = is_null

django_mongodb/query.py

@@ -85,7 +85,13 @@ def delete(self):
     def get_cursor(self):
         if self.query.low_mark == self.query.high_mark:
             return []
-        fields = [col.target.column for col in self.columns] if self.columns else None
+        fields = {}
+        for name, expr in self.columns or []:
+            try:
+                fields[expr.target.column] = 1
+            except AttributeError:
+                # Generate the MQL for an annotation.
+                fields[name] = expr.as_mql_agg(self.compiler, self.connection)
         cursor = self.collection.find(self.mongo_query, fields)
         if self.ordering:
             cursor.sort(self.ordering)