From 928770c43a892e7e6a01e5a85f5d90a5b9a63c72 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Fri, 20 Dec 2019 21:50:00 +0100
Subject: [PATCH 1/7] switching  to count_documents

---
 mongoengine/pymongo_support.py   | 28 ++++++++++++++++++++++++----
 mongoengine/queryset/base.py     | 32 +++++++++++++++++++++++++++++++-
 mongoengine/queryset/queryset.py |  1 +
 tests/test_connection.py         |  2 +-
 4 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py
index 80c0661b..1fea9525 100644
--- a/mongoengine/pymongo_support.py
+++ b/mongoengine/pymongo_support.py
@@ -10,12 +10,32 @@ 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
+
     if IS_PYMONGO_GTE_37:
-        return collection.count_documents(filter)
+        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
+        return collection.count_documents(filter=filter, **kwargs)
     else:
-        count = collection.find(filter).count()
+        cursor = collection.find(filter)
+        if limit:
+            cursor = cursor.limit(limit)
+        if skip:
+            cursor = cursor.skip(skip)
+        if hint != -1:
+            cursor = cursor.hint(hint)
+        count = cursor.count()
     return count
 
 
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 50cb37ac..d7b4007e 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -12,6 +12,7 @@ import pymongo.errors
 from pymongo.collection import ReturnDocument
 from pymongo.common import validate_read_preference
 import six
+from pymongo.errors import OperationFailure
 from six import iteritems
 
 from mongoengine import signals
@@ -26,6 +27,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
@@ -392,9 +394,37 @@ class BaseQuerySet(object):
             :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 or self._none:
             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
+
+        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
+            count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
+
         self._cursor_obj = None
         return count
 
diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py
index 4ba62d46..cc1891f6 100644
--- a/mongoengine/queryset/queryset.py
+++ b/mongoengine/queryset/queryset.py
@@ -146,6 +146,7 @@ class QuerySet(BaseQuerySet):
             return super(QuerySet, self).count(with_limit_and_skip)
 
         if self._len is None:
+            # cache the length
             self._len = super(QuerySet, self).count(with_limit_and_skip)
 
         return self._len
diff --git a/tests/test_connection.py b/tests/test_connection.py
index e40a6994..542da4f0 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):

From f93f9406ee3646eb3bac8fbeed6aa502ffd83a31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Sun, 5 Jan 2020 21:08:20 +0100
Subject: [PATCH 2/7] improve doc next to code

---
 mongoengine/queryset/base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index d7b4007e..125480a7 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -423,6 +423,7 @@ class BaseQuerySet(object):
             # 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

From 60c42dddd5a44623b696d8ce6eff7be7f955e796 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Sun, 5 Jan 2020 22:29:13 +0100
Subject: [PATCH 3/7] finalize code related to count_documents migration

---
 docs/changelog.rst             |  3 +++
 mongoengine/pymongo_support.py | 38 +++++++++++++++++++---------------
 mongoengine/queryset/base.py   | 17 +++++----------
 tests/document/test_indexes.py |  3 ++-
 4 files changed, 31 insertions(+), 30 deletions(-)

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)

From 84f3dce4920e4e212a485dcbe9f1fa4cda4f7daa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Sun, 5 Jan 2020 22:50:19 +0100
Subject: [PATCH 4/7] fix flake8 findings

---
 mongoengine/pymongo_support.py | 2 +-
 mongoengine/queryset/base.py   | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py
index 38332c13..d9c5ee27 100644
--- a/mongoengine/pymongo_support.py
+++ b/mongoengine/pymongo_support.py
@@ -29,7 +29,7 @@ def count_documents(collection, filter, skip=None, limit=None, hint=None, collat
 
     try:
         return collection.count_documents(filter=filter, **kwargs)
-    except (AttributeError, OperationFailure) as ex:
+    except (AttributeError, OperationFailure):
         # 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)
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 41b394b4..303271f5 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -12,7 +12,6 @@ import pymongo.errors
 from pymongo.collection import ReturnDocument
 from pymongo.common import validate_read_preference
 import six
-from pymongo.errors import OperationFailure
 from six import iteritems
 
 from mongoengine import signals

From e64a7a94481f02764b16ff9ce448489fad2fc321 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Tue, 7 Jan 2020 21:44:04 +0100
Subject: [PATCH 5/7] reformat with latest black

---
 mongoengine/pymongo_support.py | 4 +++-
 mongoengine/queryset/base.py   | 2 +-
 tests/document/test_indexes.py | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py
index d9c5ee27..3aef4e09 100644
--- a/mongoengine/pymongo_support.py
+++ b/mongoengine/pymongo_support.py
@@ -11,7 +11,9 @@ PYMONGO_VERSION = tuple(pymongo.version_tuple[:2])
 IS_PYMONGO_GTE_37 = PYMONGO_VERSION >= _PYMONGO_37
 
 
-def count_documents(collection, filter, skip=None, limit=None, hint=None, collation=None):
+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:
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
index 303271f5..ae45297d 100644
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -404,7 +404,7 @@ class BaseQuerySet(object):
 
         if self._limit == 0:
             # mimic the fact that historically .limit(0) sets no limit
-            kwargs.pop('limit', None)
+            kwargs.pop("limit", None)
 
         if self._hint not in (-1, None):
             kwargs["hint"] = self._hint
diff --git a/tests/document/test_indexes.py b/tests/document/test_indexes.py
index 801473b1..5133b007 100644
--- a/tests/document/test_indexes.py
+++ b/tests/document/test_indexes.py
@@ -554,7 +554,7 @@ class TestIndexes(unittest.TestCase):
         incorrect_collation = {"arndom": "wrdo"}
         with pytest.raises(OperationFailure) as exc_info:
             BlogPost.objects.collation(incorrect_collation).count()
-        assert 'Missing expected field' in str(exc_info.value)
+        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)

From 2fa48cd9e5207ad753419ace20dc0c615c7481cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Tue, 7 Jan 2020 22:24:55 +0100
Subject: [PATCH 6/7] fix for pymongo < 3.7

---
 mongoengine/pymongo_support.py | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py
index 3aef4e09..284efc2f 100644
--- a/mongoengine/pymongo_support.py
+++ b/mongoengine/pymongo_support.py
@@ -29,20 +29,22 @@ def count_documents(
     if collation is not None:
         kwargs["collation"] = collation
 
-    try:
-        return collection.count_documents(filter=filter, **kwargs)
-    except (AttributeError, OperationFailure):
-        # 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)
-        for option, option_value in kwargs.items():
-            cursor_method = getattr(cursor, option)
-            cursor = cursor_method(option_value)
-        count = cursor.count()
-    return count
+    # count_documents appeared in pymongo 3.7
+    if IS_PYMONGO_GTE_37:
+        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)
+    return cursor.count()
 
 
 def list_collection_names(db, include_system_collections=False):

From 412bed0f6d7aa42a5fa855217b0d2a3b1721ef27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Sun, 12 Jan 2020 11:04:05 +0100
Subject: [PATCH 7/7] fix bug in legacy .count due to with_limit_and_skip that
 was missing

---
 mongoengine/pymongo_support.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mongoengine/pymongo_support.py b/mongoengine/pymongo_support.py
index 284efc2f..9cf9e2ae 100644
--- a/mongoengine/pymongo_support.py
+++ b/mongoengine/pymongo_support.py
@@ -44,7 +44,8 @@ def count_documents(
     for option, option_value in kwargs.items():
         cursor_method = getattr(cursor, option)
         cursor = cursor_method(option_value)
-    return cursor.count()
+    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):