add support for order_by() on related collections

timgraham · timgraham · commit c05c3a237610 · 2024-07-23T09:11:12.000-04:00
Also add support for ordering by F and OrderBy expressions.
diff --git a/README.md b/README.md
@@ -121,6 +121,9 @@ Migrations for 'admin':
 - `Subquery`, `Exists`, and using a `QuerySet` in `QuerySet.annotate()` aren't
   supported.
 
+* Ordering a `QuerySet` by `nulls_first` or `nulls_last` isn't supported.
+  Neither is randomized ordering.
+
 - `DateTimeField` doesn't support microsecond precision, and correspondingly,
   `DurationField` stores milliseconds rather than microseconds.
 
diff --git a/django_mongodb/compiler.py b/django_mongodb/compiler.py
@@ -1,10 +1,10 @@
-from django.core.exceptions import EmptyResultSet, FieldDoesNotExist, FullResultSet
+from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError, NotSupportedError
 from django.db.models import Count, Expression
 from django.db.models.aggregates import Aggregate
-from django.db.models.constants import LOOKUP_SEP
+from django.db.models.expressions import OrderBy
 from django.db.models.sql import compiler
-from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI
+from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, ORDER_DIR
 from django.utils.functional import cached_property
 
 from .base import Cursor
@@ -199,31 +199,37 @@ def _get_ordering(self):
             if self.query.default_ordering
             else self.query.order_by
         )
-
         if not ordering:
             return self.query.standard_ordering
-
+        default_order, _ = ORDER_DIR["ASC" if self.query.standard_ordering else "DESC"]
         column_ordering = []
+        columns_seen = set()
         for order in ordering:
-            if LOOKUP_SEP in order:
-                raise NotSupportedError("Ordering can't span tables on MongoDB (%s)." % order)
             if order == "?":
                 raise NotSupportedError("Randomized ordering isn't supported by MongoDB.")
-
-            ascending = not order.startswith("-")
-            if not self.query.standard_ordering:
-                ascending = not ascending
-
-            name = order.lstrip("+-")
-            if name == "pk":
-                name = opts.pk.name
-
-            try:
-                column = opts.get_field(name).column
-            except FieldDoesNotExist:
-                # `name` is an annotation in $project.
-                column = name
-            column_ordering.append((column, ascending))
+            if hasattr(order, "resolve_expression"):
+                # order is an expression like OrderBy, F, or database function.
+                orderby = order if isinstance(order, OrderBy) else order.asc()
+                orderby = orderby.resolve_expression(self.query, allow_joins=True, reuse=None)
+                ascending = not orderby.descending
+                # If the query is reversed, ascending and descending are inverted.
+                if not self.query.standard_ordering:
+                    ascending = not ascending
+            else:
+                # order is a string like "field" or "field__other_field".
+                orderby, _ = self.find_ordering_name(
+                    order, self.query.get_meta(), default_order=default_order
+                )[0]
+                ascending = not orderby.descending
+            column = orderby.expression.as_mql(self, self.connection)
+            if isinstance(column, dict):
+                raise NotSupportedError("order_by() expression not supported.")
+            # $sort references must not include the dollar sign.
+            column = column.removeprefix("$")
+            # Don't add the same column twice.
+            if column not in columns_seen:
+                columns_seen.add(column)
+                column_ordering.append((column, ascending))
         return column_ordering
 
     @cached_property
diff --git a/django_mongodb/expressions.py b/django_mongodb/expressions.py
@@ -9,8 +9,10 @@
     Col,
     CombinedExpression,
     ExpressionWrapper,
+    F,
     NegatedExpression,
     Ref,
+    ResolvedOuterRef,
     Subquery,
     Value,
     When,
@@ -61,6 +63,10 @@ def expression_wrapper(self, compiler, connection):
     return self.expression.as_mql(compiler, connection)
 
 
+def f(self, compiler, connection):  # noqa: ARG001
+    return f"${self.name}"
+
+
 def negated_expression(self, compiler, connection):
     return {"$not": expression_wrapper(self, compiler, connection)}
 
@@ -102,9 +108,11 @@ def register_expressions():
     Col.as_mql = col
     CombinedExpression.as_mql = combined_expression
     ExpressionWrapper.as_mql = expression_wrapper
+    F.as_mql = f
     NegatedExpression.as_mql = negated_expression
     Query.as_mql = query
     Ref.as_mql = ref
+    ResolvedOuterRef.as_mql = ResolvedOuterRef.as_sql
     Subquery.as_mql = subquery
     When.as_mql = when
     Value.as_mql = value
