Merge branch 'array_filters' of https://github.com/idoshr/mongoengine into array_filters

# Conflicts:
#	docs/changelog.rst

Author: idocyabra

docs/changelog.rst

@@ -8,10 +8,17 @@ 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
-- Add option to user array_filters https://www.mongodb.com/docs/manual/reference/operator/update/positional-filtered/
+- Use estimated_documents_count OR documents_count when count is called, based on the query #2529
+- Fix no_dereference context manager which wasn't turning off auto-dereferencing correctly in some cases #2788
+- 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

@@ -540,7 +540,7 @@ data. To turn off dereferencing of the results of a query use
 You can also turn off all dereferencing for a fixed period by using the
 :class:`~mongoengine.context_managers.no_dereference` context manager::
 
-    with no_dereference(Post) as Post:
+    with no_dereference(Post):
         post = Post.objects.first()
         assert(isinstance(post.author, DBRef))
 
@@ -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

@@ -1,3 +1,4 @@
+import threading
 from contextlib import contextmanager
 
 from pymongo.read_concern import ReadConcern
@@ -15,9 +16,29 @@
     "query_counter",
     "set_write_concern",
     "set_read_write_concern",
+    "no_dereferencing_active_for_class",
 )
 
 
+thread_locals = threading.local()
+thread_locals.no_dereferencing_class = {}
+
+
+def no_dereferencing_active_for_class(cls):
+    return cls in thread_locals.no_dereferencing_class
+
+
+def _register_no_dereferencing_for_class(cls):
+    thread_locals.no_dereferencing_class.setdefault(cls, 0)
+    thread_locals.no_dereferencing_class[cls] += 1
+
+
+def _unregister_no_dereferencing_for_class(cls):
+    thread_locals.no_dereferencing_class[cls] -= 1
+    if thread_locals.no_dereferencing_class[cls] == 0:
+        thread_locals.no_dereferencing_class.pop(cls)
+
+
 class switch_db:
     """switch_db alias context manager.
 
@@ -107,7 +128,7 @@ class no_dereference:
     Turns off all dereferencing in Documents for the duration of the context
     manager::
 
-        with no_dereference(Group) as Group:
+        with no_dereference(Group):
             Group.objects.find()
     """
 
@@ -130,15 +151,17 @@ def __init__(self, cls):
 
     def __enter__(self):
         """Change the objects default and _auto_dereference values."""
+        _register_no_dereferencing_for_class(self.cls)
+
         for field in self.deref_fields:
             self.cls._fields[field]._auto_dereference = False
-        return self.cls
 
     def __exit__(self, t, value, traceback):
         """Reset the default and _auto_dereference values."""
+        _unregister_no_dereferencing_for_class(self.cls)
+
         for field in self.deref_fields:
             self.cls._fields[field]._auto_dereference = True
-        return self.cls
 
 
 class no_sub_classes:

mongoengine/pymongo_support.py

