memcached.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197
  1. import re
  2. import typing as _t
  3. from time import time
  4. from cachelib.base import BaseCache
  5. _test_memcached_key = re.compile(r"[^\x00-\x21\xff]{1,250}$").match
  6. class MemcachedCache(BaseCache):
  7. """A cache that uses memcached as backend.
  8. The first argument can either be an object that resembles the API of a
  9. :class:`memcache.Client` or a tuple/list of server addresses. In the
  10. event that a tuple/list is passed, Werkzeug tries to import the best
  11. available memcache library.
  12. This cache looks into the following packages/modules to find bindings for
  13. memcached:
  14. - ``pylibmc``
  15. - ``google.appengine.api.memcached``
  16. - ``memcached``
  17. - ``libmc``
  18. Implementation notes: This cache backend works around some limitations in
  19. memcached to simplify the interface. For example unicode keys are encoded
  20. to utf-8 on the fly. Methods such as :meth:`~BaseCache.get_dict` return
  21. the keys in the same format as passed. Furthermore all get methods
  22. silently ignore key errors to not cause problems when untrusted user data
  23. is passed to the get methods which is often the case in web applications.
  24. :param servers: a list or tuple of server addresses or alternatively
  25. a :class:`memcache.Client` or a compatible client.
  26. :param default_timeout: the default timeout that is used if no timeout is
  27. specified on :meth:`~BaseCache.set`. A timeout of
  28. 0 indicates that the cache never expires.
  29. :param key_prefix: a prefix that is added before all keys. This makes it
  30. possible to use the same memcached server for different
  31. applications. Keep in mind that
  32. :meth:`~BaseCache.clear` will also clear keys with a
  33. different prefix.
  34. """
  35. def __init__(
  36. self,
  37. servers: _t.Any = None,
  38. default_timeout: int = 300,
  39. key_prefix: _t.Optional[str] = None,
  40. ):
  41. BaseCache.__init__(self, default_timeout)
  42. if servers is None or isinstance(servers, (list, tuple)):
  43. if servers is None:
  44. servers = ["127.0.0.1:11211"]
  45. self._client = self.import_preferred_memcache_lib(servers)
  46. if self._client is None:
  47. raise RuntimeError("no memcache module found")
  48. else:
  49. # NOTE: servers is actually an already initialized memcache
  50. # client.
  51. self._client = servers
  52. self.key_prefix = key_prefix
  53. def _normalize_key(self, key: str) -> str:
  54. if self.key_prefix:
  55. key = self.key_prefix + key
  56. return key
  57. def _normalize_timeout(self, timeout: _t.Optional[int]) -> int:
  58. timeout = BaseCache._normalize_timeout(self, timeout)
  59. if timeout > 0:
  60. timeout = int(time()) + timeout
  61. return timeout
  62. def get(self, key: str) -> _t.Any:
  63. key = self._normalize_key(key)
  64. # memcached doesn't support keys longer than that. Because often
  65. # checks for so long keys can occur because it's tested from user
  66. # submitted data etc we fail silently for getting.
  67. if _test_memcached_key(key):
  68. return self._client.get(key)
  69. def get_dict(self, *keys: str) -> _t.Dict[str, _t.Any]:
  70. key_mapping = {}
  71. for key in keys:
  72. encoded_key = self._normalize_key(key)
  73. if _test_memcached_key(key):
  74. key_mapping[encoded_key] = key
  75. _keys = list(key_mapping)
  76. d = rv = self._client.get_multi(_keys) # type: _t.Dict[str, _t.Any]
  77. if self.key_prefix:
  78. rv = {}
  79. for key, value in d.items():
  80. rv[key_mapping[key]] = value
  81. if len(rv) < len(keys):
  82. for key in keys:
  83. if key not in rv:
  84. rv[key] = None
  85. return rv
  86. def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> bool:
  87. key = self._normalize_key(key)
  88. timeout = self._normalize_timeout(timeout)
  89. return bool(self._client.add(key, value, timeout))
  90. def set(
  91. self, key: str, value: _t.Any, timeout: _t.Optional[int] = None
  92. ) -> _t.Optional[bool]:
  93. key = self._normalize_key(key)
  94. timeout = self._normalize_timeout(timeout)
  95. return bool(self._client.set(key, value, timeout))
  96. def get_many(self, *keys: str) -> _t.List[_t.Any]:
  97. d = self.get_dict(*keys)
  98. return [d[key] for key in keys]
  99. def set_many(
  100. self, mapping: _t.Dict[str, _t.Any], timeout: _t.Optional[int] = None
  101. ) -> _t.List[_t.Any]:
  102. new_mapping = {}
  103. for key, value in mapping.items():
  104. key = self._normalize_key(key)
  105. new_mapping[key] = value
  106. timeout = self._normalize_timeout(timeout)
  107. failed_keys = self._client.set_multi(
  108. new_mapping, timeout
  109. ) # type: _t.List[_t.Any]
  110. k_normkey = zip(mapping.keys(), new_mapping.keys())
  111. return [k for k, nkey in k_normkey if nkey not in failed_keys]
  112. def delete(self, key: str) -> bool:
  113. key = self._normalize_key(key)
  114. if _test_memcached_key(key):
  115. return bool(self._client.delete(key))
  116. return False
  117. def delete_many(self, *keys: str) -> _t.List[_t.Any]:
  118. new_keys = []
  119. for key in keys:
  120. key = self._normalize_key(key)
  121. if _test_memcached_key(key):
  122. new_keys.append(key)
  123. self._client.delete_multi(new_keys)
  124. return [k for k in new_keys if not self.has(k)]
  125. def has(self, key: str) -> bool:
  126. key = self._normalize_key(key)
  127. if _test_memcached_key(key):
  128. return bool(self._client.append(key, ""))
  129. return False
  130. def clear(self) -> bool:
  131. return bool(self._client.flush_all())
  132. def inc(self, key: str, delta: int = 1) -> _t.Optional[int]:
  133. key = self._normalize_key(key)
  134. value = (self._client.get(key) or 0) + delta
  135. return value if self.set(key, value) else None
  136. def dec(self, key: str, delta: int = 1) -> _t.Optional[int]:
  137. key = self._normalize_key(key)
  138. value = (self._client.get(key) or 0) - delta
  139. return value if self.set(key, value) else None
  140. def import_preferred_memcache_lib(self, servers: _t.Any) -> _t.Any:
  141. """Returns an initialized memcache client. Used by the constructor."""
  142. try:
  143. import pylibmc # type: ignore
  144. except ImportError:
  145. pass
  146. else:
  147. return pylibmc.Client(servers)
  148. try:
  149. from google.appengine.api import memcache # type: ignore
  150. except ImportError:
  151. pass
  152. else:
  153. return memcache.Client()
  154. try:
  155. import memcache # type: ignore
  156. except ImportError:
  157. pass
  158. else:
  159. return memcache.Client(servers)
  160. try:
  161. import libmc # type: ignore
  162. except ImportError:
  163. pass
  164. else:
  165. return libmc.Client(servers)