diff --git a/django_mongodb/features.py b/django_mongodb/features.py
@@ -32,40 +32,25 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "lookup.tests.LookupTests.test_exact_none_transform",
         # "Save with update_fields did not affect any rows."
         "basic.tests.SelectOnSaveTests.test_select_on_save_lying_update",
-        # Lookup in order_by() not supported:
-        # argument of type '<database function>' is not iterable
+        # Order by constant not supported:
+        # AttributeError: 'Field' object has no attribute 'model'
+        "ordering.tests.OrderingTests.test_order_by_constant_value",
+        "expressions.tests.NegatedExpressionTests.test_filter",
+        "expressions_case.tests.CaseExpressionTests.test_order_by_conditional_implicit",
+        # NotSupportedError: order_by() expression not supported.
         "db_functions.comparison.test_coalesce.CoalesceTests.test_ordering",
         "db_functions.tests.FunctionTests.test_nested_function_ordering",
         "db_functions.text.test_length.LengthTests.test_ordering",
         "db_functions.text.test_strindex.StrIndexTests.test_order_by",
-        "expressions.tests.BasicExpressionsTests.test_order_by_exists",
-        "expressions.tests.BasicExpressionsTests.test_order_by_multiline_sql",
         "expressions_case.tests.CaseExpressionTests.test_order_by_conditional_explicit",
         "lookup.tests.LookupQueryingTests.test_lookup_in_order_by",
-        "ordering.tests.OrderingTests.test_default_ordering",
-        "ordering.tests.OrderingTests.test_default_ordering_by_f_expression",
-        "ordering.tests.OrderingTests.test_default_ordering_does_not_affect_group_by",
-        "ordering.tests.OrderingTests.test_order_by_constant_value",
         "ordering.tests.OrderingTests.test_order_by_expr_query_reuse",
         "ordering.tests.OrderingTests.test_order_by_expression_ref",
-        "ordering.tests.OrderingTests.test_order_by_f_expression",
-        "ordering.tests.OrderingTests.test_order_by_f_expression_duplicates",
-        "ordering.tests.OrderingTests.test_order_by_fk_attname",
-        "ordering.tests.OrderingTests.test_order_by_nulls_first",
-        "ordering.tests.OrderingTests.test_order_by_nulls_last",
         "ordering.tests.OrderingTests.test_ordering_select_related_collision",
-        "ordering.tests.OrderingTests.test_order_by_self_referential_fk",
-        "ordering.tests.OrderingTests.test_orders_nulls_first_on_filtered_subquery",
-        "ordering.tests.OrderingTests.test_related_ordering_duplicate_table_reference",
-        "ordering.tests.OrderingTests.test_reverse_ordering_pure",
-        "ordering.tests.OrderingTests.test_reverse_meta_ordering_pure",
-        "ordering.tests.OrderingTests.test_reversed_ordering",
+        "queries.tests.Queries1Tests.test_order_by_related_field_transform",
         "update.tests.AdvancedTests.test_update_ordered_by_inline_m2m_annotation",
         "update.tests.AdvancedTests.test_update_ordered_by_m2m_annotation",
         "update.tests.AdvancedTests.test_update_ordered_by_m2m_annotation_desc",
-        # 'ManyToOneRel' object has no attribute 'column'
-        "m2m_through.tests.M2mThroughTests.test_order_by_relational_field_through_model",
-        "queries.tests.Queries4Tests.test_order_by_reverse_fk",
         # pymongo: ValueError: update cannot be empty
         "update.tests.SimpleTest.test_empty_update_with_inheritance",
         "update.tests.SimpleTest.test_nonempty_update_with_inheritance",
@@ -137,6 +122,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # QuerySet.explain() not implemented:
         # https://github.com/mongodb-labs/django-mongodb/issues/28
         "queries.test_explain.ExplainUnsupportedTests.test_message",
