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 val1, val2: {"$eq": [val1, val2]},
+        "gt": lambda val1, val2: {"$gt": [val1, val2]},
+        "gte": lambda val1, val2: {"$gte": [val1, val2]},
+        "lt": lambda val1, val2: {"$lt": [val1, val2]},
+        "lte": lambda val1, val2: {"$lte": [val1, val2]},
+    }
 
     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
@@ -9,6 +9,8 @@
 from .base import Cursor
 from .query import MongoQuery, wrap_database_errors
 
+COUNT_STAR = {"__count": Count("*")}
+
 
 class SQLCompiler(compiler.SQLCompiler):
     """Base class for all Mongo compilers."""
@@ -20,7 +22,7 @@ def execute_sql(
     ):
         self.pre_sql_setup()
         # QuerySet.count()
-        if self.query.annotations == {"__count": Count("*")}:
+        if self.query.annotations == COUNT_STAR:
             return [self.get_count()]
         try:
             query = self.build_query()
@@ -57,11 +59,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 +79,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 +97,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
@@ -146,9 +146,12 @@ def build_query(self, columns=None):
     def get_columns(self):
         """Return columns 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()  # if self.query.annotations.items() != COUNT_STAR else ()
+        )
 
     def _get_ordering(self):
         """

django_mongodb/features.py

@@ -131,27 +131,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.": {
@@ -188,6 +191,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,12 @@ 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:
+                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)