@@ -37,7 +37,12 @@ def count_documents(
     # count_documents appeared in pymongo 3.7
     if PYMONGO_VERSION >= (3, 7):
         try:
-            return collection.count_documents(filter=filter, **kwargs)
+            if not filter and set(kwargs) <= {"max_time_ms"}:
+                # when no filter is provided, estimated_document_count
+                # is a lot faster as it uses the collection metadata
+                return collection.estimated_document_count(**kwargs)
+            else:
+                return collection.count_documents(filter=filter, **kwargs)
         except OperationFailure as err:
             if PYMONGO_VERSION >= (4,):
                 raise
@@ -46,15 +51,10 @@ def count_documents(
             # with .count but are no longer working with count_documents (i.e $geoNear, $near, and $nearSphere)
             # fallback to deprecated Cursor.count
             # Keeping this should be reevaluated the day pymongo removes .count entirely
-            message = str(err)
-            if not (
-                "not allowed in this context" in message
-                and (
-                    "$where" in message
-                    or "$geoNear" in message
-                    or "$near" in message
-                    or "$nearSphere" in message
-                )
+            if (
+                "$geoNear, $near, and $nearSphere are not allowed in this context"
+                not in str(err)
+                and "$where is not allowed in this context" not in str(err)
             ):
                 raise
 

mongoengine/queryset/base.py

@@ -17,6 +17,7 @@
 from mongoengine.common import _import_class
 from mongoengine.connection import get_db
 from mongoengine.context_managers import (
+    no_dereferencing_active_for_class,
     set_read_write_concern,
     set_write_concern,
     switch_db,
@@ -51,9 +52,6 @@ class BaseQuerySet:
     providing :class:`~mongoengine.Document` objects as the results.
     """
 
-    __dereference = False
-    _auto_dereference = True
-
     def __init__(self, document, collection):
         self._document = document
         self._collection_obj = collection
@@ -74,6 +72,9 @@ def __init__(self, document, collection):
         self._as_pymongo = False
         self._search_text = None
 
+        self.__dereference = False
+        self.__auto_dereference = True
+
         # If inheritance is allowed, only return instances and instances of
         # subclasses of the class being used
         if document._meta.get("allow_inheritance") is True:
@@ -801,7 +802,7 @@ def clone(self):
         return self._clone_into(self.__class__(self._document, self._collection_obj))
 
     def _clone_into(self, new_qs):
-        """Copy all of the relevant properties of this queryset to
+        """Copy all the relevant properties of this queryset to
         a new queryset (which has to be an instance of
         :class:`~mongoengine.queryset.base.BaseQuerySet`).
         """
@@ -831,7 +832,6 @@ def _clone_into(self, new_qs):
             "_empty",
             "_hint",
             "_collation",
-            "_auto_dereference",
             "_search_text",
             "_max_time_ms",
             "_comment",
@@ -842,6 +842,8 @@ def _clone_into(self, new_qs):
             val = getattr(self, prop)
             setattr(new_qs, prop, copy.copy(val))
 
+        new_qs.__auto_dereference = self._BaseQuerySet__auto_dereference
+
         if self._cursor_obj:
             new_qs._cursor_obj = self._cursor_obj.clone()
 
@@ -1117,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
@@ -1127,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
@@ -1747,10 +1757,15 @@ def _dereference(self):
             self.__dereference = _import_class("DeReference")()
         return self.__dereference
 
+    @property
+    def _auto_dereference(self):
+        should_deref = not no_dereferencing_active_for_class(self._document)
+        return should_deref and self.__auto_dereference
+
     def no_dereference(self):
         """Turn off any dereferencing for the results of this queryset."""
         queryset = self.clone()
-        queryset._auto_dereference = False
+        queryset.__auto_dereference = False
         return queryset
 
     # Helper Functions

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/document/test_indexes.py

@@ -9,7 +9,6 @@
 from mongoengine.connection import get_db
 from mongoengine.mongodb_support import (
     MONGODB_42,
-    MONGODB_70,
     get_mongodb_version,
 )
 from mongoengine.pymongo_support import PYMONGO_VERSION
@@ -451,89 +450,29 @@ class Test(Document):
         # the documents returned might have more keys in that here.
         query_plan = Test.objects(id=obj.id).exclude("a").explain()
         assert (
-            query_plan.get("queryPlanner")
-            .get("winningPlan")
-            .get("inputStage")
-            .get("stage")
-            == "IDHACK"
+            query_plan["queryPlanner"]["winningPlan"]["inputStage"]["stage"] == "IDHACK"
         )
 
         query_plan = Test.objects(id=obj.id).only("id").explain()
         assert (
-            query_plan.get("queryPlanner")
-            .get("winningPlan")
-            .get("inputStage")
-            .get("stage")
-            == "IDHACK"
+            query_plan["queryPlanner"]["winningPlan"]["inputStage"]["stage"] == "IDHACK"
         )
 
         mongo_db = get_mongodb_version()
         query_plan = Test.objects(a=1).only("a").exclude("id").explain()
-        if mongo_db < MONGODB_70:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("inputStage")
-                .get("stage")
-                == "IXSCAN"
-            )
-        else:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("queryPlan")
-                .get("inputStage")
-                .get("stage")
-                == "IXSCAN"
-            )
+        assert (
+            query_plan["queryPlanner"]["winningPlan"]["inputStage"]["stage"] == "IXSCAN"
+        )
 
         PROJECTION_STR = "PROJECTION" if mongo_db < MONGODB_42 else "PROJECTION_COVERED"
-        if mongo_db < MONGODB_70:
-            assert (
-                query_plan.get("queryPlanner").get("winningPlan").get("stage")
-                == PROJECTION_STR
-            )
-        else:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("queryPlan")
-                .get("stage")
-                == PROJECTION_STR
-            )
+        assert query_plan["queryPlanner"]["winningPlan"]["stage"] == PROJECTION_STR
 
         query_plan = Test.objects(a=1).explain()
-        if mongo_db < MONGODB_70:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("inputStage")
-                .get("stage")
-                == "IXSCAN"
-            )
-        else:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("queryPlan")
-                .get("inputStage")
-                .get("stage")
-                == "IXSCAN"
-            )
-
-        if mongo_db < MONGODB_70:
-            assert (
-                query_plan.get("queryPlanner").get("winningPlan").get("stage")
-                == "FETCH"
-            )
-        else:
-            assert (
-                query_plan.get("queryPlanner")
-                .get("winningPlan")
-                .get("queryPlan")
-                .get("stage")
-                == "FETCH"
-            )
+        assert (
+            query_plan["queryPlanner"]["winningPlan"]["inputStage"]["stage"] == "IXSCAN"
+        )
+
+        assert query_plan.get("queryPlanner").get("winningPlan").get("stage") == "FETCH"
 
     def test_index_on_id(self):
         class BlogPost(Document):