Kaynağa Gözat

Cache backend: Key cannot be unicode. Closes #504

Neil Chintomby 13 yıl önce
ebeveyn
işleme
7b35e5e55b

+ 4 - 0
Changelog

@@ -199,6 +199,10 @@ News
 * ``celeryd_multi`` now supports a ``stop_verify`` command to wait for
   processes to shutdown.
 
+* Cache backend did not work if the cache key was unicode (Issue #504).
+
+    Fix contributed by Neil Chintomby.
+
 * New setting :setting:`CELERY_RESULT_DB_SHORT_LIVED_SESSIONS` added,
   which if enabled will disable the caching of SQLAlchemy sessions
   (Issue #449).

+ 7 - 0
celery/backends/cache.py

@@ -3,6 +3,7 @@ from __future__ import absolute_import
 from ..datastructures import LRUCache
 from ..exceptions import ImproperlyConfigured
 from ..utils import cached_property
+from ..utils.encoding import str_to_bytes
 
 from .base import KeyValueStoreBackend
 
@@ -82,6 +83,12 @@ class CacheBackend(KeyValueStoreBackend):
                     "following backends: %s" % (self.backend,
                                                 ", ".join(backends.keys())))
 
+    def get_key_for_task(self, task_id):
+        return str_to_bytes(self.task_keyprefix + task_id)
+
+    def get_key_for_taskset(self, taskset_id):
+        return str_to_bytes(self.taskset_keyprefix + taskset_id)
+
     def get(self, key):
         return self.client.get(key)
 

+ 62 - 4
celery/tests/test_backends/test_cache.py

@@ -87,16 +87,26 @@ class test_CacheBackend(unittest.TestCase):
             CacheBackend(backend="unknown://")
 
 
-class MyClient(DummyClient):
+class MyMemcachedStringEncodingError(Exception):
     pass
 
 
-class test_get_best_memcache(unittest.TestCase):
+class MemcachedClient(DummyClient):
+
+    def set(self, key, value, *args, **kwargs):
+        if isinstance(key, unicode):
+            raise MyMemcachedStringEncodingError(
+                    "Keys must be str()'s, not unicode.  Convert your unicode "
+                    "strings using mystring.encode(charset)!")
+        return super(MemcachedClient, self).set(key, value, *args, **kwargs)
+
+
+class MockCacheMixin(object):
 
     @contextmanager
     def mock_memcache(self):
         memcache = types.ModuleType("memcache")
-        memcache.Client = MyClient
+        memcache.Client = MemcachedClient
         memcache.Client.__module__ = memcache.__name__
         prev, sys.modules["memcache"] = sys.modules.get("memcache"), memcache
         yield True
@@ -106,7 +116,7 @@ class test_get_best_memcache(unittest.TestCase):
     @contextmanager
     def mock_pylibmc(self):
         pylibmc = types.ModuleType("pylibmc")
-        pylibmc.Client = MyClient
+        pylibmc.Client = MemcachedClient
         pylibmc.Client.__module__ = pylibmc.__name__
         prev = sys.modules.get("pylibmc")
         sys.modules["pylibmc"] = pylibmc
@@ -114,6 +124,9 @@ class test_get_best_memcache(unittest.TestCase):
         if prev is not None:
             sys.modules["pylibmc"] = prev
 
+
+class test_get_best_memcache(unittest.TestCase, MockCacheMixin):
+
     def test_pylibmc(self):
         with reset_modules("celery.backends.cache"):
             with self.mock_pylibmc():
@@ -151,3 +164,48 @@ class test_get_best_memcache(unittest.TestCase):
         from celery.backends.cache import backends
         for name, fun in backends.items():
             self.assertTrue(fun())
+
+
+class test_memcache_key(unittest.TestCase, MockCacheMixin):
+
+    def test_memcache_unicode_key(self):
+        with self.mock_memcache():
+            with reset_modules("celery.backends.cache"):
+                with mask_modules("pylibmc"):
+                    from celery.backends import cache
+                    cache._imp = [None]
+                    task_id, result = unicode(uuid()), 42
+                    b = cache.CacheBackend(backend='memcache')
+                    b.store_result(task_id, result, status=states.SUCCESS)
+                    self.assertEqual(b.get_result(task_id), result)
+
+    def test_memcache_bytes_key(self):
+        with self.mock_memcache():
+            with reset_modules("celery.backends.cache"):
+                with mask_modules("pylibmc"):
+                    from celery.backends import cache
+                    cache._imp = [None]
+                    task_id, result = bytes(uuid()), 42
+                    b = cache.CacheBackend(backend='memcache')
+                    b.store_result(task_id, result, status=states.SUCCESS)
+                    self.assertEqual(b.get_result(task_id), result)
+
+    def test_pylibmc_unicode_key(self):
+        with reset_modules("celery.backends.cache"):
+            with self.mock_pylibmc():
+                from celery.backends import cache
+                cache._imp = [None]
+                task_id, result = unicode(uuid()), 42
+                b = cache.CacheBackend(backend='memcache')
+                b.store_result(task_id, result, status=states.SUCCESS)
+                self.assertEqual(b.get_result(task_id), result)
+
+    def test_pylibmc_bytes_key(self):
+        with reset_modules("celery.backends.cache"):
+            with self.mock_pylibmc():
+                from celery.backends import cache
+                cache._imp = [None]
+                task_id, result = bytes(uuid()), 42
+                b = cache.CacheBackend(backend='memcache')
+                b.store_result(task_id, result, status=states.SUCCESS)
+                self.assertEqual(b.get_result(task_id), result)