diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 7efb0fb6..38af72ff 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -123,6 +123,39 @@ class BaseQuerySet(object): return queryset + def __getstate__(self): + """ + Need for pickling queryset + + See https://github.com/MongoEngine/mongoengine/issues/442 + """ + + obj_dict = self.__dict__.copy() + + # don't picke collection, instead pickle collection params + obj_dict.pop("_collection_obj") + + # don't pickle cursor + obj_dict["_cursor_obj"] = None + + return obj_dict + + def __setstate__(self, obj_dict): + """ + Need for pickling queryset + + See https://github.com/MongoEngine/mongoengine/issues/442 + """ + + obj_dict["_collection_obj"] = obj_dict["_document"]._get_collection() + + # update attributes + self.__dict__.update(obj_dict) + + # forse load cursor + # self._cursor + + def __getitem__(self, key): """Support skip and limit using getitem and slicing syntax. """ diff --git a/tests/queryset/pickable.py b/tests/queryset/pickable.py new file mode 100644 index 00000000..d96e7dc6 --- /dev/null +++ b/tests/queryset/pickable.py @@ -0,0 +1,78 @@ +import pickle +import unittest +from pymongo.mongo_client import MongoClient +from mongoengine import Document, StringField, IntField +from mongoengine.connection import connect + +__author__ = 'stas' + +class Person(Document): + name = StringField() + age = IntField() + +class TestQuerysetPickable(unittest.TestCase): + """ + Test for adding pickling support for QuerySet instances + See issue https://github.com/MongoEngine/mongoengine/issues/442 + """ + def setUp(self): + super(TestQuerysetPickable, self).setUp() + + connection = connect(db="test") #type: pymongo.mongo_client.MongoClient + + connection.drop_database("test") + + self.john = Person.objects.create( + name="John", + age=21 + ) + + + def test_picke_simple_qs(self): + + qs = Person.objects.all() + + pickle.dumps(qs) + + def _get_loaded(self, qs): + s = pickle.dumps(qs) + + return pickle.loads(s) + + def test_unpickle(self): + qs = Person.objects.all() + + loadedQs = self._get_loaded(qs) + + self.assertEqual(qs.count(), loadedQs.count()) + + #can update loadedQs + loadedQs.update(age=23) + + #check + self.assertEqual(Person.objects.first().age, 23) + + def test_pickle_support_filtration(self): + Person.objects.create( + name="Alice", + age=22 + ) + + Person.objects.create( + name="Bob", + age=23 + ) + + qs = Person.objects.filter(age__gte=22) + self.assertEqual(qs.count(), 2) + + loaded = self._get_loaded(qs) + + self.assertEqual(loaded.count(), 2) + self.assertEqual(loaded.filter(name="Bob").first().age, 23) + + + + + +