From efccbb72393b030353d88216db7e85609720643e Mon Sep 17 00:00:00 2001
From: James Kassemi <james.kassemi@close.com>
Date: Fri, 17 Jan 2025 11:28:24 -0500
Subject: [PATCH] Refactor `count` to `count_documents` and
 `estimated_document_count`

PyMongo `~=4.0` removes `Cursor.count` and `Collection.count` and provides
two new count methods: `Collection.count_documents` and
`Collection.estimated_document_count`.

Both `Cursor.count` and `Collection.count` in PyMongo `~=3.0` use `count` MongoDB
command, which provide potentially inaccurate results when executd without a query
predicate, in a transaction, or on a sharded cluster.

`Collection.count_documents` issues an aggregate query, which is accurate in
all cases, but relatively expensive.

`Collection.estimated_document_count` in PyMongo `>=4.0,<4.2` uses a $collStats
aggregation to estimate the document count for the collection. In PyMongo `>4.2`
it relies on the `count` MongoDB command.

When possible, we now try using the estimated document count for count operations,
but if a limit or skip is defined, or a hint is provided, we'll use the more accurate
count.

This mirrors the behavior of the upstream mongoengine implementation.
---
 mongoengine/context_managers.py  |  2 +-
 mongoengine/queryset/queryset.py | 19 ++++++++++++++++++-
 setup.py                         |  2 +-
 tests/test_context_managers.py   |  2 +-
 4 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py
index 61eb7dc3e..a7bbbe6d8 100644
--- a/mongoengine/context_managers.py
+++ b/mongoengine/context_managers.py
@@ -220,7 +220,7 @@ def __repr__(self):
     def _get_count(self):
         """ Get the number of queries. """
         ignore_query = {"ns": {"$ne": "%s.system.indexes" % self.db.name}}
-        count = self.db.system.profile.find(ignore_query).count() - self.counter
+        count = self.db.system.profile.count_documents(filter=ignore_query) - self.counter
         self.counter += 1
         return count
 
diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py
index d5721b9ae..1fed848ed 100644
--- a/mongoengine/queryset/queryset.py
+++ b/mongoengine/queryset/queryset.py
@@ -423,9 +423,26 @@ def count(self, with_limit_and_skip=True):
 
         if with_limit_and_skip and self._len is not None:
             return self._len
-        count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
+
+        options = {}
+        if with_limit_and_skip:
+            if self._limit is not None:
+                options["limit"] = self._limit
+            if self._skip is not None:
+                options["skip"] = self._skip
+        if self._hint not in (-1, None):
+            options["hint"] = self._hint
+
+        if self._query or options:
+            count = self._cursor.collection.count_documents(
+                filter=self._query, **options
+            )
+        else:
+            count = self._cursor.collection.estimated_document_count()
+
         if with_limit_and_skip:
             self._len = count
+
         return count
 
     def delete(self, write_concern=None, _from_doc_delete=False):
diff --git a/setup.py b/setup.py
index 0391ac00b..d949d527a 100644
--- a/setup.py
+++ b/setup.py
@@ -65,7 +65,7 @@ def get_version(version_tuple):
       long_description=LONG_DESCRIPTION,
       platforms=['any'],
       classifiers=CLASSIFIERS,
-      install_requires=['pymongo>=3.0,<3.14'],
+      install_requires=['pymongo>=3.7'],
       test_suite='nose.collector',
       **extra_opts
 )
diff --git a/tests/test_context_managers.py b/tests/test_context_managers.py
index 8852e5f61..65b66cf51 100644
--- a/tests/test_context_managers.py
+++ b/tests/test_context_managers.py
@@ -196,7 +196,7 @@ def test_query_counter(self):
             self.assertEqual(0, q)
 
             for i in range(1, 51):
-                db.test.find({}).count()
+                db.test.estimated_document_count()
 
             self.assertEqual(50, q)