Merge pull request MongoEngine#2477 from lukasloetkolben/feature_allow_disk_use

Feature: "allow_disk_use"

Author: bagerard

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()
 

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()

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.