diff --git a/docs/changelog.rst b/docs/changelog.rst index 1556266b..14dfb8d0 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 - 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 diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py index 80c0661b..9cf9e2ae 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) @@ -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): diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 6ad08617..33ab6e2a 100644 --- a/mongoengine/queryset/base.py +++ b/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 diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index 8b5872f8..e2db8f0d 100644 --- a/mongoengine/queryset/queryset.py +++ b/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 diff --git a/tests/document/test_indexes.py b/tests/document/test_indexes.py index 4f728c2d..45d1cd23 100644 --- a/tests/document/test_indexes.py +++ b/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) diff --git a/tests/test_connection.py b/tests/test_connection.py index 56bc22cd..b57d4597 100644 --- a/tests/test_connection.py +++ b/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):