Merge pull request #2789 from bagerard/add_raw__order_by

add __raw__ to order_by

Author: bagerard

docs/changelog.rst

@@ -17,6 +17,7 @@ Development
 - BREAKING CHANGE: no_dereference context manager no longer returns the class in __enter__ #2788
     as it was useless and making it look like it was returning a different class although it was the same.
     Thus, it must be called like `with no_dereference(User):` and no longer `with no_dereference(User) as ...:`
+- Added __raw__ to :meth:`~mongoengine.Queryset.order_by()` to allow to provide raw pymongo 'sort' argument and get around some of the limitations #2783
 
 Changes in 0.27.0
 =================

mongoengine/context_managers.py

@@ -16,6 +16,7 @@
     "query_counter",
     "set_write_concern",
     "set_read_write_concern",
+    "no_dereferencing_active_for_class",
 )
 
 

mongoengine/queryset/base.py

@@ -1113,7 +1113,7 @@ def all_fields(self):
         )
         return queryset
 
-    def order_by(self, *keys):
+    def order_by(self, *keys, __raw__=None):
         """Order the :class:`~mongoengine.queryset.QuerySet` by the given keys.
 
         The order may be specified by prepending each of the keys by a "+" or
@@ -1123,11 +1123,19 @@ def order_by(self, *keys):
 
         :param keys: fields to order the query results by; keys may be
             prefixed with "+" or a "-" to determine the ordering direction.
+        :param __raw__: a raw pymongo "sort" argument (provided as a list of (key, direction))
+            see 'key_or_list' in `pymongo.cursor.Cursor.sort doc <https://pymongo.readthedocs.io/en/stable/api/pymongo/cursor.html#pymongo.cursor.Cursor.sort>`.
+            If both keys and __raw__ are provided, an exception is raised
         """
-        queryset = self.clone()
+        if __raw__ and keys:
+            raise OperationError("Can not use both keys and __raw__ with order_by() ")
 
+        queryset = self.clone()
         old_ordering = queryset._ordering
-        new_ordering = queryset._get_order_by(keys)
+        if __raw__:
+            new_ordering = __raw__
+        else:
+            new_ordering = queryset._get_order_by(keys)
 
         if queryset._cursor_obj:
             # If a cursor object has already been created, apply the sort to it

tests/queryset/test_queryset.py

@@ -17,6 +17,7 @@
     MONGODB_36,
     get_mongodb_version,
 )
+from mongoengine.pymongo_support import PYMONGO_VERSION
 from mongoengine.queryset import (
     DoesNotExist,
     MultipleObjectsReturned,
@@ -2720,6 +2721,44 @@ def test_order_by_chaining(self):
         ages = [p.age for p in qs]
         assert ages == [40, 30, 20]
 
+    def test_order_by_using_raw(self):
+        person_a = self.Person(name="User A", age=20)
+        person_a.save()
+        person_b = self.Person(name="User B", age=30)
+        person_b.save()
+        person_c = self.Person(name="User B", age=25)
+        person_c.save()
+        person_d = self.Person(name="User C", age=40)
+        person_d.save()
+
+        qs = self.Person.objects.order_by(__raw__=[("name", pymongo.DESCENDING)])
+        assert qs._ordering == [("name", pymongo.DESCENDING)]
+        names = [p.name for p in qs]
+        assert names == ["User C", "User B", "User B", "User A"]
+
+        names = [
+            (p.name, p.age)
+            for p in self.Person.objects.order_by(__raw__=[("name", pymongo.ASCENDING)])
+        ]
+        assert names == [("User A", 20), ("User B", 30), ("User B", 25), ("User C", 40)]
+
+        if PYMONGO_VERSION >= (4, 4):
+            # Pymongo >= 4.4 allow to mix single key with tuples inside the list
+            qs = self.Person.objects.order_by(
+                __raw__=["name", ("age", pymongo.ASCENDING)]
+            )
+            names = [(p.name, p.age) for p in qs]
+            assert names == [
+                ("User A", 20),
+                ("User B", 25),
+                ("User B", 30),
+                ("User C", 40),
+            ]
+
+    def test_order_by_using_raw_and_keys_raises_exception(self):
+        with pytest.raises(OperationError):
+            self.Person.objects.order_by("-name", __raw__=[("age", pymongo.ASCENDING)])
+
     def test_confirm_order_by_reference_wont_work(self):
         """Ordering by reference is not possible.  Use map / reduce.. or
         denormalise"""