From c8454b2f15fcff2bedaf055a384f454e761f86e1 Mon Sep 17 00:00:00 2001 From: oyiptong Date: Fri, 29 Oct 2010 20:10:38 -0400 Subject: [PATCH] Fixes SearchQuerySet not pickleable. Patch by oyiptong, tests by toastdriven. --- AUTHORS | 1 + haystack/backends/__init__.py | 3 +- haystack/query.py | 9 +++++ tests/core/tests/mocks.py | 5 +++ tests/core/tests/query.py | 48 ++++++++++++++++++++++ tests/solr_tests/tests/solr_backend.py | 55 ++++++++++++++++++++++++++ 6 files changed, 120 insertions(+), 1 deletion(-) diff --git a/AUTHORS b/AUTHORS index 1df29ad0a..4555386c0 100644 --- a/AUTHORS +++ b/AUTHORS @@ -46,3 +46,4 @@ Thanks to * Rob Hudson (robhudson) for improvements to the admin search. * apollo13 for simplifying ``SearchForm.__init__``. * Carl Meyer (carljm) for a patch regarding character primary keys. + * oyiptong for a patch regarding pickling. diff --git a/haystack/backends/__init__.py b/haystack/backends/__init__.py index f4e60a253..564e76dae 100644 --- a/haystack/backends/__init__.py +++ b/haystack/backends/__init__.py @@ -297,9 +297,10 @@ def __getstate__(self): """For pickling.""" obj_dict = self.__dict__.copy() del(obj_dict['backend']) + # Rip off the class bits as we'll be using this path when we go to load # the backend. - obj_dict['backend_used'] = ".".join(str(self.backend).replace("", "").split(".")[0:-1]) + obj_dict['backend_used'] = ".".join(str(self.backend).replace("<", "").split(".")[0:-1]) return obj_dict def __setstate__(self, obj_dict): diff --git a/haystack/query.py b/haystack/query.py index 5fbe4c6a9..ad7490e61 100644 --- a/haystack/query.py +++ b/haystack/query.py @@ -38,7 +38,16 @@ def __getstate__(self): len(self) obj_dict = self.__dict__.copy() obj_dict['_iter'] = None + del obj_dict['site'] return obj_dict + + def __setstate__(self, dict): + """ + For unpickling. + """ + self.__dict__ = dict + from haystack import site as main_site + self.site = main_site def __repr__(self): data = list(self[:REPR_OUTPUT_SIZE]) diff --git a/tests/core/tests/mocks.py b/tests/core/tests/mocks.py index 20eb653a7..18fcda7e8 100644 --- a/tests/core/tests/mocks.py +++ b/tests/core/tests/mocks.py @@ -112,3 +112,8 @@ def run_mlt(self): results = self.backend.more_like_this(self._mlt_instance, final_query) self._results = results['results'][self.start_offset:self.end_offset] self._hit_count = results['hits'] + + +# For pickling tests. +SearchBackend = MockSearchBackend +SearchQuery = MockSearchQuery diff --git a/tests/core/tests/query.py b/tests/core/tests/query.py index 45d2ee442..52f80b332 100644 --- a/tests/core/tests/query.py +++ b/tests/core/tests/query.py @@ -18,6 +18,16 @@ except NameError: from sets import Set as set +test_pickling = True + +try: + import cPickle as pickle +except ImportError: + try: + import pickle + except ImportError: + test_pickling = False + class SQTestCase(TestCase): def test_split_expression(self): @@ -661,3 +671,41 @@ def test_dictionary_lookup(self): EmptySearchQuerySets can be used in templates. """ self.assertRaises(TypeError, lambda: self.esqs['count']) + + +if test_pickling: + class PickleSearchQuerySetTestCase(TestCase): + def setUp(self): + super(PickleSearchQuerySetTestCase, self).setUp() + self.bsqs = SearchQuerySet(query=DummySearchQuery(backend=DummySearchBackend())) + self.msqs = SearchQuerySet(query=MockSearchQuery(backend=MockSearchBackend())) + self.mmsqs = SearchQuerySet(query=MockSearchQuery(backend=MixedMockSearchBackend())) + + # Stow. + self.old_debug = settings.DEBUG + settings.DEBUG = True + self.old_site = haystack.site + test_site = SearchSite() + test_site.register(MockModel) + test_site.register(CharPKMockModel) + haystack.site = test_site + + backends.reset_search_queries() + + def tearDown(self): + # Restore. + haystack.site = self.old_site + settings.DEBUG = self.old_debug + super(PickleSearchQuerySetTestCase, self).tearDown() + + def test_pickling(self): + results = self.msqs.all() + + for res in results: + # Make sure the cache is full. + pass + + in_a_pickle = pickle.dumps(results) + like_a_cuke = pickle.loads(in_a_pickle) + self.assertEqual(len(like_a_cuke), len(results)) + self.assertEqual(like_a_cuke[0].id, results[0].id) diff --git a/tests/solr_tests/tests/solr_backend.py b/tests/solr_tests/tests/solr_backend.py index 21c900035..0224f1fe2 100644 --- a/tests/solr_tests/tests/solr_backend.py +++ b/tests/solr_tests/tests/solr_backend.py @@ -15,6 +15,16 @@ except NameError: from sets import Set as set +test_pickling = True + +try: + import cPickle as pickle +except ImportError: + try: + import pickle + except ImportError: + test_pickling = False + def clear_solr_index(): # Wipe it clean. @@ -794,3 +804,48 @@ def test_round_trip(self): self.assertEqual(result.created, datetime.datetime(2009, 11, 21, 21, 31, 00)) self.assertEqual(result.tags, ['staff', 'outdoor', 'activist', 'scientist']) self.assertEqual(result.sites, [3, 5, 1]) + + +if test_pickling: + class LiveSolrPickleTestCase(TestCase): + fixtures = ['bulk_data.json'] + + def setUp(self): + super(LiveSolrPickleTestCase, self).setUp() + + # Wipe it clean. + clear_solr_index() + + # With the models registered, you get the proper bits. + import haystack + from haystack.sites import SearchSite + + # Stow. + self.old_site = haystack.site + test_site = SearchSite() + test_site.register(MockModel, SolrMockModelSearchIndex) + test_site.register(AnotherMockModel, SolrAnotherMockModelSearchIndex) + haystack.site = test_site + + self.sqs = SearchQuerySet() + + test_site.get_index(MockModel).update() + test_site.get_index(AnotherMockModel).update() + + def tearDown(self): + # Restore. + import haystack + haystack.site = self.old_site + super(LiveSolrPickleTestCase, self).tearDown() + + def test_pickling(self): + results = self.sqs.all() + + for res in results: + # Make sure the cache is full. + pass + + in_a_pickle = pickle.dumps(results) + like_a_cuke = pickle.loads(in_a_pickle) + self.assertEqual(len(like_a_cuke), len(results)) + self.assertEqual(like_a_cuke[0].id, results[0].id)