Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions mongoengine/queryset/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,39 @@ def __call__(self, q_obj=None, class_check=True, read_preference=None,

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.
"""
Expand Down
78 changes: 78 additions & 0 deletions tests/queryset/pickable.py
Original file line number Diff line number Diff line change
@@ -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)