Repository: MongoEngine/mongoengine

Issue:
'Queryset' object has no attribute 'allowDiskUse'
Hello.
I got a QueryExceededMemoryLimitNoDiskUseAllowed Error, while using something like this: `Email.objects(hash__in='1234').search_text("hello").order_by('sent_date').no_cache().timeout(False)`

Is there a way to allow Disk use in my specific example?


Discussion:
I don't think the error is from mongoengine but it is probably from pymongo (perhaps even mongoDB). If you can identify how to change that in pymongo, we can then see if we can add support in the mongoengine config
But are you sure you need to process all the data in 1 shot ? That won't scale. The queryset gives you a generator so if you can treat this in a loop without needing to hold everything in memory, perhaps you can avoid that situation
In pymongo you have the following option:
`db.collection.find(<match>).sort(<sort>).allow_disk_use(True)`

They added that in Version 3.11.0
https://pymongo.readthedocs.io/en/stable/changelog.html
I'd accept a PR if you feel like adding support for that. If you need inspiration, what you want is quite similar to how `Doc.objects.timeout(True) `is behaving (https://github.com/MongoEngine/mongoengine/blob/2e284b93b88fb9e64fc9749ca14df2d9f6629165/mongoengine/queryset/base.py#L1168)

Resolving patch:
diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
--- a/mongoengine/queryset/base.py
+++ b/mongoengine/queryset/base.py
@@ -64,6 +64,7 @@ def __init__(self, document, collection):
         self._ordering = None
         self._snapshot = False
         self._timeout = True
+        self._allow_disk_use = False
         self._read_preference = None
         self._read_concern = None
         self._iter = False
@@ -799,6 +800,7 @@ def _clone_into(self, new_qs):
             "_ordering",
             "_snapshot",
             "_timeout",
+            "_allow_disk_use",
             "_read_preference",
             "_read_concern",
             "_iter",
@@ -1165,6 +1167,16 @@ def snapshot(self, enabled):
         queryset._snapshot = enabled
         return queryset
 
+    def allow_disk_use(self, enabled):
+        """Enable or disable the use of temporary files on disk while processing a blocking sort operation.
+         (To store data exceeding the 100 megabyte system memory limit)
+
+        :param enabled: whether or not temporary files on disk are used
+        """
+        queryset = self.clone()
+        queryset._allow_disk_use = enabled
+        return queryset
+
     def timeout(self, enabled):
         """Enable or disable the default mongod timeout when querying. (no_cursor_timeout option)
 
@@ -1604,6 +1616,9 @@ def _cursor_args(self):
         if not self._timeout:
             cursor_args["no_cursor_timeout"] = True
 
+        if self._allow_disk_use:
+            cursor_args["allow_disk_use"] = True
+
         if self._loaded_fields:
             cursor_args[fields_name] = self._loaded_fields.as_dict()
 


Test patch:
diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py
--- a/tests/queryset/test_queryset.py
+++ b/tests/queryset/test_queryset.py
@@ -21,6 +21,7 @@
     QuerySetManager,
     queryset_manager,
 )
+from tests.utils import requires_mongodb_gte_44
 
 
 class db_ops_tracker(query_counter):
@@ -5656,6 +5657,31 @@ def test_no_cursor_timeout(self):
         qs = self.Person.objects().timeout(False)
         assert qs._cursor_args == {"no_cursor_timeout": True}
 
+    @requires_mongodb_gte_44
+    def test_allow_disk_use(self):
+        qs = self.Person.objects()
+        assert qs._cursor_args == {}
+
+        qs = self.Person.objects().allow_disk_use(False)
+        assert qs._cursor_args == {}
+
+        qs = self.Person.objects().allow_disk_use(True)
+        assert qs._cursor_args == {"allow_disk_use": True}
+
+        # Test if allow_disk_use changes the results
+        self.Person.drop_collection()
+        self.Person.objects.create(name="Foo", age=12)
+        self.Person.objects.create(name="Baz", age=17)
+        self.Person.objects.create(name="Bar", age=13)
+
+        qs_disk = self.Person.objects().order_by("age").allow_disk_use(True)
+        qs = self.Person.objects().order_by("age")
+
+        assert qs_disk.count() == qs.count()
+
+        for index in range(qs_disk.count()):
+            assert qs_disk[index] == qs[index]
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/utils.py b/tests/utils.py
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,3 +1,4 @@
+import operator
 import unittest
 
 import pytest
@@ -33,6 +34,10 @@ def get_as_pymongo(doc):
     return doc.__class__.objects.as_pymongo().get(id=doc.id)
 
 
+def requires_mongodb_gte_44(func):
+    return _decorated_with_ver_requirement(func, (4, 4), oper=operator.ge)
+
+
 def _decorated_with_ver_requirement(func, mongo_version_req, oper):
     """Return a MongoDB version requirement decorator.