Merge branch 'master' into array_filters

Author: idoshr

docs/changelog.rst

@@ -9,6 +9,7 @@ Development
 - (Fill this out as you fix issues and develop your features).
 - Fix for uuidRepresentation not read when provided in URI #2741
 - Add option to user array_filters https://www.mongodb.com/docs/manual/reference/operator/update/positional-filtered/
+- Fix combination of __raw__ and mongoengine syntax #2773
 - Add tests against MongoDB 6.0 and MongoDB 7.0 in the pipeline
 - Fix validate() not being called when inheritance is used in EmbeddedDocument and validate is overriden #2784
 - Add support for readPreferenceTags in connection parameters #2644
@@ -17,6 +18,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
 =================

docs/guide/querying.rst

@@ -605,6 +605,7 @@ There are several different "modifiers" that you may use with these methods:
 * ``add_to_set`` -- add value to a list only if its not in the list already
 * ``rename`` -- rename the key name
 
+.. _need to add upsert=True: http://docs.mongodb.org/manual/reference/operator/update/setOnInsert
 .. _depending on the value: http://docs.mongodb.org/manual/reference/operator/update/pop/
 
 The syntax for atomic updates is similar to the querying syntax, but the

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

@@ -1119,7 +1119,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
@@ -1129,11 +1129,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

mongoengine/queryset/transform.py

@@ -62,14 +62,23 @@
 )
 
 
+def handle_raw_query(value, mongo_query):
+    """Combine a raw query with an existing one"""
+    for op, v in value.items():
+        if op not in mongo_query:
+            mongo_query[op] = v
+        elif op in mongo_query and isinstance(mongo_query[op], dict):
+            mongo_query[op].update(v)
+
+
 # TODO make this less complex
 def query(_doc_cls=None, **kwargs):
     """Transform a query from Django-style format to Mongo format."""
     mongo_query = {}
     merge_query = defaultdict(list)
     for key, value in sorted(kwargs.items()):
         if key == "__raw__":
-            mongo_query.update(value)
+            handle_raw_query(value, mongo_query)
             continue
 
         parts = key.rsplit("__")
@@ -234,7 +243,7 @@ def update(_doc_cls=None, **update):
 
     for key, value in update.items():
         if key == "__raw__":
-            mongo_update.update(value)
+            handle_raw_query(value, mongo_update)
             continue
 
         parts = key.split("__")

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,
@@ -2328,6 +2329,46 @@ class BlogPost(Document):
         post.reload()
         assert post.slug == "When test test it"
 
+    def test_combination_of_mongoengine_and__raw__(self):
+        """Ensure that the '__raw__' update/query works in combination with mongoengine syntax correctly."""
+
+        class BlogPost(Document):
+            slug = StringField()
+            foo = StringField()
+            tags = ListField(StringField())
+
+        BlogPost.drop_collection()
+
+        post = BlogPost(slug="test", foo="bar")
+        post.save()
+
+        BlogPost.objects(slug="test").update(
+            foo="baz",
+            __raw__={"$set": {"slug": "test test"}},
+        )
+        post.reload()
+        assert post.slug == "test test"
+        assert post.foo == "baz"
+
+        assert BlogPost.objects(foo="baz", __raw__={"slug": "test test"}).count() == 1
+        assert (
+            BlogPost.objects(foo__ne="bar", __raw__={"slug": {"$ne": "test"}}).count()
+            == 1
+        )
+        assert (
+            BlogPost.objects(foo="baz", __raw__={"slug": {"$ne": "test test"}}).count()
+            == 0
+        )
+        assert (
+            BlogPost.objects(foo__ne="baz", __raw__={"slug": "test test"}).count() == 0
+        )
+        assert (
+            BlogPost.objects(
+                foo__ne="baz", __raw__={"slug": {"$ne": "test test"}}
+            ).count()
+            == 0
+        )
+
     def test_add_to_set_each(self):
         class Item(Document):
             name = StringField(required=True)
@@ -2739,6 +2780,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"""