finalize code related to count_documents migration

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
 
 Changes in 0.19.1
 =================

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)
 
@@ -16,25 +17,28 @@ def count_documents(collection, filter, skip=None, limit=None, hint=None, collat
     if limit == 0:
         return 0  # Pymongo raises an OperationFailure if called with limit=0
 
-    if IS_PYMONGO_GTE_37:
+    kwargs = {}
-        kwargs = {}
+    if skip is not None:
-        if skip is not None:
+        kwargs["skip"] = skip
-            kwargs["skip"] = skip
+    if limit is not None:
-        if limit is not None:
+        kwargs["limit"] = limit
-            kwargs["limit"] = limit
+    if hint not in (-1, None):
-        if collation is not None:
+        kwargs["hint"] = hint
-            kwargs["collation"] = collation
+    if collation is not None:
-        if hint not in (-1, None):
+        kwargs["collation"] = collation
-            kwargs["hint"] = hint
+
+    try:
         return collection.count_documents(filter=filter, **kwargs)
-    else:
+    except (AttributeError, OperationFailure) as ex:
+        # AttributeError - count_documents appeared in pymongo 3.7
+        # 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
         cursor = collection.find(filter)
-        if limit:
+        for option, option_value in kwargs.items():
-            cursor = cursor.limit(limit)
+            cursor_method = getattr(cursor, option)
-        if skip:
+            cursor = cursor_method(option_value)
-            cursor = cursor.skip(skip)
-        if hint != -1:
-            cursor = cursor.hint(hint)
         count = cursor.count()
     return count
 

mongoengine/queryset/base.py

@@ -413,18 +413,11 @@ class BaseQuerySet(object):
         if self._collation:
             kwargs["collation"] = self._collation
 
-        try:
+        count = count_documents(
-            count = count_documents(
+            collection=self._cursor.collection,
-                collection=self._cursor.collection,
+            filter=self._cursor._Cursor__spec,
-                filter=self._cursor._Cursor__spec,
+            **kwargs
-                **kwargs
+        )
-            )
-        except 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
-            count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
 
         self._cursor_obj = None
         return count

tests/document/test_indexes.py

@@ -552,8 +552,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)