Browse Source

Memcache keys are unicode on Py3. Closes #1697

Ask Solem 11 years ago
parent
commit
1540b00008
3 changed files with 42 additions and 24 deletions
  1. 15 7
      celery/backends/base.py
  2. 23 13
      celery/backends/cache.py
  3. 4 4
      celery/tests/backends/test_cache.py

+ 15 - 7
celery/backends/base.py

@@ -324,11 +324,19 @@ BaseDictBackend = BaseBackend  # XXX compat
 
 
 class KeyValueStoreBackend(BaseBackend):
-    task_keyprefix = ensure_bytes('celery-task-meta-')
-    group_keyprefix = ensure_bytes('celery-taskset-meta-')
-    chord_keyprefix = ensure_bytes('chord-unlock-')
+    key_t = ensure_bytes
+    task_keyprefix = 'celery-task-meta-'
+    group_keyprefix = 'celery-taskset-meta-'
+    chord_keyprefix = 'chord-unlock-'
     implements_incr = False
 
+    def __init__(self, *args, **kwargs):
+        super(KeyValueStoreBackend, self).__init__(*args, **kwargs)
+        self.key_t = self.__class__.key_t.__func__  # remove binding
+        self.task_keyprefix = self.key_t(self.task_keyprefix)
+        self.group_keyprefix = self.key_t(self.group_keyprefix)
+        self.chord_keyprefix = self.key_t(self.chord_keyprefix)
+
     def get(self, key):
         raise NotImplementedError('Must implement the get method.')
 
@@ -349,19 +357,19 @@ class KeyValueStoreBackend(BaseBackend):
 
     def get_key_for_task(self, task_id):
         """Get the cache key for a task by id."""
-        return self.task_keyprefix + ensure_bytes(task_id)
+        return self.task_keyprefix + self.key_t(task_id)
 
     def get_key_for_group(self, group_id):
         """Get the cache key for a group by id."""
-        return self.group_keyprefix + ensure_bytes(group_id)
+        return self.group_keyprefix + self.key_t(group_id)
 
     def get_key_for_chord(self, group_id):
         """Get the cache key for the chord waiting on group with given id."""
-        return self.chord_keyprefix + ensure_bytes(group_id)
+        return self.chord_keyprefix + self.key_t(group_id)
 
     def _strip_prefix(self, key):
         """Takes bytes, emits string."""
-        key = ensure_bytes(key)
+        key = self.key_t(key)
         for prefix in self.task_keyprefix, self.group_keyprefix:
             if key.startswith(prefix):
                 return bytes_to_str(key[len(prefix):])

+ 23 - 13
celery/backends/cache.py

@@ -8,7 +8,10 @@
 """
 from __future__ import absolute_import
 
+import sys
+
 from kombu.utils import cached_property
+from kombu.utils.encoding import bytes_to_str, ensure_bytes
 
 from celery.exceptions import ImproperlyConfigured
 from celery.utils.functional import LRUCache
@@ -19,6 +22,8 @@ __all__ = ['CacheBackend']
 
 _imp = [None]
 
+PY3 = sys.version_info[0] == 3
+
 REQUIRES_BACKEND = """\
 The memcached backend requires either pylibmc or python-memcached.\
 """
@@ -31,7 +36,7 @@ Please use one of the following backends instead: {1}\
 
 def import_best_memcache():
     if _imp[0] is None:
-        is_pylibmc = False
+        is_pylibmc, memcache_key_t = False, ensure_bytes
         try:
             import pylibmc as memcache
             is_pylibmc = True
@@ -40,17 +45,22 @@ def import_best_memcache():
                 import memcache  # noqa
             except ImportError:
                 raise ImproperlyConfigured(REQUIRES_BACKEND)
-        _imp[0] = (is_pylibmc, memcache)
+            else:
+                memcache_key_t = bytes_to_str if PY3 else ensure_bytes
+        _imp[0] = (is_pylibmc, memcache, memcache_key_t)
     return _imp[0]
 
 
 def get_best_memcache(*args, **kwargs):
-    behaviors = kwargs.pop('behaviors', None)
-    is_pylibmc, memcache = import_best_memcache()
-    client = memcache.Client(*args, **kwargs)
-    if is_pylibmc and behaviors is not None:
-        client.behaviors = behaviors
-    return client
+    is_pylibmc, memcache, key_t = import_best_memcache()
+    Client = _Client = memcache.Client
+
+    if not is_pylibmc:
+        def Client(*args, **kwargs):  # noqa
+            kwargs.pop('behaviors', None)
+            return _Client(*args, **kwargs)
+
+    return Client, key_t
 
 
 class DummyClient(object):
@@ -75,10 +85,10 @@ class DummyClient(object):
         return self.cache.incr(key, delta)
 
 
-backends = {'memcache': lambda: get_best_memcache,
-            'memcached': lambda: get_best_memcache,
-            'pylibmc': lambda: get_best_memcache,
-            'memory': lambda: DummyClient}
+backends = {'memcache': get_best_memcache,
+            'memcached': get_best_memcache,
+            'pylibmc': get_best_memcache,
+            'memory': lambda: (DummyClient, ensure_bytes)}
 
 
 class CacheBackend(KeyValueStoreBackend):
@@ -100,7 +110,7 @@ class CacheBackend(KeyValueStoreBackend):
             self.servers = servers.rstrip('/').split(';')
         self.expires = self.prepare_expires(expires, type=int)
         try:
-            self.Client = backends[self.backend]()
+            self.Client, self.key_t = backends[self.backend]()
         except KeyError:
             raise ImproperlyConfigured(UNKNOWN_BACKEND.format(
                 self.backend, ', '.join(backends)))

+ 4 - 4
celery/tests/backends/test_cache.py

@@ -164,7 +164,7 @@ class test_get_best_memcache(AppCase, MockCacheMixin):
             with reset_modules('celery.backends.cache'):
                 from celery.backends import cache
                 cache._imp = [None]
-                self.assertEqual(cache.get_best_memcache().__module__,
+                self.assertEqual(cache.get_best_memcache()[0].__module__,
                                  'pylibmc')
 
     def test_memcache(self):
@@ -173,7 +173,7 @@ class test_get_best_memcache(AppCase, MockCacheMixin):
                 with mask_modules('pylibmc'):
                     from celery.backends import cache
                     cache._imp = [None]
-                    self.assertEqual(cache.get_best_memcache().__module__,
+                    self.assertEqual(cache.get_best_memcache()[0]().__module__,
                                      'memcache')
 
     def test_no_implementations(self):
@@ -189,9 +189,9 @@ class test_get_best_memcache(AppCase, MockCacheMixin):
             with reset_modules('celery.backends.cache'):
                 from celery.backends import cache
                 cache._imp = [None]
-                cache.get_best_memcache(behaviors={'foo': 'bar'})
+                cache.get_best_memcache()[0](behaviors={'foo': 'bar'})
                 self.assertTrue(cache._imp[0])
-                cache.get_best_memcache()
+                cache.get_best_memcache()[0]()
 
     def test_backends(self):
         from celery.backends.cache import backends