diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index c4a7e771..ae8cd407 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -64,6 +64,7 @@ class BaseQuerySet: 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 @@ class BaseQuerySet: "_ordering", "_snapshot", "_timeout", + "_allow_disk_use", "_read_preference", "_read_concern", "_iter", @@ -1165,6 +1167,16 @@ class BaseQuerySet: 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 @@ class BaseQuerySet: 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() diff --git a/tests/queryset/test_queryset.py b/tests/queryset/test_queryset.py index 079447bc..4d281c60 100644 --- a/tests/queryset/test_queryset.py +++ b/tests/queryset/test_queryset.py @@ -21,6 +21,7 @@ from mongoengine.queryset import ( QuerySetManager, queryset_manager, ) +from tests.utils import requires_mongodb_gte_44 class db_ops_tracker(query_counter): @@ -5656,6 +5657,31 @@ class TestQueryset(unittest.TestCase): 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 index 0899c208..adb0bdb4 100644 --- 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.