Add timeout to PyMongo.count_documents()

kobotoolbox · May 19, 2021 · e3e5124 · e3e5124
1 parent 3ac4270
commit e3e5124
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/onadata/apps/viewer/models/parsed_instance.py b/onadata/apps/viewer/models/parsed_instance.py
@@ -97,7 +97,13 @@ def query_mongo(cls, username, id_string, query, fields, sort, start=0,
         query = cls._get_mongo_cursor_query(query, hide_deleted, username, id_string)
 
         if count:
-            return [{"count": xform_instances.count_documents(query)}]
+            return [
+                {
+                    'count': xform_instances.count_documents(
+                        query, maxTimeMS=settings.MONGO_DB_MAX_TIME_MS
+                    )
+                }
+            ]
 
         cursor = cls._get_mongo_cursor(query, fields)
 
@@ -149,7 +155,13 @@ def query_mongo_minimal(
 
         query = cls._get_mongo_cursor_query(query, hide_deleted)
         if count:
-            return [{"count": xform_instances.count_documents(query)}]
+            return [
+                {
+                    'count': xform_instances.count_documents(
+                        query, maxTimeMS=settings.MONGO_DB_MAX_TIME_MS
+                    )
+                }
+            ]
 
         cursor = cls._get_mongo_cursor(query, fields)
 
@@ -172,7 +184,13 @@ def query_mongo_no_paging(cls, query, fields, count=False, hide_deleted=True):
         query = cls._get_mongo_cursor_query(query, hide_deleted)
 
         if count:
-            return [{"count": xform_instances.count_documents(query)}]
+            return [
+                {
+                    'count': xform_instances.count_documents(
+                        query, maxTimeMS=settings.MONGO_DB_MAX_TIME_MS
+                    )
+                }
+            ]
 
         return cls._get_mongo_cursor(query, fields)