default.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778
  1. import random
  2. import re
  3. import socket
  4. from collections import OrderedDict
  5. from datetime import datetime
  6. from typing import Any, Dict, Iterator, List, Optional, Union
  7. from django.conf import settings
  8. from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache, get_key_func
  9. from django.core.exceptions import ImproperlyConfigured
  10. from django.utils.module_loading import import_string
  11. from redis import Redis
  12. from redis.exceptions import ConnectionError, ResponseError, TimeoutError
  13. from .. import pool
  14. from ..exceptions import CompressorError, ConnectionInterrupted
  15. from ..util import CacheKey
  16. _main_exceptions = (TimeoutError, ResponseError, ConnectionError, socket.timeout)
  17. special_re = re.compile("([*?[])")
  18. def glob_escape(s: str) -> str:
  19. return special_re.sub(r"[\1]", s)
  20. class DefaultClient:
  21. def __init__(self, server, params: Dict[str, Any], backend: BaseCache) -> None:
  22. self._backend = backend
  23. self._server = server
  24. self._params = params
  25. self.reverse_key = get_key_func(
  26. params.get("REVERSE_KEY_FUNCTION")
  27. or "django_redis.util.default_reverse_key"
  28. )
  29. if not self._server:
  30. raise ImproperlyConfigured("Missing connections string")
  31. if not isinstance(self._server, (list, tuple, set)):
  32. self._server = self._server.split(",")
  33. self._clients: List[Optional[Redis]] = [None] * len(self._server)
  34. self._options = params.get("OPTIONS", {})
  35. self._replica_read_only = self._options.get("REPLICA_READ_ONLY", True)
  36. serializer_path = self._options.get(
  37. "SERIALIZER", "django_redis.serializers.pickle.PickleSerializer"
  38. )
  39. serializer_cls = import_string(serializer_path)
  40. compressor_path = self._options.get(
  41. "COMPRESSOR", "django_redis.compressors.identity.IdentityCompressor"
  42. )
  43. compressor_cls = import_string(compressor_path)
  44. self._serializer = serializer_cls(options=self._options)
  45. self._compressor = compressor_cls(options=self._options)
  46. self.connection_factory = pool.get_connection_factory(options=self._options)
  47. def __contains__(self, key: Any) -> bool:
  48. return self.has_key(key)
  49. def get_next_client_index(
  50. self, write: bool = True, tried: Optional[List[int]] = None
  51. ) -> int:
  52. """
  53. Return a next index for read client. This function implements a default
  54. behavior for get a next read client for a replication setup.
  55. Overwrite this function if you want a specific
  56. behavior.
  57. """
  58. if tried is None:
  59. tried = list()
  60. if tried and len(tried) < len(self._server):
  61. not_tried = [i for i in range(0, len(self._server)) if i not in tried]
  62. return random.choice(not_tried)
  63. if write or len(self._server) == 1:
  64. return 0
  65. return random.randint(1, len(self._server) - 1)
  66. def get_client(
  67. self,
  68. write: bool = True,
  69. tried: Optional[List[int]] = None,
  70. show_index: bool = False,
  71. ):
  72. """
  73. Method used for obtain a raw redis client.
  74. This function is used by almost all cache backend
  75. operations for obtain a native redis client/connection
  76. instance.
  77. """
  78. index = self.get_next_client_index(write=write, tried=tried)
  79. if self._clients[index] is None:
  80. self._clients[index] = self.connect(index)
  81. if show_index:
  82. return self._clients[index], index
  83. else:
  84. return self._clients[index]
  85. def connect(self, index: int = 0) -> Redis:
  86. """
  87. Given a connection index, returns a new raw redis client/connection
  88. instance. Index is used for replication setups and indicates that
  89. connection string should be used. In normal setups, index is 0.
  90. """
  91. return self.connection_factory.connect(self._server[index])
  92. def disconnect(self, index=0, client=None):
  93. """delegates the connection factory to disconnect the client"""
  94. if not client:
  95. client = self._clients[index]
  96. return self.connection_factory.disconnect(client) if client else None
  97. def set(
  98. self,
  99. key: Any,
  100. value: Any,
  101. timeout: Optional[float] = DEFAULT_TIMEOUT,
  102. version: Optional[int] = None,
  103. client: Optional[Redis] = None,
  104. nx: bool = False,
  105. xx: bool = False,
  106. ) -> bool:
  107. """
  108. Persist a value to the cache, and set an optional expiration time.
  109. Also supports optional nx parameter. If set to True - will use redis
  110. setnx instead of set.
  111. """
  112. nkey = self.make_key(key, version=version)
  113. nvalue = self.encode(value)
  114. if timeout is DEFAULT_TIMEOUT:
  115. timeout = self._backend.default_timeout
  116. original_client = client
  117. tried: List[int] = []
  118. while True:
  119. try:
  120. if client is None:
  121. client, index = self.get_client(
  122. write=True, tried=tried, show_index=True
  123. )
  124. if timeout is not None:
  125. # Convert to milliseconds
  126. timeout = int(timeout * 1000)
  127. if timeout <= 0:
  128. if nx:
  129. # Using negative timeouts when nx is True should
  130. # not expire (in our case delete) the value if it exists.
  131. # Obviously expire not existent value is noop.
  132. return not self.has_key(key, version=version, client=client)
  133. else:
  134. # redis doesn't support negative timeouts in ex flags
  135. # so it seems that it's better to just delete the key
  136. # than to set it and than expire in a pipeline
  137. return bool(
  138. self.delete(key, client=client, version=version)
  139. )
  140. return bool(client.set(nkey, nvalue, nx=nx, px=timeout, xx=xx))
  141. except _main_exceptions as e:
  142. if (
  143. not original_client
  144. and not self._replica_read_only
  145. and len(tried) < len(self._server)
  146. ):
  147. tried.append(index)
  148. client = None
  149. continue
  150. raise ConnectionInterrupted(connection=client) from e
  151. def incr_version(
  152. self,
  153. key: Any,
  154. delta: int = 1,
  155. version: Optional[int] = None,
  156. client: Optional[Redis] = None,
  157. ) -> int:
  158. """
  159. Adds delta to the cache version for the supplied key. Returns the
  160. new version.
  161. """
  162. if client is None:
  163. client = self.get_client(write=True)
  164. if version is None:
  165. version = self._backend.version
  166. old_key = self.make_key(key, version)
  167. value = self.get(old_key, version=version, client=client)
  168. try:
  169. ttl = self.ttl(old_key, version=version, client=client)
  170. except _main_exceptions as e:
  171. raise ConnectionInterrupted(connection=client) from e
  172. if value is None:
  173. raise ValueError("Key '%s' not found" % key)
  174. if isinstance(key, CacheKey):
  175. new_key = self.make_key(key.original_key(), version=version + delta)
  176. else:
  177. new_key = self.make_key(key, version=version + delta)
  178. self.set(new_key, value, timeout=ttl, client=client)
  179. self.delete(old_key, client=client)
  180. return version + delta
  181. def add(
  182. self,
  183. key: Any,
  184. value: Any,
  185. timeout: Any = DEFAULT_TIMEOUT,
  186. version: Optional[Any] = None,
  187. client: Optional[Redis] = None,
  188. ) -> bool:
  189. """
  190. Add a value to the cache, failing if the key already exists.
  191. Returns ``True`` if the object was added, ``False`` if not.
  192. """
  193. return self.set(key, value, timeout, version=version, client=client, nx=True)
  194. def get(
  195. self,
  196. key: Any,
  197. default=None,
  198. version: Optional[int] = None,
  199. client: Optional[Redis] = None,
  200. ) -> Any:
  201. """
  202. Retrieve a value from the cache.
  203. Returns decoded value if key is found, the default if not.
  204. """
  205. if client is None:
  206. client = self.get_client(write=False)
  207. key = self.make_key(key, version=version)
  208. try:
  209. value = client.get(key)
  210. except _main_exceptions as e:
  211. raise ConnectionInterrupted(connection=client) from e
  212. if value is None:
  213. return default
  214. return self.decode(value)
  215. def persist(
  216. self, key: Any, version: Optional[int] = None, client: Optional[Redis] = None
  217. ) -> bool:
  218. if client is None:
  219. client = self.get_client(write=True)
  220. key = self.make_key(key, version=version)
  221. return client.persist(key)
  222. def expire(
  223. self,
  224. key: Any,
  225. timeout,
  226. version: Optional[int] = None,
  227. client: Optional[Redis] = None,
  228. ) -> bool:
  229. if client is None:
  230. client = self.get_client(write=True)
  231. key = self.make_key(key, version=version)
  232. return client.expire(key, timeout)
  233. def pexpire(self, key, timeout, version=None, client=None) -> bool:
  234. if client is None:
  235. client = self.get_client(write=True)
  236. key = self.make_key(key, version=version)
  237. # Temporary casting until https://github.com/redis/redis-py/issues/1664
  238. # is fixed.
  239. return bool(client.pexpire(key, timeout))
  240. def pexpire_at(
  241. self,
  242. key: Any,
  243. when: Union[datetime, int],
  244. version: Optional[int] = None,
  245. client: Optional[Redis] = None,
  246. ) -> bool:
  247. """
  248. Set an expire flag on a ``key`` to ``when``, which can be represented
  249. as an integer indicating unix time or a Python datetime object.
  250. """
  251. if client is None:
  252. client = self.get_client(write=True)
  253. key = self.make_key(key, version=version)
  254. return bool(client.pexpireat(key, when))
  255. def expire_at(
  256. self,
  257. key: Any,
  258. when: Union[datetime, int],
  259. version: Optional[int] = None,
  260. client: Optional[Redis] = None,
  261. ) -> bool:
  262. """
  263. Set an expire flag on a ``key`` to ``when``, which can be represented
  264. as an integer indicating unix time or a Python datetime object.
  265. """
  266. if client is None:
  267. client = self.get_client(write=True)
  268. key = self.make_key(key, version=version)
  269. return client.expireat(key, when)
  270. def lock(
  271. self,
  272. key,
  273. version: Optional[int] = None,
  274. timeout=None,
  275. sleep=0.1,
  276. blocking_timeout=None,
  277. client: Optional[Redis] = None,
  278. thread_local=True,
  279. ):
  280. if client is None:
  281. client = self.get_client(write=True)
  282. key = self.make_key(key, version=version)
  283. return client.lock(
  284. key,
  285. timeout=timeout,
  286. sleep=sleep,
  287. blocking_timeout=blocking_timeout,
  288. thread_local=thread_local,
  289. )
  290. def delete(
  291. self,
  292. key: Any,
  293. version: Optional[int] = None,
  294. prefix: Optional[str] = None,
  295. client: Optional[Redis] = None,
  296. ) -> int:
  297. """
  298. Remove a key from the cache.
  299. """
  300. if client is None:
  301. client = self.get_client(write=True)
  302. try:
  303. return client.delete(self.make_key(key, version=version, prefix=prefix))
  304. except _main_exceptions as e:
  305. raise ConnectionInterrupted(connection=client) from e
  306. def delete_pattern(
  307. self,
  308. pattern: str,
  309. version: Optional[int] = None,
  310. prefix: Optional[str] = None,
  311. client: Optional[Redis] = None,
  312. itersize: Optional[int] = None,
  313. ) -> int:
  314. """
  315. Remove all keys matching pattern.
  316. """
  317. if client is None:
  318. client = self.get_client(write=True)
  319. pattern = self.make_pattern(pattern, version=version, prefix=prefix)
  320. try:
  321. count = 0
  322. pipeline = client.pipeline()
  323. for key in client.scan_iter(match=pattern, count=itersize):
  324. pipeline.delete(key)
  325. count += 1
  326. pipeline.execute()
  327. return count
  328. except _main_exceptions as e:
  329. raise ConnectionInterrupted(connection=client) from e
  330. def delete_many(
  331. self, keys, version: Optional[int] = None, client: Optional[Redis] = None
  332. ):
  333. """
  334. Remove multiple keys at once.
  335. """
  336. if client is None:
  337. client = self.get_client(write=True)
  338. keys = [self.make_key(k, version=version) for k in keys]
  339. if not keys:
  340. return
  341. try:
  342. return client.delete(*keys)
  343. except _main_exceptions as e:
  344. raise ConnectionInterrupted(connection=client) from e
  345. def clear(self, client: Optional[Redis] = None) -> None:
  346. """
  347. Flush all cache keys.
  348. """
  349. if client is None:
  350. client = self.get_client(write=True)
  351. try:
  352. client.flushdb()
  353. except _main_exceptions as e:
  354. raise ConnectionInterrupted(connection=client) from e
  355. def decode(self, value: Union[bytes, int]) -> Any:
  356. """
  357. Decode the given value.
  358. """
  359. try:
  360. value = int(value)
  361. except (ValueError, TypeError):
  362. try:
  363. value = self._compressor.decompress(value)
  364. except CompressorError:
  365. # Handle little values, chosen to be not compressed
  366. pass
  367. value = self._serializer.loads(value)
  368. return value
  369. def encode(self, value: Any) -> Union[bytes, Any]:
  370. """
  371. Encode the given value.
  372. """
  373. if isinstance(value, bool) or not isinstance(value, int):
  374. value = self._serializer.dumps(value)
  375. value = self._compressor.compress(value)
  376. return value
  377. return value
  378. def get_many(
  379. self, keys, version: Optional[int] = None, client: Optional[Redis] = None
  380. ) -> OrderedDict:
  381. """
  382. Retrieve many keys.
  383. """
  384. if client is None:
  385. client = self.get_client(write=False)
  386. if not keys:
  387. return OrderedDict()
  388. recovered_data = OrderedDict()
  389. map_keys = OrderedDict((self.make_key(k, version=version), k) for k in keys)
  390. try:
  391. results = client.mget(*map_keys)
  392. except _main_exceptions as e:
  393. raise ConnectionInterrupted(connection=client) from e
  394. for key, value in zip(map_keys, results):
  395. if value is None:
  396. continue
  397. recovered_data[map_keys[key]] = self.decode(value)
  398. return recovered_data
  399. def set_many(
  400. self,
  401. data: Dict[Any, Any],
  402. timeout: Optional[float] = DEFAULT_TIMEOUT,
  403. version: Optional[int] = None,
  404. client: Optional[Redis] = None,
  405. ) -> None:
  406. """
  407. Set a bunch of values in the cache at once from a dict of key/value
  408. pairs. This is much more efficient than calling set() multiple times.
  409. If timeout is given, that timeout will be used for the key; otherwise
  410. the default cache timeout will be used.
  411. """
  412. if client is None:
  413. client = self.get_client(write=True)
  414. try:
  415. pipeline = client.pipeline()
  416. for key, value in data.items():
  417. self.set(key, value, timeout, version=version, client=pipeline)
  418. pipeline.execute()
  419. except _main_exceptions as e:
  420. raise ConnectionInterrupted(connection=client) from e
  421. def _incr(
  422. self,
  423. key: Any,
  424. delta: int = 1,
  425. version: Optional[int] = None,
  426. client: Optional[Redis] = None,
  427. ignore_key_check: bool = False,
  428. ) -> int:
  429. if client is None:
  430. client = self.get_client(write=True)
  431. key = self.make_key(key, version=version)
  432. try:
  433. try:
  434. # if key expired after exists check, then we get
  435. # key with wrong value and ttl -1.
  436. # use lua script for atomicity
  437. if not ignore_key_check:
  438. lua = """
  439. local exists = redis.call('EXISTS', KEYS[1])
  440. if (exists == 1) then
  441. return redis.call('INCRBY', KEYS[1], ARGV[1])
  442. else return false end
  443. """
  444. else:
  445. lua = """
  446. return redis.call('INCRBY', KEYS[1], ARGV[1])
  447. """
  448. value = client.eval(lua, 1, key, delta)
  449. if value is None:
  450. raise ValueError("Key '%s' not found" % key)
  451. except ResponseError:
  452. # if cached value or total value is greater than 64 bit signed
  453. # integer.
  454. # elif int is encoded. so redis sees the data as string.
  455. # In this situations redis will throw ResponseError
  456. # try to keep TTL of key
  457. timeout = self.ttl(key, version=version, client=client)
  458. # returns -2 if the key does not exist
  459. # means, that key have expired
  460. if timeout == -2:
  461. raise ValueError("Key '%s' not found" % key)
  462. value = self.get(key, version=version, client=client) + delta
  463. self.set(key, value, version=version, timeout=timeout, client=client)
  464. except _main_exceptions as e:
  465. raise ConnectionInterrupted(connection=client) from e
  466. return value
  467. def incr(
  468. self,
  469. key: Any,
  470. delta: int = 1,
  471. version: Optional[int] = None,
  472. client: Optional[Redis] = None,
  473. ignore_key_check: bool = False,
  474. ) -> int:
  475. """
  476. Add delta to value in the cache. If the key does not exist, raise a
  477. ValueError exception. if ignore_key_check=True then the key will be
  478. created and set to the delta value by default.
  479. """
  480. return self._incr(
  481. key=key,
  482. delta=delta,
  483. version=version,
  484. client=client,
  485. ignore_key_check=ignore_key_check,
  486. )
  487. def decr(
  488. self,
  489. key: Any,
  490. delta: int = 1,
  491. version: Optional[int] = None,
  492. client: Optional[Redis] = None,
  493. ) -> int:
  494. """
  495. Decreace delta to value in the cache. If the key does not exist, raise a
  496. ValueError exception.
  497. """
  498. return self._incr(key=key, delta=-delta, version=version, client=client)
  499. def ttl(
  500. self, key: Any, version: Optional[int] = None, client: Optional[Redis] = None
  501. ) -> Optional[int]:
  502. """
  503. Executes TTL redis command and return the "time-to-live" of specified key.
  504. If key is a non volatile key, it returns None.
  505. """
  506. if client is None:
  507. client = self.get_client(write=False)
  508. key = self.make_key(key, version=version)
  509. if not client.exists(key):
  510. return 0
  511. t = client.ttl(key)
  512. if t >= 0:
  513. return t
  514. elif t == -1:
  515. return None
  516. elif t == -2:
  517. return 0
  518. else:
  519. # Should never reach here
  520. return None
  521. def pttl(self, key, version=None, client=None):
  522. """
  523. Executes PTTL redis command and return the "time-to-live" of specified key.
  524. If key is a non volatile key, it returns None.
  525. """
  526. if client is None:
  527. client = self.get_client(write=False)
  528. key = self.make_key(key, version=version)
  529. if not client.exists(key):
  530. return 0
  531. t = client.pttl(key)
  532. if t >= 0:
  533. return t
  534. elif t == -1:
  535. return None
  536. elif t == -2:
  537. return 0
  538. else:
  539. # Should never reach here
  540. return None
  541. def has_key(
  542. self, key: Any, version: Optional[int] = None, client: Optional[Redis] = None
  543. ) -> bool:
  544. """
  545. Test if key exists.
  546. """
  547. if client is None:
  548. client = self.get_client(write=False)
  549. key = self.make_key(key, version=version)
  550. try:
  551. return client.exists(key) == 1
  552. except _main_exceptions as e:
  553. raise ConnectionInterrupted(connection=client) from e
  554. def iter_keys(
  555. self,
  556. search: str,
  557. itersize: Optional[int] = None,
  558. client: Optional[Redis] = None,
  559. version: Optional[int] = None,
  560. ) -> Iterator[str]:
  561. """
  562. Same as keys, but uses redis >= 2.8 cursors
  563. for make memory efficient keys iteration.
  564. """
  565. if client is None:
  566. client = self.get_client(write=False)
  567. pattern = self.make_pattern(search, version=version)
  568. for item in client.scan_iter(match=pattern, count=itersize):
  569. yield self.reverse_key(item.decode())
  570. def keys(
  571. self, search: str, version: Optional[int] = None, client: Optional[Redis] = None
  572. ) -> List[Any]:
  573. """
  574. Execute KEYS command and return matched results.
  575. Warning: this can return huge number of results, in
  576. this case, it strongly recommended use iter_keys
  577. for it.
  578. """
  579. if client is None:
  580. client = self.get_client(write=False)
  581. pattern = self.make_pattern(search, version=version)
  582. try:
  583. return [self.reverse_key(k.decode()) for k in client.keys(pattern)]
  584. except _main_exceptions as e:
  585. raise ConnectionInterrupted(connection=client) from e
  586. def make_key(
  587. self, key: Any, version: Optional[Any] = None, prefix: Optional[str] = None
  588. ) -> CacheKey:
  589. if isinstance(key, CacheKey):
  590. return key
  591. if prefix is None:
  592. prefix = self._backend.key_prefix
  593. if version is None:
  594. version = self._backend.version
  595. return CacheKey(self._backend.key_func(key, prefix, version))
  596. def make_pattern(
  597. self, pattern: str, version: Optional[int] = None, prefix: Optional[str] = None
  598. ) -> CacheKey:
  599. if isinstance(pattern, CacheKey):
  600. return pattern
  601. if prefix is None:
  602. prefix = self._backend.key_prefix
  603. prefix = glob_escape(prefix)
  604. if version is None:
  605. version = self._backend.version
  606. version_str = glob_escape(str(version))
  607. return CacheKey(self._backend.key_func(pattern, prefix, version_str))
  608. def close(self, **kwargs):
  609. close_flag = self._options.get(
  610. "CLOSE_CONNECTION",
  611. getattr(settings, "DJANGO_REDIS_CLOSE_CONNECTION", False),
  612. )
  613. if close_flag:
  614. self.do_close_clients()
  615. def do_close_clients(self):
  616. """default implementation: Override in custom client"""
  617. num_clients = len(self._clients)
  618. for idx in range(num_clients):
  619. self.disconnect(index=idx)
  620. self._clients = [None] * num_clients
  621. def touch(
  622. self,
  623. key: Any,
  624. timeout: Optional[float] = DEFAULT_TIMEOUT,
  625. version: Optional[int] = None,
  626. client: Optional[Redis] = None,
  627. ) -> bool:
  628. """
  629. Sets a new expiration for a key.
  630. """
  631. if timeout is DEFAULT_TIMEOUT:
  632. timeout = self._backend.default_timeout
  633. if client is None:
  634. client = self.get_client(write=True)
  635. key = self.make_key(key, version=version)
  636. if timeout is None:
  637. return bool(client.persist(key))
  638. else:
  639. # Convert to milliseconds
  640. timeout = int(timeout * 1000)
  641. return bool(client.pexpire(key, timeout))