diff --git a/docs/changelog.rst b/docs/changelog.rst index b8e6ae56..d11fd347 100644 --- a/docs/changelog.rst +++ b/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 ================= diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py index 1fea9525..38332c13 100644 --- a/mongoengine/pymongo_support.py +++ b/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 = {} - if skip is not None: - kwargs["skip"] = skip - if limit is not None: - kwargs["limit"] = limit - if collation is not None: - kwargs["collation"] = collation - if hint not in (-1, None): - kwargs["hint"] = hint + 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 + + 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: - cursor = cursor.limit(limit) - if skip: - cursor = cursor.skip(skip) - if hint != -1: - cursor = cursor.hint(hint) + for option, option_value in kwargs.items(): + cursor_method = getattr(cursor, option) + cursor = cursor_method(option_value) count = cursor.count() return count diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 125480a7..41b394b4 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -413,18 +413,11 @@ class BaseQuerySet(object): if self._collation: kwargs["collation"] = self._collation - try: - count = count_documents( - collection=self._cursor.collection, - filter=self._cursor._Cursor__spec, - **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) + count = count_documents( + collection=self._cursor.collection, + filter=self._cursor._Cursor__spec, + **kwargs + ) self._cursor_obj = None return count diff --git a/tests/document/test_indexes.py b/tests/document/test_indexes.py index be857b59..801473b1 100644 --- a/tests/document/test_indexes.py +++ b/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)