+        # filter() on related model + update() doesn't work.
+        "queries.tests.Queries5Tests.test_ticket9848",
     }
     # $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
     _django_test_expected_failures_bitwise = {
@@ -320,6 +307,7 @@ def django_test_expected_failures(self):
             "expressions.tests.BasicExpressionsTests.test_boolean_expression_in_Q",
             "expressions.tests.BasicExpressionsTests.test_case_in_filter_if_boolean_output_field",
             "expressions.tests.BasicExpressionsTests.test_exists_in_filter",
+            "expressions.tests.BasicExpressionsTests.test_order_by_exists",
             "expressions.tests.BasicExpressionsTests.test_subquery",
             "expressions.tests.ExistsTests.test_filter_by_empty_exists",
             "expressions.tests.ExistsTests.test_negated_empty_exists",
@@ -438,6 +426,7 @@ def django_test_expected_failures(self):
             "expressions.tests.FieldTransformTests.test_month_aggregation",
             "expressions_case.tests.CaseDocumentationExamples.test_conditional_aggregation_example",
             "model_fields.test_jsonfield.TestQuerying.test_ordering_grouping_by_count",
+            "ordering.tests.OrderingTests.test_default_ordering_does_not_affect_group_by",
             "queries.tests.Queries1Tests.test_ticket_20250",
             "queries.tests.ValuesQuerysetTests.test_named_values_list_expression_with_default_alias",
         },
@@ -514,6 +503,11 @@ def django_test_expected_failures(self):
             "queries.tests.ValuesQuerysetTests.test_named_values_list_without_fields",
             "select_related.tests.SelectRelatedTests.test_select_related_with_extra",
         },
+        "Ordering a QuerySet by null_first/nulls_last is not supported on MongoDB.": {
+            "ordering.tests.OrderingTests.test_order_by_nulls_first",
+            "ordering.tests.OrderingTests.test_order_by_nulls_last",
+            "ordering.tests.OrderingTests.test_orders_nulls_first_on_filtered_subquery",
+        },
         "QuerySet.update() crash: Unrecognized expression '$count'": {
             "update.tests.AdvancedTests.test_update_annotated_multi_table_queryset",
         },
@@ -529,6 +523,7 @@ def django_test_expected_failures(self):
             "delete_regress.tests.DeleteLockingTest.test_concurrent_delete",
             "expressions.tests.BasicExpressionsTests.test_annotate_values_filter",
             "expressions.tests.BasicExpressionsTests.test_filtering_on_rawsql_that_is_boolean",
+            "expressions.tests.BasicExpressionsTests.test_order_by_multiline_sql",
             "model_fields.test_jsonfield.TestQuerying.test_key_sql_injection_escape",
             "model_fields.test_jsonfield.TestQuerying.test_key_transform_raw_expression",
             "model_fields.test_jsonfield.TestQuerying.test_nested_key_transform_raw_expression",
@@ -617,18 +612,6 @@ def django_test_expected_failures(self):
         "Randomized ordering isn't supported by MongoDB.": {
             "ordering.tests.OrderingTests.test_random_ordering",
         },
-        # https://github.com/mongodb-labs/django-mongodb/issues/34
-        "Ordering can't span tables on MongoDB": {
-            "queries.tests.ConditionalTests.test_infinite_loop",
-            "queries.tests.NullableRelOrderingTests.test_join_already_in_query",
-            "queries.tests.Queries1Tests.test_order_by_related_field_transform",
-            "queries.tests.Queries1Tests.test_ticket7181",
-            "queries.tests.Queries1Tests.test_tickets_2076_7256",
-            "queries.tests.Queries1Tests.test_tickets_2874_3002",
-            "queries.tests.Queries5Tests.test_ordering",
-            "queries.tests.Queries5Tests.test_ticket9848",
-            "queries.tests.Ticket14056Tests.test_ticket_14056",
-        },
         "Queries without a collection aren't supported on MongoDB.": {
             "queries.test_q.QCheckTests",
             "queries.test_query.TestQueryNoModel",
diff --git a/django_mongodb/operations.py b/django_mongodb/operations.py
@@ -6,9 +6,9 @@
 
 from bson.decimal128 import Decimal128
 from django.conf import settings
-from django.db import DataError
+from django.db import DataError, NotSupportedError
 from django.db.backends.base.operations import BaseDatabaseOperations
-from django.db.models.expressions import Combinable
+from django.db.models.expressions import Combinable, OrderBy
 from django.utils import timezone
 from django.utils.regex_helper import _lazy_re_compile
 
@@ -140,6 +140,11 @@ def convert_uuidfield_value(self, value, expression, connection):
             value = uuid.UUID(value)
         return value
 
+    def check_expression_support(self, expression):
+        if isinstance(expression, OrderBy) and (expression.nulls_first or expression.nulls_last):
+            option = "null_first" if expression.nulls_first else "nulls_last"
+            raise NotSupportedError(f"Ordering a QuerySet by {option} is not supported on MongoDB.")
+
     def combine_expression(self, connector, sub_expressions):
         lhs, rhs = sub_expressions
         if connector == Combinable.BITLEFTSHIFT:
