Merge pull request #2529 from bagerard/fix_performance_issue_count

Use estimated_documents_count OR documents_count based on query

Author: bagerard

docs/changelog.rst

@@ -11,6 +11,7 @@ Development
 - Add tests against MongoDB 6.0 and MongoDB 7.0 in the pipeline
 - Fix validate() not being called when inheritance is used in EmbeddedDocument and validate is overriden #2784
 - Add support for readPreferenceTags in connection parameters #2644
+- Use estimated_documents_count OR documents_count when count is called, based on the query #2529
 
 Changes in 0.27.0
 =================

mongoengine/pymongo_support.py

@@ -37,7 +37,12 @@ def count_documents(
     # count_documents appeared in pymongo 3.7
     if PYMONGO_VERSION >= (3, 7):
         try:
-            return collection.count_documents(filter=filter, **kwargs)
+            if not filter and set(kwargs) <= {"max_time_ms"}:
+                # when no filter is provided, estimated_document_count
+                # is a lot faster as it uses the collection metadata
+                return collection.estimated_document_count(**kwargs)
+            else:
+                return collection.count_documents(filter=filter, **kwargs)
         except OperationFailure as err:
             if PYMONGO_VERSION >= (4,):
                 raise
@@ -46,15 +51,10 @@ def count_documents(
             # 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
-            message = str(err)
-            if not (
-                "not allowed in this context" in message
-                and (
-                    "$where" in message
-                    or "$geoNear" in message
-                    or "$near" in message
-                    or "$nearSphere" in message
-                )
+            if (
+                "$geoNear, $near, and $nearSphere are not allowed in this context"
+                not in str(err)
+                and "$where is not allowed in this context" not in str(err)
             ):
                 raise
 

tests/test_context_managers.py

@@ -14,9 +14,10 @@
     switch_db,
 )
 from mongoengine.pymongo_support import count_documents
+from tests.utils import MongoDBTestCase
 
 
-class TestContextManagers:
+class TestContextManagers(MongoDBTestCase):
     def test_set_write_concern(self):
         connect("mongoenginetest")