Merge pull request #103 from django-cache-machine/safe-pickle

safely pickle CachingQuerySet.timeout when set to DEFAULT_TIMEOUT

caching/base.py

@@ -138,10 +138,30 @@ def cache_objects(self, objects):
 
 class CachingQuerySet(models.query.QuerySet):
 
+    _default_timeout_pickle_key = '__DEFAULT_TIMEOUT__'
+
     def __init__(self, *args, **kw):
         super(CachingQuerySet, self).__init__(*args, **kw)
         self.timeout = DEFAULT_TIMEOUT
 
+    def __getstate__(self):
+        """
+        Safely pickle our timeout if it's a DEFAULT_TIMEOUT. This is not needed
+        by cache-machine itself, but by application code that may re-cache objects
+        retrieved using cache-machine.
+        """
+        state = dict()
+        state.update(self.__dict__)
+        if self.timeout == DEFAULT_TIMEOUT:
+            state['timeout'] = self._default_timeout_pickle_key
+        return state
+
+    def __setstate__(self, state):
+        """ Safely unpickle our timeout if it's a DEFAULT_TIMEOUT. """
+        self.__dict__.update(state)
+        if self.timeout == self._default_timeout_pickle_key:
+            self.timeout = DEFAULT_TIMEOUT
+
     def flush_key(self):
         return flush_key(self.query_key())
 

tests/test_cache.py

@@ -2,6 +2,7 @@
 
 import django
 import jinja2
+import pickle
 
 from django.conf import settings
 from django.test import TestCase
@@ -13,7 +14,7 @@
     import mock
 from nose.tools import eq_
 
-from caching import base, invalidation, config
+from caching import base, invalidation, config, compat
 from .testapp.models import Addon, User
 
 cache = invalidation.cache
@@ -533,3 +534,25 @@ def test_invalidate_on_create_disabled(self):
         assert not any([u.from_cache for u in users])
         User.objects.create(name='spam')
         assert all([u.from_cache for u in User.objects.all()])
+
+    def test_pickle_queryset(self):
+        """
+        Test for CacheingQuerySet.__getstate__ and CachingQuerySet.__setstate__.
+        """
+        # Make sure CachingQuerySet.timeout, when set to DEFAULT_TIMEOUT, can be safely
+        # pickled/unpickled on/from different Python processes which may have different
+        # underlying values for DEFAULT_TIMEOUT:
+        q1 = Addon.objects.all()
+        assert q1.timeout == compat.DEFAULT_TIMEOUT
+        pickled = pickle.dumps(q1)
+        new_timeout = object()
+        with mock.patch('caching.base.DEFAULT_TIMEOUT', new_timeout):
+            q2 = pickle.loads(pickled)
+            assert q2.timeout == new_timeout
+        # Make sure values other than DEFAULT_TIMEOUT remain unaffected:
+        q1 = Addon.objects.cache(10).all()
+        assert q1.timeout == 10
+        pickled = pickle.dumps(q1)
+        with mock.patch('caching.base.DEFAULT_TIMEOUT', new_timeout):
+            q2 = pickle.loads(pickled)
+            assert q2.timeout == 10