Merge pull request #2243 from bagerard/fix_count_documents_deprecation

Fix count documents deprecation

docs/changelog.rst

@@ -6,6 +6,9 @@ Changelog
 Development
 ===========
 - (Fill this out as you fix issues and develop your features).
+- When using pymongo >= 3.7, make use of Collection.count_documents instead of Collection.count
+    and Cursor.count that got deprecated in pymongo >= 3.7.
+    This should have a negative impact on performance of count see Issue #2219
 - Fix a bug that made the queryset drop the read_preference after clone().
 - Fix the behavior of Doc.objects.limit(0) which should return all documents (similar to mongodb) #2311
 

mongoengine/pymongo_support.py

@@ -2,6 +2,7 @@
 Helper functions, constants, and types to aid with PyMongo v2.7 - v3.x support.
 """
 import pymongo
+from pymongo.errors import OperationFailure
 
 _PYMONGO_37 = (3, 7)
 
@@ -10,13 +11,41 @@ PYMONGO_VERSION = tuple(pymongo.version_tuple[:2])
 IS_PYMONGO_GTE_37 = PYMONGO_VERSION >= _PYMONGO_37
 
 
-def count_documents(collection, filter):
-    """Pymongo>3.7 deprecates count in favour of count_documents"""
+def count_documents(
+    collection, filter, skip=None, limit=None, hint=None, collation=None
+):
+    """Pymongo>3.7 deprecates count in favour of count_documents
+    """
+    if limit == 0:
+        return 0  # Pymongo raises an OperationFailure if called with limit=0
+
+    kwargs = {}
+    if skip is not None:
+        kwargs["skip"] = skip
+    if limit is not None:
+        kwargs["limit"] = limit
+    if hint not in (-1, None):
+        kwargs["hint"] = hint
+    if collation is not None:
+        kwargs["collation"] = collation
+
+    # count_documents appeared in pymongo 3.7
     if IS_PYMONGO_GTE_37:
-        return collection.count_documents(filter)
-    else:
-        count = collection.find(filter).count()
-    return count
+        try:
+            return collection.count_documents(filter=filter, **kwargs)
+        except OperationFailure:
+            # OperationFailure - accounts for some operators that used to work
+            # 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
+            pass
+
+    cursor = collection.find(filter)
+    for option, option_value in kwargs.items():
+        cursor_method = getattr(cursor, option)
+        cursor = cursor_method(option_value)
+    with_limit_and_skip = "skip" in kwargs or "limit" in kwargs
+    return cursor.count(with_limit_and_skip=with_limit_and_skip)
 
 
 def list_collection_names(db, include_system_collections=False):

mongoengine/queryset/base.py

@@ -29,6 +29,7 @@ from mongoengine.errors import (
     NotUniqueError,
     OperationError,
 )
+from mongoengine.pymongo_support import count_documents
 from mongoengine.queryset import transform
 from mongoengine.queryset.field_list import QueryFieldList
 from mongoengine.queryset.visitor import Q, QNode
@@ -405,6 +406,8 @@ class BaseQuerySet:
             :meth:`skip` that has been applied to this cursor into account when
             getting the count
         """
+        # mimic the fact that setting .limit(0) in pymongo sets no limit
+        # https://docs.mongodb.com/manual/reference/method/cursor.limit/#zero-value
         if (
             self._limit == 0
             and with_limit_and_skip is False
@@ -412,7 +415,27 @@ class BaseQuerySet:
             or self._empty
         ):
             return 0
-        count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
+
+        kwargs = (
+            {"limit": self._limit, "skip": self._skip} if with_limit_and_skip else {}
+        )
+
+        if self._limit == 0:
+            # mimic the fact that historically .limit(0) sets no limit
+            kwargs.pop("limit", None)
+
+        if self._hint not in (-1, None):
+            kwargs["hint"] = self._hint
+
+        if self._collation:
+            kwargs["collation"] = self._collation
+
+        count = count_documents(
+            collection=self._cursor.collection,
+            filter=self._cursor._Cursor__spec,
+            **kwargs
+        )
+
         self._cursor_obj = None
         return count
 

mongoengine/queryset/queryset.py

@@ -144,6 +144,7 @@ class QuerySet(BaseQuerySet):
             return super().count(with_limit_and_skip)
 
         if self._len is None:
+            # cache the length
             self._len = super().count(with_limit_and_skip)
 
         return self._len

tests/document/test_indexes.py

@@ -551,8 +551,9 @@ class TestIndexes(unittest.TestCase):
         assert 5 == query_result.count()
 
         incorrect_collation = {"arndom": "wrdo"}
-        with pytest.raises(OperationFailure):
+        with pytest.raises(OperationFailure) as exc_info:
             BlogPost.objects.collation(incorrect_collation).count()
+        assert "Missing expected field" in str(exc_info.value)
 
         query_result = BlogPost.objects.collation({}).order_by("name")
         assert [x.name for x in query_result] == sorted(names)

tests/test_connection.py

@@ -282,7 +282,7 @@ class ConnectionTest(unittest.TestCase):
         # database won't exist until we save a document
         some_document.save()
         assert conn.get_default_database().name == "mongoenginetest"
-        assert conn.database_names()[0] == "mongoenginetest"
+        assert conn.list_database_names()[0] == "mongoenginetest"
 
     @require_mongomock
     def test_connect_with_host_list(self):