cluster.py 91 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486
  1. import random
  2. import socket
  3. import sys
  4. import threading
  5. import time
  6. from collections import OrderedDict
  7. from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  8. from redis._parsers import CommandsParser, Encoder
  9. from redis._parsers.helpers import parse_scan
  10. from redis.backoff import default_backoff
  11. from redis.client import CaseInsensitiveDict, PubSub, Redis
  12. from redis.commands import READ_COMMANDS, RedisClusterCommands
  13. from redis.commands.helpers import list_or_args
  14. from redis.connection import ConnectionPool, DefaultParser, parse_url
  15. from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
  16. from redis.exceptions import (
  17. AskError,
  18. AuthenticationError,
  19. ClusterCrossSlotError,
  20. ClusterDownError,
  21. ClusterError,
  22. ConnectionError,
  23. DataError,
  24. MasterDownError,
  25. MovedError,
  26. RedisClusterException,
  27. RedisError,
  28. ResponseError,
  29. SlotNotCoveredError,
  30. TimeoutError,
  31. TryAgainError,
  32. )
  33. from redis.lock import Lock
  34. from redis.retry import Retry
  35. from redis.utils import (
  36. HIREDIS_AVAILABLE,
  37. dict_merge,
  38. list_keys_to_dict,
  39. merge_result,
  40. safe_str,
  41. str_if_bytes,
  42. )
  43. def get_node_name(host: str, port: Union[str, int]) -> str:
  44. return f"{host}:{port}"
  45. def get_connection(redis_node, *args, **options):
  46. return redis_node.connection or redis_node.connection_pool.get_connection(
  47. args[0], **options
  48. )
  49. def parse_scan_result(command, res, **options):
  50. cursors = {}
  51. ret = []
  52. for node_name, response in res.items():
  53. cursor, r = parse_scan(response, **options)
  54. cursors[node_name] = cursor
  55. ret += r
  56. return cursors, ret
  57. def parse_pubsub_numsub(command, res, **options):
  58. numsub_d = OrderedDict()
  59. for numsub_tups in res.values():
  60. for channel, numsubbed in numsub_tups:
  61. try:
  62. numsub_d[channel] += numsubbed
  63. except KeyError:
  64. numsub_d[channel] = numsubbed
  65. ret_numsub = [(channel, numsub) for channel, numsub in numsub_d.items()]
  66. return ret_numsub
  67. def parse_cluster_slots(
  68. resp: Any, **options: Any
  69. ) -> Dict[Tuple[int, int], Dict[str, Any]]:
  70. current_host = options.get("current_host", "")
  71. def fix_server(*args: Any) -> Tuple[str, Any]:
  72. return str_if_bytes(args[0]) or current_host, args[1]
  73. slots = {}
  74. for slot in resp:
  75. start, end, primary = slot[:3]
  76. replicas = slot[3:]
  77. slots[start, end] = {
  78. "primary": fix_server(*primary),
  79. "replicas": [fix_server(*replica) for replica in replicas],
  80. }
  81. return slots
  82. def parse_cluster_shards(resp, **options):
  83. """
  84. Parse CLUSTER SHARDS response.
  85. """
  86. if isinstance(resp[0], dict):
  87. return resp
  88. shards = []
  89. for x in resp:
  90. shard = {"slots": [], "nodes": []}
  91. for i in range(0, len(x[1]), 2):
  92. shard["slots"].append((x[1][i], (x[1][i + 1])))
  93. nodes = x[3]
  94. for node in nodes:
  95. dict_node = {}
  96. for i in range(0, len(node), 2):
  97. dict_node[node[i]] = node[i + 1]
  98. shard["nodes"].append(dict_node)
  99. shards.append(shard)
  100. return shards
  101. def parse_cluster_myshardid(resp, **options):
  102. """
  103. Parse CLUSTER MYSHARDID response.
  104. """
  105. return resp.decode("utf-8")
  106. PRIMARY = "primary"
  107. REPLICA = "replica"
  108. SLOT_ID = "slot-id"
  109. REDIS_ALLOWED_KEYS = (
  110. "charset",
  111. "connection_class",
  112. "connection_pool",
  113. "connection_pool_class",
  114. "client_name",
  115. "credential_provider",
  116. "db",
  117. "decode_responses",
  118. "encoding",
  119. "encoding_errors",
  120. "errors",
  121. "host",
  122. "lib_name",
  123. "lib_version",
  124. "max_connections",
  125. "nodes_flag",
  126. "redis_connect_func",
  127. "password",
  128. "port",
  129. "queue_class",
  130. "retry",
  131. "retry_on_timeout",
  132. "protocol",
  133. "socket_connect_timeout",
  134. "socket_keepalive",
  135. "socket_keepalive_options",
  136. "socket_timeout",
  137. "ssl",
  138. "ssl_ca_certs",
  139. "ssl_ca_data",
  140. "ssl_certfile",
  141. "ssl_cert_reqs",
  142. "ssl_keyfile",
  143. "ssl_password",
  144. "unix_socket_path",
  145. "username",
  146. )
  147. KWARGS_DISABLED_KEYS = ("host", "port")
  148. def cleanup_kwargs(**kwargs):
  149. """
  150. Remove unsupported or disabled keys from kwargs
  151. """
  152. connection_kwargs = {
  153. k: v
  154. for k, v in kwargs.items()
  155. if k in REDIS_ALLOWED_KEYS and k not in KWARGS_DISABLED_KEYS
  156. }
  157. return connection_kwargs
  158. class ClusterParser(DefaultParser):
  159. EXCEPTION_CLASSES = dict_merge(
  160. DefaultParser.EXCEPTION_CLASSES,
  161. {
  162. "ASK": AskError,
  163. "TRYAGAIN": TryAgainError,
  164. "MOVED": MovedError,
  165. "CLUSTERDOWN": ClusterDownError,
  166. "CROSSSLOT": ClusterCrossSlotError,
  167. "MASTERDOWN": MasterDownError,
  168. },
  169. )
  170. class AbstractRedisCluster:
  171. RedisClusterRequestTTL = 16
  172. PRIMARIES = "primaries"
  173. REPLICAS = "replicas"
  174. ALL_NODES = "all"
  175. RANDOM = "random"
  176. DEFAULT_NODE = "default-node"
  177. NODE_FLAGS = {PRIMARIES, REPLICAS, ALL_NODES, RANDOM, DEFAULT_NODE}
  178. COMMAND_FLAGS = dict_merge(
  179. list_keys_to_dict(
  180. [
  181. "ACL CAT",
  182. "ACL DELUSER",
  183. "ACL DRYRUN",
  184. "ACL GENPASS",
  185. "ACL GETUSER",
  186. "ACL HELP",
  187. "ACL LIST",
  188. "ACL LOG",
  189. "ACL LOAD",
  190. "ACL SAVE",
  191. "ACL SETUSER",
  192. "ACL USERS",
  193. "ACL WHOAMI",
  194. "AUTH",
  195. "CLIENT LIST",
  196. "CLIENT SETINFO",
  197. "CLIENT SETNAME",
  198. "CLIENT GETNAME",
  199. "CONFIG SET",
  200. "CONFIG REWRITE",
  201. "CONFIG RESETSTAT",
  202. "TIME",
  203. "PUBSUB CHANNELS",
  204. "PUBSUB NUMPAT",
  205. "PUBSUB NUMSUB",
  206. "PUBSUB SHARDCHANNELS",
  207. "PUBSUB SHARDNUMSUB",
  208. "PING",
  209. "INFO",
  210. "SHUTDOWN",
  211. "KEYS",
  212. "DBSIZE",
  213. "BGSAVE",
  214. "SLOWLOG GET",
  215. "SLOWLOG LEN",
  216. "SLOWLOG RESET",
  217. "WAIT",
  218. "WAITAOF",
  219. "SAVE",
  220. "MEMORY PURGE",
  221. "MEMORY MALLOC-STATS",
  222. "MEMORY STATS",
  223. "LASTSAVE",
  224. "CLIENT TRACKINGINFO",
  225. "CLIENT PAUSE",
  226. "CLIENT UNPAUSE",
  227. "CLIENT UNBLOCK",
  228. "CLIENT ID",
  229. "CLIENT REPLY",
  230. "CLIENT GETREDIR",
  231. "CLIENT INFO",
  232. "CLIENT KILL",
  233. "READONLY",
  234. "CLUSTER INFO",
  235. "CLUSTER MEET",
  236. "CLUSTER MYSHARDID",
  237. "CLUSTER NODES",
  238. "CLUSTER REPLICAS",
  239. "CLUSTER RESET",
  240. "CLUSTER SET-CONFIG-EPOCH",
  241. "CLUSTER SLOTS",
  242. "CLUSTER SHARDS",
  243. "CLUSTER COUNT-FAILURE-REPORTS",
  244. "CLUSTER KEYSLOT",
  245. "COMMAND",
  246. "COMMAND COUNT",
  247. "COMMAND LIST",
  248. "COMMAND GETKEYS",
  249. "CONFIG GET",
  250. "DEBUG",
  251. "RANDOMKEY",
  252. "READONLY",
  253. "READWRITE",
  254. "TIME",
  255. "TFUNCTION LOAD",
  256. "TFUNCTION DELETE",
  257. "TFUNCTION LIST",
  258. "TFCALL",
  259. "TFCALLASYNC",
  260. "GRAPH.CONFIG",
  261. "LATENCY HISTORY",
  262. "LATENCY LATEST",
  263. "LATENCY RESET",
  264. "MODULE LIST",
  265. "MODULE LOAD",
  266. "MODULE UNLOAD",
  267. "MODULE LOADEX",
  268. ],
  269. DEFAULT_NODE,
  270. ),
  271. list_keys_to_dict(
  272. [
  273. "FLUSHALL",
  274. "FLUSHDB",
  275. "FUNCTION DELETE",
  276. "FUNCTION FLUSH",
  277. "FUNCTION LIST",
  278. "FUNCTION LOAD",
  279. "FUNCTION RESTORE",
  280. "REDISGEARS_2.REFRESHCLUSTER",
  281. "SCAN",
  282. "SCRIPT EXISTS",
  283. "SCRIPT FLUSH",
  284. "SCRIPT LOAD",
  285. ],
  286. PRIMARIES,
  287. ),
  288. list_keys_to_dict(["FUNCTION DUMP"], RANDOM),
  289. list_keys_to_dict(
  290. [
  291. "CLUSTER COUNTKEYSINSLOT",
  292. "CLUSTER DELSLOTS",
  293. "CLUSTER DELSLOTSRANGE",
  294. "CLUSTER GETKEYSINSLOT",
  295. "CLUSTER SETSLOT",
  296. ],
  297. SLOT_ID,
  298. ),
  299. )
  300. SEARCH_COMMANDS = (
  301. [
  302. "FT.CREATE",
  303. "FT.SEARCH",
  304. "FT.AGGREGATE",
  305. "FT.EXPLAIN",
  306. "FT.EXPLAINCLI",
  307. "FT,PROFILE",
  308. "FT.ALTER",
  309. "FT.DROPINDEX",
  310. "FT.ALIASADD",
  311. "FT.ALIASUPDATE",
  312. "FT.ALIASDEL",
  313. "FT.TAGVALS",
  314. "FT.SUGADD",
  315. "FT.SUGGET",
  316. "FT.SUGDEL",
  317. "FT.SUGLEN",
  318. "FT.SYNUPDATE",
  319. "FT.SYNDUMP",
  320. "FT.SPELLCHECK",
  321. "FT.DICTADD",
  322. "FT.DICTDEL",
  323. "FT.DICTDUMP",
  324. "FT.INFO",
  325. "FT._LIST",
  326. "FT.CONFIG",
  327. "FT.ADD",
  328. "FT.DEL",
  329. "FT.DROP",
  330. "FT.GET",
  331. "FT.MGET",
  332. "FT.SYNADD",
  333. ],
  334. )
  335. CLUSTER_COMMANDS_RESPONSE_CALLBACKS = {
  336. "CLUSTER SLOTS": parse_cluster_slots,
  337. "CLUSTER SHARDS": parse_cluster_shards,
  338. "CLUSTER MYSHARDID": parse_cluster_myshardid,
  339. }
  340. RESULT_CALLBACKS = dict_merge(
  341. list_keys_to_dict(["PUBSUB NUMSUB", "PUBSUB SHARDNUMSUB"], parse_pubsub_numsub),
  342. list_keys_to_dict(
  343. ["PUBSUB NUMPAT"], lambda command, res: sum(list(res.values()))
  344. ),
  345. list_keys_to_dict(
  346. ["KEYS", "PUBSUB CHANNELS", "PUBSUB SHARDCHANNELS"], merge_result
  347. ),
  348. list_keys_to_dict(
  349. [
  350. "PING",
  351. "CONFIG SET",
  352. "CONFIG REWRITE",
  353. "CONFIG RESETSTAT",
  354. "CLIENT SETNAME",
  355. "BGSAVE",
  356. "SLOWLOG RESET",
  357. "SAVE",
  358. "MEMORY PURGE",
  359. "CLIENT PAUSE",
  360. "CLIENT UNPAUSE",
  361. ],
  362. lambda command, res: all(res.values()) if isinstance(res, dict) else res,
  363. ),
  364. list_keys_to_dict(
  365. ["DBSIZE", "WAIT"],
  366. lambda command, res: sum(res.values()) if isinstance(res, dict) else res,
  367. ),
  368. list_keys_to_dict(
  369. ["CLIENT UNBLOCK"], lambda command, res: 1 if sum(res.values()) > 0 else 0
  370. ),
  371. list_keys_to_dict(["SCAN"], parse_scan_result),
  372. list_keys_to_dict(
  373. ["SCRIPT LOAD"], lambda command, res: list(res.values()).pop()
  374. ),
  375. list_keys_to_dict(
  376. ["SCRIPT EXISTS"], lambda command, res: [all(k) for k in zip(*res.values())]
  377. ),
  378. list_keys_to_dict(["SCRIPT FLUSH"], lambda command, res: all(res.values())),
  379. )
  380. ERRORS_ALLOW_RETRY = (ConnectionError, TimeoutError, ClusterDownError)
  381. def replace_default_node(self, target_node: "ClusterNode" = None) -> None:
  382. """Replace the default cluster node.
  383. A random cluster node will be chosen if target_node isn't passed, and primaries
  384. will be prioritized. The default node will not be changed if there are no other
  385. nodes in the cluster.
  386. Args:
  387. target_node (ClusterNode, optional): Target node to replace the default
  388. node. Defaults to None.
  389. """
  390. if target_node:
  391. self.nodes_manager.default_node = target_node
  392. else:
  393. curr_node = self.get_default_node()
  394. primaries = [node for node in self.get_primaries() if node != curr_node]
  395. if primaries:
  396. # Choose a primary if the cluster contains different primaries
  397. self.nodes_manager.default_node = random.choice(primaries)
  398. else:
  399. # Otherwise, hoose a primary if the cluster contains different primaries
  400. replicas = [node for node in self.get_replicas() if node != curr_node]
  401. if replicas:
  402. self.nodes_manager.default_node = random.choice(replicas)
  403. class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
  404. @classmethod
  405. def from_url(cls, url, **kwargs):
  406. """
  407. Return a Redis client object configured from the given URL
  408. For example::
  409. redis://[[username]:[password]]@localhost:6379/0
  410. rediss://[[username]:[password]]@localhost:6379/0
  411. unix://[username@]/path/to/socket.sock?db=0[&password=password]
  412. Three URL schemes are supported:
  413. - `redis://` creates a TCP socket connection. See more at:
  414. <https://www.iana.org/assignments/uri-schemes/prov/redis>
  415. - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
  416. <https://www.iana.org/assignments/uri-schemes/prov/rediss>
  417. - ``unix://``: creates a Unix Domain Socket connection.
  418. The username, password, hostname, path and all querystring values
  419. are passed through urllib.parse.unquote in order to replace any
  420. percent-encoded values with their corresponding characters.
  421. There are several ways to specify a database number. The first value
  422. found will be used:
  423. 1. A ``db`` querystring option, e.g. redis://localhost?db=0
  424. 2. If using the redis:// or rediss:// schemes, the path argument
  425. of the url, e.g. redis://localhost/0
  426. 3. A ``db`` keyword argument to this function.
  427. If none of these options are specified, the default db=0 is used.
  428. All querystring options are cast to their appropriate Python types.
  429. Boolean arguments can be specified with string values "True"/"False"
  430. or "Yes"/"No". Values that cannot be properly cast cause a
  431. ``ValueError`` to be raised. Once parsed, the querystring arguments
  432. and keyword arguments are passed to the ``ConnectionPool``'s
  433. class initializer. In the case of conflicting arguments, querystring
  434. arguments always win.
  435. """
  436. return cls(url=url, **kwargs)
  437. def __init__(
  438. self,
  439. host: Optional[str] = None,
  440. port: int = 6379,
  441. startup_nodes: Optional[List["ClusterNode"]] = None,
  442. cluster_error_retry_attempts: int = 3,
  443. retry: Optional["Retry"] = None,
  444. require_full_coverage: bool = False,
  445. reinitialize_steps: int = 5,
  446. read_from_replicas: bool = False,
  447. dynamic_startup_nodes: bool = True,
  448. url: Optional[str] = None,
  449. address_remap: Optional[Callable[[str, int], Tuple[str, int]]] = None,
  450. **kwargs,
  451. ):
  452. """
  453. Initialize a new RedisCluster client.
  454. :param startup_nodes:
  455. List of nodes from which initial bootstrapping can be done
  456. :param host:
  457. Can be used to point to a startup node
  458. :param port:
  459. Can be used to point to a startup node
  460. :param require_full_coverage:
  461. When set to False (default value): the client will not require a
  462. full coverage of the slots. However, if not all slots are covered,
  463. and at least one node has 'cluster-require-full-coverage' set to
  464. 'yes,' the server will throw a ClusterDownError for some key-based
  465. commands. See -
  466. https://redis.io/topics/cluster-tutorial#redis-cluster-configuration-parameters
  467. When set to True: all slots must be covered to construct the
  468. cluster client. If not all slots are covered, RedisClusterException
  469. will be thrown.
  470. :param read_from_replicas:
  471. Enable read from replicas in READONLY mode. You can read possibly
  472. stale data.
  473. When set to true, read commands will be assigned between the
  474. primary and its replications in a Round-Robin manner.
  475. :param dynamic_startup_nodes:
  476. Set the RedisCluster's startup nodes to all of the discovered nodes.
  477. If true (default value), the cluster's discovered nodes will be used to
  478. determine the cluster nodes-slots mapping in the next topology refresh.
  479. It will remove the initial passed startup nodes if their endpoints aren't
  480. listed in the CLUSTER SLOTS output.
  481. If you use dynamic DNS endpoints for startup nodes but CLUSTER SLOTS lists
  482. specific IP addresses, it is best to set it to false.
  483. :param cluster_error_retry_attempts:
  484. Number of times to retry before raising an error when
  485. :class:`~.TimeoutError` or :class:`~.ConnectionError` or
  486. :class:`~.ClusterDownError` are encountered
  487. :param reinitialize_steps:
  488. Specifies the number of MOVED errors that need to occur before
  489. reinitializing the whole cluster topology. If a MOVED error occurs
  490. and the cluster does not need to be reinitialized on this current
  491. error handling, only the MOVED slot will be patched with the
  492. redirected node.
  493. To reinitialize the cluster on every MOVED error, set
  494. reinitialize_steps to 1.
  495. To avoid reinitializing the cluster on moved errors, set
  496. reinitialize_steps to 0.
  497. :param address_remap:
  498. An optional callable which, when provided with an internal network
  499. address of a node, e.g. a `(host, port)` tuple, will return the address
  500. where the node is reachable. This can be used to map the addresses at
  501. which the nodes _think_ they are, to addresses at which a client may
  502. reach them, such as when they sit behind a proxy.
  503. :**kwargs:
  504. Extra arguments that will be sent into Redis instance when created
  505. (See Official redis-py doc for supported kwargs
  506. [https://github.com/andymccurdy/redis-py/blob/master/redis/client.py])
  507. Some kwargs are not supported and will raise a
  508. RedisClusterException:
  509. - db (Redis do not support database SELECT in cluster mode)
  510. """
  511. if startup_nodes is None:
  512. startup_nodes = []
  513. if "db" in kwargs:
  514. # Argument 'db' is not possible to use in cluster mode
  515. raise RedisClusterException(
  516. "Argument 'db' is not possible to use in cluster mode"
  517. )
  518. # Get the startup node/s
  519. from_url = False
  520. if url is not None:
  521. from_url = True
  522. url_options = parse_url(url)
  523. if "path" in url_options:
  524. raise RedisClusterException(
  525. "RedisCluster does not currently support Unix Domain "
  526. "Socket connections"
  527. )
  528. if "db" in url_options and url_options["db"] != 0:
  529. # Argument 'db' is not possible to use in cluster mode
  530. raise RedisClusterException(
  531. "A ``db`` querystring option can only be 0 in cluster mode"
  532. )
  533. kwargs.update(url_options)
  534. host = kwargs.get("host")
  535. port = kwargs.get("port", port)
  536. startup_nodes.append(ClusterNode(host, port))
  537. elif host is not None and port is not None:
  538. startup_nodes.append(ClusterNode(host, port))
  539. elif len(startup_nodes) == 0:
  540. # No startup node was provided
  541. raise RedisClusterException(
  542. "RedisCluster requires at least one node to discover the "
  543. "cluster. Please provide one of the followings:\n"
  544. "1. host and port, for example:\n"
  545. " RedisCluster(host='localhost', port=6379)\n"
  546. "2. list of startup nodes, for example:\n"
  547. " RedisCluster(startup_nodes=[ClusterNode('localhost', 6379),"
  548. " ClusterNode('localhost', 6378)])"
  549. )
  550. # Update the connection arguments
  551. # Whenever a new connection is established, RedisCluster's on_connect
  552. # method should be run
  553. # If the user passed on_connect function we'll save it and run it
  554. # inside the RedisCluster.on_connect() function
  555. self.user_on_connect_func = kwargs.pop("redis_connect_func", None)
  556. kwargs.update({"redis_connect_func": self.on_connect})
  557. kwargs = cleanup_kwargs(**kwargs)
  558. if retry:
  559. self.retry = retry
  560. kwargs.update({"retry": self.retry})
  561. else:
  562. kwargs.update({"retry": Retry(default_backoff(), 0)})
  563. self.encoder = Encoder(
  564. kwargs.get("encoding", "utf-8"),
  565. kwargs.get("encoding_errors", "strict"),
  566. kwargs.get("decode_responses", False),
  567. )
  568. self.cluster_error_retry_attempts = cluster_error_retry_attempts
  569. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  570. self.node_flags = self.__class__.NODE_FLAGS.copy()
  571. self.read_from_replicas = read_from_replicas
  572. self.reinitialize_counter = 0
  573. self.reinitialize_steps = reinitialize_steps
  574. self.nodes_manager = NodesManager(
  575. startup_nodes=startup_nodes,
  576. from_url=from_url,
  577. require_full_coverage=require_full_coverage,
  578. dynamic_startup_nodes=dynamic_startup_nodes,
  579. address_remap=address_remap,
  580. **kwargs,
  581. )
  582. self.cluster_response_callbacks = CaseInsensitiveDict(
  583. self.__class__.CLUSTER_COMMANDS_RESPONSE_CALLBACKS
  584. )
  585. self.result_callbacks = CaseInsensitiveDict(self.__class__.RESULT_CALLBACKS)
  586. self.commands_parser = CommandsParser(self)
  587. self._lock = threading.Lock()
  588. def __enter__(self):
  589. return self
  590. def __exit__(self, exc_type, exc_value, traceback):
  591. self.close()
  592. def __del__(self):
  593. self.close()
  594. def disconnect_connection_pools(self):
  595. for node in self.get_nodes():
  596. if node.redis_connection:
  597. try:
  598. node.redis_connection.connection_pool.disconnect()
  599. except OSError:
  600. # Client was already disconnected. do nothing
  601. pass
  602. def on_connect(self, connection):
  603. """
  604. Initialize the connection, authenticate and select a database and send
  605. READONLY if it is set during object initialization.
  606. """
  607. connection.set_parser(ClusterParser)
  608. connection.on_connect()
  609. if self.read_from_replicas:
  610. # Sending READONLY command to server to configure connection as
  611. # readonly. Since each cluster node may change its server type due
  612. # to a failover, we should establish a READONLY connection
  613. # regardless of the server type. If this is a primary connection,
  614. # READONLY would not affect executing write commands.
  615. connection.send_command("READONLY")
  616. if str_if_bytes(connection.read_response()) != "OK":
  617. raise ConnectionError("READONLY command failed")
  618. if self.user_on_connect_func is not None:
  619. self.user_on_connect_func(connection)
  620. def get_redis_connection(self, node):
  621. if not node.redis_connection:
  622. with self._lock:
  623. if not node.redis_connection:
  624. self.nodes_manager.create_redis_connections([node])
  625. return node.redis_connection
  626. def get_node(self, host=None, port=None, node_name=None):
  627. return self.nodes_manager.get_node(host, port, node_name)
  628. def get_primaries(self):
  629. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  630. def get_replicas(self):
  631. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  632. def get_random_node(self):
  633. return random.choice(list(self.nodes_manager.nodes_cache.values()))
  634. def get_nodes(self):
  635. return list(self.nodes_manager.nodes_cache.values())
  636. def get_node_from_key(self, key, replica=False):
  637. """
  638. Get the node that holds the key's slot.
  639. If replica set to True but the slot doesn't have any replicas, None is
  640. returned.
  641. """
  642. slot = self.keyslot(key)
  643. slot_cache = self.nodes_manager.slots_cache.get(slot)
  644. if slot_cache is None or len(slot_cache) == 0:
  645. raise SlotNotCoveredError(f'Slot "{slot}" is not covered by the cluster.')
  646. if replica and len(self.nodes_manager.slots_cache[slot]) < 2:
  647. return None
  648. elif replica:
  649. node_idx = 1
  650. else:
  651. # primary
  652. node_idx = 0
  653. return slot_cache[node_idx]
  654. def get_default_node(self):
  655. """
  656. Get the cluster's default node
  657. """
  658. return self.nodes_manager.default_node
  659. def set_default_node(self, node):
  660. """
  661. Set the default node of the cluster.
  662. :param node: 'ClusterNode'
  663. :return True if the default node was set, else False
  664. """
  665. if node is None or self.get_node(node_name=node.name) is None:
  666. return False
  667. self.nodes_manager.default_node = node
  668. return True
  669. def get_retry(self) -> Optional["Retry"]:
  670. return self.retry
  671. def set_retry(self, retry: "Retry") -> None:
  672. self.retry = retry
  673. for node in self.get_nodes():
  674. node.redis_connection.set_retry(retry)
  675. def monitor(self, target_node=None):
  676. """
  677. Returns a Monitor object for the specified target node.
  678. The default cluster node will be selected if no target node was
  679. specified.
  680. Monitor is useful for handling the MONITOR command to the redis server.
  681. next_command() method returns one command from monitor
  682. listen() method yields commands from monitor.
  683. """
  684. if target_node is None:
  685. target_node = self.get_default_node()
  686. if target_node.redis_connection is None:
  687. raise RedisClusterException(
  688. f"Cluster Node {target_node.name} has no redis_connection"
  689. )
  690. return target_node.redis_connection.monitor()
  691. def pubsub(self, node=None, host=None, port=None, **kwargs):
  692. """
  693. Allows passing a ClusterNode, or host&port, to get a pubsub instance
  694. connected to the specified node
  695. """
  696. return ClusterPubSub(self, node=node, host=host, port=port, **kwargs)
  697. def pipeline(self, transaction=None, shard_hint=None):
  698. """
  699. Cluster impl:
  700. Pipelines do not work in cluster mode the same way they
  701. do in normal mode. Create a clone of this object so
  702. that simulating pipelines will work correctly. Each
  703. command will be called directly when used and
  704. when calling execute() will only return the result stack.
  705. """
  706. if shard_hint:
  707. raise RedisClusterException("shard_hint is deprecated in cluster mode")
  708. if transaction:
  709. raise RedisClusterException("transaction is deprecated in cluster mode")
  710. return ClusterPipeline(
  711. nodes_manager=self.nodes_manager,
  712. commands_parser=self.commands_parser,
  713. startup_nodes=self.nodes_manager.startup_nodes,
  714. result_callbacks=self.result_callbacks,
  715. cluster_response_callbacks=self.cluster_response_callbacks,
  716. cluster_error_retry_attempts=self.cluster_error_retry_attempts,
  717. read_from_replicas=self.read_from_replicas,
  718. reinitialize_steps=self.reinitialize_steps,
  719. lock=self._lock,
  720. )
  721. def lock(
  722. self,
  723. name,
  724. timeout=None,
  725. sleep=0.1,
  726. blocking=True,
  727. blocking_timeout=None,
  728. lock_class=None,
  729. thread_local=True,
  730. ):
  731. """
  732. Return a new Lock object using key ``name`` that mimics
  733. the behavior of threading.Lock.
  734. If specified, ``timeout`` indicates a maximum life for the lock.
  735. By default, it will remain locked until release() is called.
  736. ``sleep`` indicates the amount of time to sleep per loop iteration
  737. when the lock is in blocking mode and another client is currently
  738. holding the lock.
  739. ``blocking`` indicates whether calling ``acquire`` should block until
  740. the lock has been acquired or to fail immediately, causing ``acquire``
  741. to return False and the lock not being acquired. Defaults to True.
  742. Note this value can be overridden by passing a ``blocking``
  743. argument to ``acquire``.
  744. ``blocking_timeout`` indicates the maximum amount of time in seconds to
  745. spend trying to acquire the lock. A value of ``None`` indicates
  746. continue trying forever. ``blocking_timeout`` can be specified as a
  747. float or integer, both representing the number of seconds to wait.
  748. ``lock_class`` forces the specified lock implementation. Note that as
  749. of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
  750. a Lua-based lock). So, it's unlikely you'll need this parameter, unless
  751. you have created your own custom lock class.
  752. ``thread_local`` indicates whether the lock token is placed in
  753. thread-local storage. By default, the token is placed in thread local
  754. storage so that a thread only sees its token, not a token set by
  755. another thread. Consider the following timeline:
  756. time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
  757. thread-1 sets the token to "abc"
  758. time: 1, thread-2 blocks trying to acquire `my-lock` using the
  759. Lock instance.
  760. time: 5, thread-1 has not yet completed. redis expires the lock
  761. key.
  762. time: 5, thread-2 acquired `my-lock` now that it's available.
  763. thread-2 sets the token to "xyz"
  764. time: 6, thread-1 finishes its work and calls release(). if the
  765. token is *not* stored in thread local storage, then
  766. thread-1 would see the token value as "xyz" and would be
  767. able to successfully release the thread-2's lock.
  768. In some use cases it's necessary to disable thread local storage. For
  769. example, if you have code where one thread acquires a lock and passes
  770. that lock instance to a worker thread to release later. If thread
  771. local storage isn't disabled in this case, the worker thread won't see
  772. the token set by the thread that acquired the lock. Our assumption
  773. is that these cases aren't common and as such default to using
  774. thread local storage."""
  775. if lock_class is None:
  776. lock_class = Lock
  777. return lock_class(
  778. self,
  779. name,
  780. timeout=timeout,
  781. sleep=sleep,
  782. blocking=blocking,
  783. blocking_timeout=blocking_timeout,
  784. thread_local=thread_local,
  785. )
  786. def set_response_callback(self, command, callback):
  787. """Set a custom Response Callback"""
  788. self.cluster_response_callbacks[command] = callback
  789. def _determine_nodes(self, *args, **kwargs) -> List["ClusterNode"]:
  790. # Determine which nodes should be executed the command on.
  791. # Returns a list of target nodes.
  792. command = args[0].upper()
  793. if len(args) >= 2 and f"{args[0]} {args[1]}".upper() in self.command_flags:
  794. command = f"{args[0]} {args[1]}".upper()
  795. nodes_flag = kwargs.pop("nodes_flag", None)
  796. if nodes_flag is not None:
  797. # nodes flag passed by the user
  798. command_flag = nodes_flag
  799. else:
  800. # get the nodes group for this command if it was predefined
  801. command_flag = self.command_flags.get(command)
  802. if command_flag == self.__class__.RANDOM:
  803. # return a random node
  804. return [self.get_random_node()]
  805. elif command_flag == self.__class__.PRIMARIES:
  806. # return all primaries
  807. return self.get_primaries()
  808. elif command_flag == self.__class__.REPLICAS:
  809. # return all replicas
  810. return self.get_replicas()
  811. elif command_flag == self.__class__.ALL_NODES:
  812. # return all nodes
  813. return self.get_nodes()
  814. elif command_flag == self.__class__.DEFAULT_NODE:
  815. # return the cluster's default node
  816. return [self.nodes_manager.default_node]
  817. elif command in self.__class__.SEARCH_COMMANDS[0]:
  818. return [self.nodes_manager.default_node]
  819. else:
  820. # get the node that holds the key's slot
  821. slot = self.determine_slot(*args)
  822. node = self.nodes_manager.get_node_from_slot(
  823. slot, self.read_from_replicas and command in READ_COMMANDS
  824. )
  825. return [node]
  826. def _should_reinitialized(self):
  827. # To reinitialize the cluster on every MOVED error,
  828. # set reinitialize_steps to 1.
  829. # To avoid reinitializing the cluster on moved errors, set
  830. # reinitialize_steps to 0.
  831. if self.reinitialize_steps == 0:
  832. return False
  833. else:
  834. return self.reinitialize_counter % self.reinitialize_steps == 0
  835. def keyslot(self, key):
  836. """
  837. Calculate keyslot for a given key.
  838. See Keys distribution model in https://redis.io/topics/cluster-spec
  839. """
  840. k = self.encoder.encode(key)
  841. return key_slot(k)
  842. def _get_command_keys(self, *args):
  843. """
  844. Get the keys in the command. If the command has no keys in in, None is
  845. returned.
  846. NOTE: Due to a bug in redis<7.0, this function does not work properly
  847. for EVAL or EVALSHA when the `numkeys` arg is 0.
  848. - issue: https://github.com/redis/redis/issues/9493
  849. - fix: https://github.com/redis/redis/pull/9733
  850. So, don't use this function with EVAL or EVALSHA.
  851. """
  852. redis_conn = self.get_default_node().redis_connection
  853. return self.commands_parser.get_keys(redis_conn, *args)
  854. def determine_slot(self, *args):
  855. """
  856. Figure out what slot to use based on args.
  857. Raises a RedisClusterException if there's a missing key and we can't
  858. determine what slots to map the command to; or, if the keys don't
  859. all map to the same key slot.
  860. """
  861. command = args[0]
  862. if self.command_flags.get(command) == SLOT_ID:
  863. # The command contains the slot ID
  864. return args[1]
  865. # Get the keys in the command
  866. # EVAL and EVALSHA are common enough that it's wasteful to go to the
  867. # redis server to parse the keys. Besides, there is a bug in redis<7.0
  868. # where `self._get_command_keys()` fails anyway. So, we special case
  869. # EVAL/EVALSHA.
  870. if command.upper() in ("EVAL", "EVALSHA"):
  871. # command syntax: EVAL "script body" num_keys ...
  872. if len(args) <= 2:
  873. raise RedisClusterException(f"Invalid args in command: {args}")
  874. num_actual_keys = int(args[2])
  875. eval_keys = args[3 : 3 + num_actual_keys]
  876. # if there are 0 keys, that means the script can be run on any node
  877. # so we can just return a random slot
  878. if len(eval_keys) == 0:
  879. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  880. keys = eval_keys
  881. else:
  882. keys = self._get_command_keys(*args)
  883. if keys is None or len(keys) == 0:
  884. # FCALL can call a function with 0 keys, that means the function
  885. # can be run on any node so we can just return a random slot
  886. if command.upper() in ("FCALL", "FCALL_RO"):
  887. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  888. raise RedisClusterException(
  889. "No way to dispatch this command to Redis Cluster. "
  890. "Missing key.\nYou can execute the command by specifying "
  891. f"target nodes.\nCommand: {args}"
  892. )
  893. # single key command
  894. if len(keys) == 1:
  895. return self.keyslot(keys[0])
  896. # multi-key command; we need to make sure all keys are mapped to
  897. # the same slot
  898. slots = {self.keyslot(key) for key in keys}
  899. if len(slots) != 1:
  900. raise RedisClusterException(
  901. f"{command} - all keys must map to the same key slot"
  902. )
  903. return slots.pop()
  904. def get_encoder(self):
  905. """
  906. Get the connections' encoder
  907. """
  908. return self.encoder
  909. def get_connection_kwargs(self):
  910. """
  911. Get the connections' key-word arguments
  912. """
  913. return self.nodes_manager.connection_kwargs
  914. def _is_nodes_flag(self, target_nodes):
  915. return isinstance(target_nodes, str) and target_nodes in self.node_flags
  916. def _parse_target_nodes(self, target_nodes):
  917. if isinstance(target_nodes, list):
  918. nodes = target_nodes
  919. elif isinstance(target_nodes, ClusterNode):
  920. # Supports passing a single ClusterNode as a variable
  921. nodes = [target_nodes]
  922. elif isinstance(target_nodes, dict):
  923. # Supports dictionaries of the format {node_name: node}.
  924. # It enables to execute commands with multi nodes as follows:
  925. # rc.cluster_save_config(rc.get_primaries())
  926. nodes = target_nodes.values()
  927. else:
  928. raise TypeError(
  929. "target_nodes type can be one of the following: "
  930. "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
  931. "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
  932. f"The passed type is {type(target_nodes)}"
  933. )
  934. return nodes
  935. def execute_command(self, *args, **kwargs):
  936. """
  937. Wrapper for ERRORS_ALLOW_RETRY error handling.
  938. It will try the number of times specified by the config option
  939. "self.cluster_error_retry_attempts" which defaults to 3 unless manually
  940. configured.
  941. If it reaches the number of times, the command will raise the exception
  942. Key argument :target_nodes: can be passed with the following types:
  943. nodes_flag: PRIMARIES, REPLICAS, ALL_NODES, RANDOM
  944. ClusterNode
  945. list<ClusterNode>
  946. dict<Any, ClusterNode>
  947. """
  948. target_nodes_specified = False
  949. is_default_node = False
  950. target_nodes = None
  951. passed_targets = kwargs.pop("target_nodes", None)
  952. if passed_targets is not None and not self._is_nodes_flag(passed_targets):
  953. target_nodes = self._parse_target_nodes(passed_targets)
  954. target_nodes_specified = True
  955. # If an error that allows retrying was thrown, the nodes and slots
  956. # cache were reinitialized. We will retry executing the command with
  957. # the updated cluster setup only when the target nodes can be
  958. # determined again with the new cache tables. Therefore, when target
  959. # nodes were passed to this function, we cannot retry the command
  960. # execution since the nodes may not be valid anymore after the tables
  961. # were reinitialized. So in case of passed target nodes,
  962. # retry_attempts will be set to 0.
  963. retry_attempts = (
  964. 0 if target_nodes_specified else self.cluster_error_retry_attempts
  965. )
  966. # Add one for the first execution
  967. execute_attempts = 1 + retry_attempts
  968. for _ in range(execute_attempts):
  969. try:
  970. res = {}
  971. if not target_nodes_specified:
  972. # Determine the nodes to execute the command on
  973. target_nodes = self._determine_nodes(
  974. *args, **kwargs, nodes_flag=passed_targets
  975. )
  976. if not target_nodes:
  977. raise RedisClusterException(
  978. f"No targets were found to execute {args} command on"
  979. )
  980. if (
  981. len(target_nodes) == 1
  982. and target_nodes[0] == self.get_default_node()
  983. ):
  984. is_default_node = True
  985. for node in target_nodes:
  986. res[node.name] = self._execute_command(node, *args, **kwargs)
  987. # Return the processed result
  988. return self._process_result(args[0], res, **kwargs)
  989. except Exception as e:
  990. if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  991. if is_default_node:
  992. # Replace the default cluster node
  993. self.replace_default_node()
  994. # The nodes and slots cache were reinitialized.
  995. # Try again with the new cluster setup.
  996. retry_attempts -= 1
  997. continue
  998. else:
  999. # raise the exception
  1000. raise e
  1001. def _execute_command(self, target_node, *args, **kwargs):
  1002. """
  1003. Send a command to a node in the cluster
  1004. """
  1005. command = args[0]
  1006. redis_node = None
  1007. connection = None
  1008. redirect_addr = None
  1009. asking = False
  1010. moved = False
  1011. ttl = int(self.RedisClusterRequestTTL)
  1012. while ttl > 0:
  1013. ttl -= 1
  1014. try:
  1015. if asking:
  1016. target_node = self.get_node(node_name=redirect_addr)
  1017. elif moved:
  1018. # MOVED occurred and the slots cache was updated,
  1019. # refresh the target node
  1020. slot = self.determine_slot(*args)
  1021. target_node = self.nodes_manager.get_node_from_slot(
  1022. slot, self.read_from_replicas and command in READ_COMMANDS
  1023. )
  1024. moved = False
  1025. redis_node = self.get_redis_connection(target_node)
  1026. connection = get_connection(redis_node, *args, **kwargs)
  1027. if asking:
  1028. connection.send_command("ASKING")
  1029. redis_node.parse_response(connection, "ASKING", **kwargs)
  1030. asking = False
  1031. connection.send_command(*args)
  1032. response = redis_node.parse_response(connection, command, **kwargs)
  1033. if command in self.cluster_response_callbacks:
  1034. response = self.cluster_response_callbacks[command](
  1035. response, **kwargs
  1036. )
  1037. return response
  1038. except AuthenticationError:
  1039. raise
  1040. except (ConnectionError, TimeoutError) as e:
  1041. # Connection retries are being handled in the node's
  1042. # Retry object.
  1043. # ConnectionError can also be raised if we couldn't get a
  1044. # connection from the pool before timing out, so check that
  1045. # this is an actual connection before attempting to disconnect.
  1046. if connection is not None:
  1047. connection.disconnect()
  1048. # Remove the failed node from the startup nodes before we try
  1049. # to reinitialize the cluster
  1050. self.nodes_manager.startup_nodes.pop(target_node.name, None)
  1051. # Reset the cluster node's connection
  1052. target_node.redis_connection = None
  1053. self.nodes_manager.initialize()
  1054. raise e
  1055. except MovedError as e:
  1056. # First, we will try to patch the slots/nodes cache with the
  1057. # redirected node output and try again. If MovedError exceeds
  1058. # 'reinitialize_steps' number of times, we will force
  1059. # reinitializing the tables, and then try again.
  1060. # 'reinitialize_steps' counter will increase faster when
  1061. # the same client object is shared between multiple threads. To
  1062. # reduce the frequency you can set this variable in the
  1063. # RedisCluster constructor.
  1064. self.reinitialize_counter += 1
  1065. if self._should_reinitialized():
  1066. self.nodes_manager.initialize()
  1067. # Reset the counter
  1068. self.reinitialize_counter = 0
  1069. else:
  1070. self.nodes_manager.update_moved_exception(e)
  1071. moved = True
  1072. except TryAgainError:
  1073. if ttl < self.RedisClusterRequestTTL / 2:
  1074. time.sleep(0.05)
  1075. except AskError as e:
  1076. redirect_addr = get_node_name(host=e.host, port=e.port)
  1077. asking = True
  1078. except ClusterDownError as e:
  1079. # ClusterDownError can occur during a failover and to get
  1080. # self-healed, we will try to reinitialize the cluster layout
  1081. # and retry executing the command
  1082. time.sleep(0.25)
  1083. self.nodes_manager.initialize()
  1084. raise e
  1085. except ResponseError:
  1086. raise
  1087. except Exception as e:
  1088. if connection:
  1089. connection.disconnect()
  1090. raise e
  1091. finally:
  1092. if connection is not None:
  1093. redis_node.connection_pool.release(connection)
  1094. raise ClusterError("TTL exhausted.")
  1095. def close(self):
  1096. try:
  1097. with self._lock:
  1098. if self.nodes_manager:
  1099. self.nodes_manager.close()
  1100. except AttributeError:
  1101. # RedisCluster's __init__ can fail before nodes_manager is set
  1102. pass
  1103. def _process_result(self, command, res, **kwargs):
  1104. """
  1105. Process the result of the executed command.
  1106. The function would return a dict or a single value.
  1107. :type command: str
  1108. :type res: dict
  1109. `res` should be in the following format:
  1110. Dict<node_name, command_result>
  1111. """
  1112. if command in self.result_callbacks:
  1113. return self.result_callbacks[command](command, res, **kwargs)
  1114. elif len(res) == 1:
  1115. # When we execute the command on a single node, we can
  1116. # remove the dictionary and return a single response
  1117. return list(res.values())[0]
  1118. else:
  1119. return res
  1120. def load_external_module(self, funcname, func):
  1121. """
  1122. This function can be used to add externally defined redis modules,
  1123. and their namespaces to the redis client.
  1124. ``funcname`` - A string containing the name of the function to create
  1125. ``func`` - The function, being added to this class.
  1126. """
  1127. setattr(self, funcname, func)
  1128. class ClusterNode:
  1129. def __init__(self, host, port, server_type=None, redis_connection=None):
  1130. if host == "localhost":
  1131. host = socket.gethostbyname(host)
  1132. self.host = host
  1133. self.port = port
  1134. self.name = get_node_name(host, port)
  1135. self.server_type = server_type
  1136. self.redis_connection = redis_connection
  1137. def __repr__(self):
  1138. return (
  1139. f"[host={self.host},"
  1140. f"port={self.port},"
  1141. f"name={self.name},"
  1142. f"server_type={self.server_type},"
  1143. f"redis_connection={self.redis_connection}]"
  1144. )
  1145. def __eq__(self, obj):
  1146. return isinstance(obj, ClusterNode) and obj.name == self.name
  1147. def __del__(self):
  1148. if self.redis_connection is not None:
  1149. self.redis_connection.close()
  1150. class LoadBalancer:
  1151. """
  1152. Round-Robin Load Balancing
  1153. """
  1154. def __init__(self, start_index: int = 0) -> None:
  1155. self.primary_to_idx = {}
  1156. self.start_index = start_index
  1157. def get_server_index(self, primary: str, list_size: int) -> int:
  1158. server_index = self.primary_to_idx.setdefault(primary, self.start_index)
  1159. # Update the index
  1160. self.primary_to_idx[primary] = (server_index + 1) % list_size
  1161. return server_index
  1162. def reset(self) -> None:
  1163. self.primary_to_idx.clear()
  1164. class NodesManager:
  1165. def __init__(
  1166. self,
  1167. startup_nodes,
  1168. from_url=False,
  1169. require_full_coverage=False,
  1170. lock=None,
  1171. dynamic_startup_nodes=True,
  1172. connection_pool_class=ConnectionPool,
  1173. address_remap: Optional[Callable[[str, int], Tuple[str, int]]] = None,
  1174. **kwargs,
  1175. ):
  1176. self.nodes_cache = {}
  1177. self.slots_cache = {}
  1178. self.startup_nodes = {}
  1179. self.default_node = None
  1180. self.populate_startup_nodes(startup_nodes)
  1181. self.from_url = from_url
  1182. self._require_full_coverage = require_full_coverage
  1183. self._dynamic_startup_nodes = dynamic_startup_nodes
  1184. self.connection_pool_class = connection_pool_class
  1185. self.address_remap = address_remap
  1186. self._moved_exception = None
  1187. self.connection_kwargs = kwargs
  1188. self.read_load_balancer = LoadBalancer()
  1189. if lock is None:
  1190. lock = threading.Lock()
  1191. self._lock = lock
  1192. self.initialize()
  1193. def get_node(self, host=None, port=None, node_name=None):
  1194. """
  1195. Get the requested node from the cluster's nodes.
  1196. nodes.
  1197. :return: ClusterNode if the node exists, else None
  1198. """
  1199. if host and port:
  1200. # the user passed host and port
  1201. if host == "localhost":
  1202. host = socket.gethostbyname(host)
  1203. return self.nodes_cache.get(get_node_name(host=host, port=port))
  1204. elif node_name:
  1205. return self.nodes_cache.get(node_name)
  1206. else:
  1207. return None
  1208. def update_moved_exception(self, exception):
  1209. self._moved_exception = exception
  1210. def _update_moved_slots(self):
  1211. """
  1212. Update the slot's node with the redirected one
  1213. """
  1214. e = self._moved_exception
  1215. redirected_node = self.get_node(host=e.host, port=e.port)
  1216. if redirected_node is not None:
  1217. # The node already exists
  1218. if redirected_node.server_type is not PRIMARY:
  1219. # Update the node's server type
  1220. redirected_node.server_type = PRIMARY
  1221. else:
  1222. # This is a new node, we will add it to the nodes cache
  1223. redirected_node = ClusterNode(e.host, e.port, PRIMARY)
  1224. self.nodes_cache[redirected_node.name] = redirected_node
  1225. if redirected_node in self.slots_cache[e.slot_id]:
  1226. # The MOVED error resulted from a failover, and the new slot owner
  1227. # had previously been a replica.
  1228. old_primary = self.slots_cache[e.slot_id][0]
  1229. # Update the old primary to be a replica and add it to the end of
  1230. # the slot's node list
  1231. old_primary.server_type = REPLICA
  1232. self.slots_cache[e.slot_id].append(old_primary)
  1233. # Remove the old replica, which is now a primary, from the slot's
  1234. # node list
  1235. self.slots_cache[e.slot_id].remove(redirected_node)
  1236. # Override the old primary with the new one
  1237. self.slots_cache[e.slot_id][0] = redirected_node
  1238. if self.default_node == old_primary:
  1239. # Update the default node with the new primary
  1240. self.default_node = redirected_node
  1241. else:
  1242. # The new slot owner is a new server, or a server from a different
  1243. # shard. We need to remove all current nodes from the slot's list
  1244. # (including replications) and add just the new node.
  1245. self.slots_cache[e.slot_id] = [redirected_node]
  1246. # Reset moved_exception
  1247. self._moved_exception = None
  1248. def get_node_from_slot(self, slot, read_from_replicas=False, server_type=None):
  1249. """
  1250. Gets a node that servers this hash slot
  1251. """
  1252. if self._moved_exception:
  1253. with self._lock:
  1254. if self._moved_exception:
  1255. self._update_moved_slots()
  1256. if self.slots_cache.get(slot) is None or len(self.slots_cache[slot]) == 0:
  1257. raise SlotNotCoveredError(
  1258. f'Slot "{slot}" not covered by the cluster. '
  1259. f'"require_full_coverage={self._require_full_coverage}"'
  1260. )
  1261. if read_from_replicas is True:
  1262. # get the server index in a Round-Robin manner
  1263. primary_name = self.slots_cache[slot][0].name
  1264. node_idx = self.read_load_balancer.get_server_index(
  1265. primary_name, len(self.slots_cache[slot])
  1266. )
  1267. elif (
  1268. server_type is None
  1269. or server_type == PRIMARY
  1270. or len(self.slots_cache[slot]) == 1
  1271. ):
  1272. # return a primary
  1273. node_idx = 0
  1274. else:
  1275. # return a replica
  1276. # randomly choose one of the replicas
  1277. node_idx = random.randint(1, len(self.slots_cache[slot]) - 1)
  1278. return self.slots_cache[slot][node_idx]
  1279. def get_nodes_by_server_type(self, server_type):
  1280. """
  1281. Get all nodes with the specified server type
  1282. :param server_type: 'primary' or 'replica'
  1283. :return: list of ClusterNode
  1284. """
  1285. return [
  1286. node
  1287. for node in self.nodes_cache.values()
  1288. if node.server_type == server_type
  1289. ]
  1290. def populate_startup_nodes(self, nodes):
  1291. """
  1292. Populate all startup nodes and filters out any duplicates
  1293. """
  1294. for n in nodes:
  1295. self.startup_nodes[n.name] = n
  1296. def check_slots_coverage(self, slots_cache):
  1297. # Validate if all slots are covered or if we should try next
  1298. # startup node
  1299. for i in range(0, REDIS_CLUSTER_HASH_SLOTS):
  1300. if i not in slots_cache:
  1301. return False
  1302. return True
  1303. def create_redis_connections(self, nodes):
  1304. """
  1305. This function will create a redis connection to all nodes in :nodes:
  1306. """
  1307. for node in nodes:
  1308. if node.redis_connection is None:
  1309. node.redis_connection = self.create_redis_node(
  1310. host=node.host, port=node.port, **self.connection_kwargs
  1311. )
  1312. def create_redis_node(self, host, port, **kwargs):
  1313. if self.from_url:
  1314. # Create a redis node with a costumed connection pool
  1315. kwargs.update({"host": host})
  1316. kwargs.update({"port": port})
  1317. r = Redis(connection_pool=self.connection_pool_class(**kwargs))
  1318. else:
  1319. r = Redis(host=host, port=port, **kwargs)
  1320. return r
  1321. def _get_or_create_cluster_node(self, host, port, role, tmp_nodes_cache):
  1322. node_name = get_node_name(host, port)
  1323. # check if we already have this node in the tmp_nodes_cache
  1324. target_node = tmp_nodes_cache.get(node_name)
  1325. if target_node is None:
  1326. # before creating a new cluster node, check if the cluster node already
  1327. # exists in the current nodes cache and has a valid connection so we can
  1328. # reuse it
  1329. target_node = self.nodes_cache.get(node_name)
  1330. if target_node is None or target_node.redis_connection is None:
  1331. # create new cluster node for this cluster
  1332. target_node = ClusterNode(host, port, role)
  1333. if target_node.server_type != role:
  1334. target_node.server_type = role
  1335. return target_node
  1336. def initialize(self):
  1337. """
  1338. Initializes the nodes cache, slots cache and redis connections.
  1339. :startup_nodes:
  1340. Responsible for discovering other nodes in the cluster
  1341. """
  1342. self.reset()
  1343. tmp_nodes_cache = {}
  1344. tmp_slots = {}
  1345. disagreements = []
  1346. startup_nodes_reachable = False
  1347. fully_covered = False
  1348. kwargs = self.connection_kwargs
  1349. exception = None
  1350. for startup_node in self.startup_nodes.values():
  1351. try:
  1352. if startup_node.redis_connection:
  1353. r = startup_node.redis_connection
  1354. else:
  1355. # Create a new Redis connection
  1356. r = self.create_redis_node(
  1357. startup_node.host, startup_node.port, **kwargs
  1358. )
  1359. self.startup_nodes[startup_node.name].redis_connection = r
  1360. # Make sure cluster mode is enabled on this node
  1361. if bool(r.info().get("cluster_enabled")) is False:
  1362. raise RedisClusterException(
  1363. "Cluster mode is not enabled on this node"
  1364. )
  1365. cluster_slots = str_if_bytes(r.execute_command("CLUSTER SLOTS"))
  1366. startup_nodes_reachable = True
  1367. except Exception as e:
  1368. # Try the next startup node.
  1369. # The exception is saved and raised only if we have no more nodes.
  1370. exception = e
  1371. continue
  1372. # CLUSTER SLOTS command results in the following output:
  1373. # [[slot_section[from_slot,to_slot,master,replica1,...,replicaN]]]
  1374. # where each node contains the following list: [IP, port, node_id]
  1375. # Therefore, cluster_slots[0][2][0] will be the IP address of the
  1376. # primary node of the first slot section.
  1377. # If there's only one server in the cluster, its ``host`` is ''
  1378. # Fix it to the host in startup_nodes
  1379. if (
  1380. len(cluster_slots) == 1
  1381. and len(cluster_slots[0][2][0]) == 0
  1382. and len(self.startup_nodes) == 1
  1383. ):
  1384. cluster_slots[0][2][0] = startup_node.host
  1385. for slot in cluster_slots:
  1386. primary_node = slot[2]
  1387. host = str_if_bytes(primary_node[0])
  1388. if host == "":
  1389. host = startup_node.host
  1390. port = int(primary_node[1])
  1391. host, port = self.remap_host_port(host, port)
  1392. target_node = self._get_or_create_cluster_node(
  1393. host, port, PRIMARY, tmp_nodes_cache
  1394. )
  1395. # add this node to the nodes cache
  1396. tmp_nodes_cache[target_node.name] = target_node
  1397. for i in range(int(slot[0]), int(slot[1]) + 1):
  1398. if i not in tmp_slots:
  1399. tmp_slots[i] = []
  1400. tmp_slots[i].append(target_node)
  1401. replica_nodes = [slot[j] for j in range(3, len(slot))]
  1402. for replica_node in replica_nodes:
  1403. host = str_if_bytes(replica_node[0])
  1404. port = replica_node[1]
  1405. host, port = self.remap_host_port(host, port)
  1406. target_replica_node = self._get_or_create_cluster_node(
  1407. host, port, REPLICA, tmp_nodes_cache
  1408. )
  1409. tmp_slots[i].append(target_replica_node)
  1410. # add this node to the nodes cache
  1411. tmp_nodes_cache[
  1412. target_replica_node.name
  1413. ] = target_replica_node
  1414. else:
  1415. # Validate that 2 nodes want to use the same slot cache
  1416. # setup
  1417. tmp_slot = tmp_slots[i][0]
  1418. if tmp_slot.name != target_node.name:
  1419. disagreements.append(
  1420. f"{tmp_slot.name} vs {target_node.name} on slot: {i}"
  1421. )
  1422. if len(disagreements) > 5:
  1423. raise RedisClusterException(
  1424. f"startup_nodes could not agree on a valid "
  1425. f'slots cache: {", ".join(disagreements)}'
  1426. )
  1427. fully_covered = self.check_slots_coverage(tmp_slots)
  1428. if fully_covered:
  1429. # Don't need to continue to the next startup node if all
  1430. # slots are covered
  1431. break
  1432. if not startup_nodes_reachable:
  1433. raise RedisClusterException(
  1434. f"Redis Cluster cannot be connected. Please provide at least "
  1435. f"one reachable node: {str(exception)}"
  1436. ) from exception
  1437. # Create Redis connections to all nodes
  1438. self.create_redis_connections(list(tmp_nodes_cache.values()))
  1439. # Check if the slots are not fully covered
  1440. if not fully_covered and self._require_full_coverage:
  1441. # Despite the requirement that the slots be covered, there
  1442. # isn't a full coverage
  1443. raise RedisClusterException(
  1444. f"All slots are not covered after query all startup_nodes. "
  1445. f"{len(tmp_slots)} of {REDIS_CLUSTER_HASH_SLOTS} "
  1446. f"covered..."
  1447. )
  1448. # Set the tmp variables to the real variables
  1449. self.nodes_cache = tmp_nodes_cache
  1450. self.slots_cache = tmp_slots
  1451. # Set the default node
  1452. self.default_node = self.get_nodes_by_server_type(PRIMARY)[0]
  1453. if self._dynamic_startup_nodes:
  1454. # Populate the startup nodes with all discovered nodes
  1455. self.startup_nodes = tmp_nodes_cache
  1456. # If initialize was called after a MovedError, clear it
  1457. self._moved_exception = None
  1458. def close(self):
  1459. self.default_node = None
  1460. for node in self.nodes_cache.values():
  1461. if node.redis_connection:
  1462. node.redis_connection.close()
  1463. def reset(self):
  1464. try:
  1465. self.read_load_balancer.reset()
  1466. except TypeError:
  1467. # The read_load_balancer is None, do nothing
  1468. pass
  1469. def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
  1470. """
  1471. Remap the host and port returned from the cluster to a different
  1472. internal value. Useful if the client is not connecting directly
  1473. to the cluster.
  1474. """
  1475. if self.address_remap:
  1476. return self.address_remap((host, port))
  1477. return host, port
  1478. class ClusterPubSub(PubSub):
  1479. """
  1480. Wrapper for PubSub class.
  1481. IMPORTANT: before using ClusterPubSub, read about the known limitations
  1482. with pubsub in Cluster mode and learn how to workaround them:
  1483. https://redis-py-cluster.readthedocs.io/en/stable/pubsub.html
  1484. """
  1485. def __init__(
  1486. self,
  1487. redis_cluster,
  1488. node=None,
  1489. host=None,
  1490. port=None,
  1491. push_handler_func=None,
  1492. **kwargs,
  1493. ):
  1494. """
  1495. When a pubsub instance is created without specifying a node, a single
  1496. node will be transparently chosen for the pubsub connection on the
  1497. first command execution. The node will be determined by:
  1498. 1. Hashing the channel name in the request to find its keyslot
  1499. 2. Selecting a node that handles the keyslot: If read_from_replicas is
  1500. set to true, a replica can be selected.
  1501. :type redis_cluster: RedisCluster
  1502. :type node: ClusterNode
  1503. :type host: str
  1504. :type port: int
  1505. """
  1506. self.node = None
  1507. self.set_pubsub_node(redis_cluster, node, host, port)
  1508. connection_pool = (
  1509. None
  1510. if self.node is None
  1511. else redis_cluster.get_redis_connection(self.node).connection_pool
  1512. )
  1513. self.cluster = redis_cluster
  1514. self.node_pubsub_mapping = {}
  1515. self._pubsubs_generator = self._pubsubs_generator()
  1516. super().__init__(
  1517. connection_pool=connection_pool,
  1518. encoder=redis_cluster.encoder,
  1519. push_handler_func=push_handler_func,
  1520. **kwargs,
  1521. )
  1522. def set_pubsub_node(self, cluster, node=None, host=None, port=None):
  1523. """
  1524. The pubsub node will be set according to the passed node, host and port
  1525. When none of the node, host, or port are specified - the node is set
  1526. to None and will be determined by the keyslot of the channel in the
  1527. first command to be executed.
  1528. RedisClusterException will be thrown if the passed node does not exist
  1529. in the cluster.
  1530. If host is passed without port, or vice versa, a DataError will be
  1531. thrown.
  1532. :type cluster: RedisCluster
  1533. :type node: ClusterNode
  1534. :type host: str
  1535. :type port: int
  1536. """
  1537. if node is not None:
  1538. # node is passed by the user
  1539. self._raise_on_invalid_node(cluster, node, node.host, node.port)
  1540. pubsub_node = node
  1541. elif host is not None and port is not None:
  1542. # host and port passed by the user
  1543. node = cluster.get_node(host=host, port=port)
  1544. self._raise_on_invalid_node(cluster, node, host, port)
  1545. pubsub_node = node
  1546. elif any([host, port]) is True:
  1547. # only 'host' or 'port' passed
  1548. raise DataError("Passing a host requires passing a port, and vice versa")
  1549. else:
  1550. # nothing passed by the user. set node to None
  1551. pubsub_node = None
  1552. self.node = pubsub_node
  1553. def get_pubsub_node(self):
  1554. """
  1555. Get the node that is being used as the pubsub connection
  1556. """
  1557. return self.node
  1558. def _raise_on_invalid_node(self, redis_cluster, node, host, port):
  1559. """
  1560. Raise a RedisClusterException if the node is None or doesn't exist in
  1561. the cluster.
  1562. """
  1563. if node is None or redis_cluster.get_node(node_name=node.name) is None:
  1564. raise RedisClusterException(
  1565. f"Node {host}:{port} doesn't exist in the cluster"
  1566. )
  1567. def execute_command(self, *args):
  1568. """
  1569. Execute a subscribe/unsubscribe command.
  1570. Taken code from redis-py and tweak to make it work within a cluster.
  1571. """
  1572. # NOTE: don't parse the response in this function -- it could pull a
  1573. # legitimate message off the stack if the connection is already
  1574. # subscribed to one or more channels
  1575. if self.connection is None:
  1576. if self.connection_pool is None:
  1577. if len(args) > 1:
  1578. # Hash the first channel and get one of the nodes holding
  1579. # this slot
  1580. channel = args[1]
  1581. slot = self.cluster.keyslot(channel)
  1582. node = self.cluster.nodes_manager.get_node_from_slot(
  1583. slot, self.cluster.read_from_replicas
  1584. )
  1585. else:
  1586. # Get a random node
  1587. node = self.cluster.get_random_node()
  1588. self.node = node
  1589. redis_connection = self.cluster.get_redis_connection(node)
  1590. self.connection_pool = redis_connection.connection_pool
  1591. self.connection = self.connection_pool.get_connection(
  1592. "pubsub", self.shard_hint
  1593. )
  1594. # register a callback that re-subscribes to any channels we
  1595. # were listening to when we were disconnected
  1596. self.connection._register_connect_callback(self.on_connect)
  1597. if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
  1598. self.connection._parser.set_push_handler(self.push_handler_func)
  1599. connection = self.connection
  1600. self._execute(connection, connection.send_command, *args)
  1601. def _get_node_pubsub(self, node):
  1602. try:
  1603. return self.node_pubsub_mapping[node.name]
  1604. except KeyError:
  1605. pubsub = node.redis_connection.pubsub(
  1606. push_handler_func=self.push_handler_func
  1607. )
  1608. self.node_pubsub_mapping[node.name] = pubsub
  1609. return pubsub
  1610. def _sharded_message_generator(self):
  1611. for _ in range(len(self.node_pubsub_mapping)):
  1612. pubsub = next(self._pubsubs_generator)
  1613. message = pubsub.get_message()
  1614. if message is not None:
  1615. return message
  1616. return None
  1617. def _pubsubs_generator(self):
  1618. while True:
  1619. for pubsub in self.node_pubsub_mapping.values():
  1620. yield pubsub
  1621. def get_sharded_message(
  1622. self, ignore_subscribe_messages=False, timeout=0.0, target_node=None
  1623. ):
  1624. if target_node:
  1625. message = self.node_pubsub_mapping[target_node.name].get_message(
  1626. ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout
  1627. )
  1628. else:
  1629. message = self._sharded_message_generator()
  1630. if message is None:
  1631. return None
  1632. elif str_if_bytes(message["type"]) == "sunsubscribe":
  1633. if message["channel"] in self.pending_unsubscribe_shard_channels:
  1634. self.pending_unsubscribe_shard_channels.remove(message["channel"])
  1635. self.shard_channels.pop(message["channel"], None)
  1636. node = self.cluster.get_node_from_key(message["channel"])
  1637. if self.node_pubsub_mapping[node.name].subscribed is False:
  1638. self.node_pubsub_mapping.pop(node.name)
  1639. if not self.channels and not self.patterns and not self.shard_channels:
  1640. # There are no subscriptions anymore, set subscribed_event flag
  1641. # to false
  1642. self.subscribed_event.clear()
  1643. if self.ignore_subscribe_messages or ignore_subscribe_messages:
  1644. return None
  1645. return message
  1646. def ssubscribe(self, *args, **kwargs):
  1647. if args:
  1648. args = list_or_args(args[0], args[1:])
  1649. s_channels = dict.fromkeys(args)
  1650. s_channels.update(kwargs)
  1651. for s_channel, handler in s_channels.items():
  1652. node = self.cluster.get_node_from_key(s_channel)
  1653. pubsub = self._get_node_pubsub(node)
  1654. if handler:
  1655. pubsub.ssubscribe(**{s_channel: handler})
  1656. else:
  1657. pubsub.ssubscribe(s_channel)
  1658. self.shard_channels.update(pubsub.shard_channels)
  1659. self.pending_unsubscribe_shard_channels.difference_update(
  1660. self._normalize_keys({s_channel: None})
  1661. )
  1662. if pubsub.subscribed and not self.subscribed:
  1663. self.subscribed_event.set()
  1664. self.health_check_response_counter = 0
  1665. def sunsubscribe(self, *args):
  1666. if args:
  1667. args = list_or_args(args[0], args[1:])
  1668. else:
  1669. args = self.shard_channels
  1670. for s_channel in args:
  1671. node = self.cluster.get_node_from_key(s_channel)
  1672. p = self._get_node_pubsub(node)
  1673. p.sunsubscribe(s_channel)
  1674. self.pending_unsubscribe_shard_channels.update(
  1675. p.pending_unsubscribe_shard_channels
  1676. )
  1677. def get_redis_connection(self):
  1678. """
  1679. Get the Redis connection of the pubsub connected node.
  1680. """
  1681. if self.node is not None:
  1682. return self.node.redis_connection
  1683. def disconnect(self):
  1684. """
  1685. Disconnect the pubsub connection.
  1686. """
  1687. if self.connection:
  1688. self.connection.disconnect()
  1689. for pubsub in self.node_pubsub_mapping.values():
  1690. pubsub.connection.disconnect()
  1691. class ClusterPipeline(RedisCluster):
  1692. """
  1693. Support for Redis pipeline
  1694. in cluster mode
  1695. """
  1696. ERRORS_ALLOW_RETRY = (
  1697. ConnectionError,
  1698. TimeoutError,
  1699. MovedError,
  1700. AskError,
  1701. TryAgainError,
  1702. )
  1703. def __init__(
  1704. self,
  1705. nodes_manager: "NodesManager",
  1706. commands_parser: "CommandsParser",
  1707. result_callbacks: Optional[Dict[str, Callable]] = None,
  1708. cluster_response_callbacks: Optional[Dict[str, Callable]] = None,
  1709. startup_nodes: Optional[List["ClusterNode"]] = None,
  1710. read_from_replicas: bool = False,
  1711. cluster_error_retry_attempts: int = 3,
  1712. reinitialize_steps: int = 5,
  1713. lock=None,
  1714. **kwargs,
  1715. ):
  1716. """ """
  1717. self.command_stack = []
  1718. self.nodes_manager = nodes_manager
  1719. self.commands_parser = commands_parser
  1720. self.refresh_table_asap = False
  1721. self.result_callbacks = (
  1722. result_callbacks or self.__class__.RESULT_CALLBACKS.copy()
  1723. )
  1724. self.startup_nodes = startup_nodes if startup_nodes else []
  1725. self.read_from_replicas = read_from_replicas
  1726. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  1727. self.cluster_response_callbacks = cluster_response_callbacks
  1728. self.cluster_error_retry_attempts = cluster_error_retry_attempts
  1729. self.reinitialize_counter = 0
  1730. self.reinitialize_steps = reinitialize_steps
  1731. self.encoder = Encoder(
  1732. kwargs.get("encoding", "utf-8"),
  1733. kwargs.get("encoding_errors", "strict"),
  1734. kwargs.get("decode_responses", False),
  1735. )
  1736. if lock is None:
  1737. lock = threading.Lock()
  1738. self._lock = lock
  1739. def __repr__(self):
  1740. """ """
  1741. return f"{type(self).__name__}"
  1742. def __enter__(self):
  1743. """ """
  1744. return self
  1745. def __exit__(self, exc_type, exc_value, traceback):
  1746. """ """
  1747. self.reset()
  1748. def __del__(self):
  1749. try:
  1750. self.reset()
  1751. except Exception:
  1752. pass
  1753. def __len__(self):
  1754. """ """
  1755. return len(self.command_stack)
  1756. def __bool__(self):
  1757. "Pipeline instances should always evaluate to True on Python 3+"
  1758. return True
  1759. def execute_command(self, *args, **kwargs):
  1760. """
  1761. Wrapper function for pipeline_execute_command
  1762. """
  1763. return self.pipeline_execute_command(*args, **kwargs)
  1764. def pipeline_execute_command(self, *args, **options):
  1765. """
  1766. Appends the executed command to the pipeline's command stack
  1767. """
  1768. self.command_stack.append(
  1769. PipelineCommand(args, options, len(self.command_stack))
  1770. )
  1771. return self
  1772. def raise_first_error(self, stack):
  1773. """
  1774. Raise the first exception on the stack
  1775. """
  1776. for c in stack:
  1777. r = c.result
  1778. if isinstance(r, Exception):
  1779. self.annotate_exception(r, c.position + 1, c.args)
  1780. raise r
  1781. def annotate_exception(self, exception, number, command):
  1782. """
  1783. Provides extra context to the exception prior to it being handled
  1784. """
  1785. cmd = " ".join(map(safe_str, command))
  1786. msg = (
  1787. f"Command # {number} ({cmd}) of pipeline "
  1788. f"caused error: {exception.args[0]}"
  1789. )
  1790. exception.args = (msg,) + exception.args[1:]
  1791. def execute(self, raise_on_error=True):
  1792. """
  1793. Execute all the commands in the current pipeline
  1794. """
  1795. stack = self.command_stack
  1796. try:
  1797. return self.send_cluster_commands(stack, raise_on_error)
  1798. finally:
  1799. self.reset()
  1800. def reset(self):
  1801. """
  1802. Reset back to empty pipeline.
  1803. """
  1804. self.command_stack = []
  1805. self.scripts = set()
  1806. # TODO: Implement
  1807. # make sure to reset the connection state in the event that we were
  1808. # watching something
  1809. # if self.watching and self.connection:
  1810. # try:
  1811. # # call this manually since our unwatch or
  1812. # # immediate_execute_command methods can call reset()
  1813. # self.connection.send_command('UNWATCH')
  1814. # self.connection.read_response()
  1815. # except ConnectionError:
  1816. # # disconnect will also remove any previous WATCHes
  1817. # self.connection.disconnect()
  1818. # clean up the other instance attributes
  1819. self.watching = False
  1820. self.explicit_transaction = False
  1821. # TODO: Implement
  1822. # we can safely return the connection to the pool here since we're
  1823. # sure we're no longer WATCHing anything
  1824. # if self.connection:
  1825. # self.connection_pool.release(self.connection)
  1826. # self.connection = None
  1827. def send_cluster_commands(
  1828. self, stack, raise_on_error=True, allow_redirections=True
  1829. ):
  1830. """
  1831. Wrapper for CLUSTERDOWN error handling.
  1832. If the cluster reports it is down it is assumed that:
  1833. - connection_pool was disconnected
  1834. - connection_pool was reseted
  1835. - refereh_table_asap set to True
  1836. It will try the number of times specified by
  1837. the config option "self.cluster_error_retry_attempts"
  1838. which defaults to 3 unless manually configured.
  1839. If it reaches the number of times, the command will
  1840. raises ClusterDownException.
  1841. """
  1842. if not stack:
  1843. return []
  1844. retry_attempts = self.cluster_error_retry_attempts
  1845. while True:
  1846. try:
  1847. return self._send_cluster_commands(
  1848. stack,
  1849. raise_on_error=raise_on_error,
  1850. allow_redirections=allow_redirections,
  1851. )
  1852. except (ClusterDownError, ConnectionError) as e:
  1853. if retry_attempts > 0:
  1854. # Try again with the new cluster setup. All other errors
  1855. # should be raised.
  1856. retry_attempts -= 1
  1857. pass
  1858. else:
  1859. raise e
  1860. def _send_cluster_commands(
  1861. self, stack, raise_on_error=True, allow_redirections=True
  1862. ):
  1863. """
  1864. Send a bunch of cluster commands to the redis cluster.
  1865. `allow_redirections` If the pipeline should follow
  1866. `ASK` & `MOVED` responses automatically. If set
  1867. to false it will raise RedisClusterException.
  1868. """
  1869. # the first time sending the commands we send all of
  1870. # the commands that were queued up.
  1871. # if we have to run through it again, we only retry
  1872. # the commands that failed.
  1873. attempt = sorted(stack, key=lambda x: x.position)
  1874. is_default_node = False
  1875. # build a list of node objects based on node names we need to
  1876. nodes = {}
  1877. # as we move through each command that still needs to be processed,
  1878. # we figure out the slot number that command maps to, then from
  1879. # the slot determine the node.
  1880. for c in attempt:
  1881. while True:
  1882. # refer to our internal node -> slot table that
  1883. # tells us where a given command should route to.
  1884. # (it might be possible we have a cached node that no longer
  1885. # exists in the cluster, which is why we do this in a loop)
  1886. passed_targets = c.options.pop("target_nodes", None)
  1887. if passed_targets and not self._is_nodes_flag(passed_targets):
  1888. target_nodes = self._parse_target_nodes(passed_targets)
  1889. else:
  1890. target_nodes = self._determine_nodes(
  1891. *c.args, node_flag=passed_targets
  1892. )
  1893. if not target_nodes:
  1894. raise RedisClusterException(
  1895. f"No targets were found to execute {c.args} command on"
  1896. )
  1897. if len(target_nodes) > 1:
  1898. raise RedisClusterException(
  1899. f"Too many targets for command {c.args}"
  1900. )
  1901. node = target_nodes[0]
  1902. if node == self.get_default_node():
  1903. is_default_node = True
  1904. # now that we know the name of the node
  1905. # ( it's just a string in the form of host:port )
  1906. # we can build a list of commands for each node.
  1907. node_name = node.name
  1908. if node_name not in nodes:
  1909. redis_node = self.get_redis_connection(node)
  1910. try:
  1911. connection = get_connection(redis_node, c.args)
  1912. except ConnectionError:
  1913. # Connection retries are being handled in the node's
  1914. # Retry object. Reinitialize the node -> slot table.
  1915. self.nodes_manager.initialize()
  1916. if is_default_node:
  1917. self.replace_default_node()
  1918. raise
  1919. nodes[node_name] = NodeCommands(
  1920. redis_node.parse_response,
  1921. redis_node.connection_pool,
  1922. connection,
  1923. )
  1924. nodes[node_name].append(c)
  1925. break
  1926. # send the commands in sequence.
  1927. # we write to all the open sockets for each node first,
  1928. # before reading anything
  1929. # this allows us to flush all the requests out across the
  1930. # network essentially in parallel
  1931. # so that we can read them all in parallel as they come back.
  1932. # we dont' multiplex on the sockets as they come available,
  1933. # but that shouldn't make too much difference.
  1934. node_commands = nodes.values()
  1935. for n in node_commands:
  1936. n.write()
  1937. for n in node_commands:
  1938. n.read()
  1939. # release all of the redis connections we allocated earlier
  1940. # back into the connection pool.
  1941. # we used to do this step as part of a try/finally block,
  1942. # but it is really dangerous to
  1943. # release connections back into the pool if for some
  1944. # reason the socket has data still left in it
  1945. # from a previous operation. The write and
  1946. # read operations already have try/catch around them for
  1947. # all known types of errors including connection
  1948. # and socket level errors.
  1949. # So if we hit an exception, something really bad
  1950. # happened and putting any oF
  1951. # these connections back into the pool is a very bad idea.
  1952. # the socket might have unread buffer still sitting in it,
  1953. # and then the next time we read from it we pass the
  1954. # buffered result back from a previous command and
  1955. # every single request after to that connection will always get
  1956. # a mismatched result.
  1957. for n in nodes.values():
  1958. n.connection_pool.release(n.connection)
  1959. # if the response isn't an exception it is a
  1960. # valid response from the node
  1961. # we're all done with that command, YAY!
  1962. # if we have more commands to attempt, we've run into problems.
  1963. # collect all the commands we are allowed to retry.
  1964. # (MOVED, ASK, or connection errors or timeout errors)
  1965. attempt = sorted(
  1966. (
  1967. c
  1968. for c in attempt
  1969. if isinstance(c.result, ClusterPipeline.ERRORS_ALLOW_RETRY)
  1970. ),
  1971. key=lambda x: x.position,
  1972. )
  1973. if attempt and allow_redirections:
  1974. # RETRY MAGIC HAPPENS HERE!
  1975. # send these remaing commands one at a time using `execute_command`
  1976. # in the main client. This keeps our retry logic
  1977. # in one place mostly,
  1978. # and allows us to be more confident in correctness of behavior.
  1979. # at this point any speed gains from pipelining have been lost
  1980. # anyway, so we might as well make the best
  1981. # attempt to get the correct behavior.
  1982. #
  1983. # The client command will handle retries for each
  1984. # individual command sequentially as we pass each
  1985. # one into `execute_command`. Any exceptions
  1986. # that bubble out should only appear once all
  1987. # retries have been exhausted.
  1988. #
  1989. # If a lot of commands have failed, we'll be setting the
  1990. # flag to rebuild the slots table from scratch.
  1991. # So MOVED errors should correct themselves fairly quickly.
  1992. self.reinitialize_counter += 1
  1993. if self._should_reinitialized():
  1994. self.nodes_manager.initialize()
  1995. if is_default_node:
  1996. self.replace_default_node()
  1997. for c in attempt:
  1998. try:
  1999. # send each command individually like we
  2000. # do in the main client.
  2001. c.result = super().execute_command(*c.args, **c.options)
  2002. except RedisError as e:
  2003. c.result = e
  2004. # turn the response back into a simple flat array that corresponds
  2005. # to the sequence of commands issued in the stack in pipeline.execute()
  2006. response = []
  2007. for c in sorted(stack, key=lambda x: x.position):
  2008. if c.args[0] in self.cluster_response_callbacks:
  2009. c.result = self.cluster_response_callbacks[c.args[0]](
  2010. c.result, **c.options
  2011. )
  2012. response.append(c.result)
  2013. if raise_on_error:
  2014. self.raise_first_error(stack)
  2015. return response
  2016. def _fail_on_redirect(self, allow_redirections):
  2017. """ """
  2018. if not allow_redirections:
  2019. raise RedisClusterException(
  2020. "ASK & MOVED redirection not allowed in this pipeline"
  2021. )
  2022. def exists(self, *keys):
  2023. return self.execute_command("EXISTS", *keys)
  2024. def eval(self):
  2025. """ """
  2026. raise RedisClusterException("method eval() is not implemented")
  2027. def multi(self):
  2028. """ """
  2029. raise RedisClusterException("method multi() is not implemented")
  2030. def immediate_execute_command(self, *args, **options):
  2031. """ """
  2032. raise RedisClusterException(
  2033. "method immediate_execute_command() is not implemented"
  2034. )
  2035. def _execute_transaction(self, *args, **kwargs):
  2036. """ """
  2037. raise RedisClusterException("method _execute_transaction() is not implemented")
  2038. def load_scripts(self):
  2039. """ """
  2040. raise RedisClusterException("method load_scripts() is not implemented")
  2041. def watch(self, *names):
  2042. """ """
  2043. raise RedisClusterException("method watch() is not implemented")
  2044. def unwatch(self):
  2045. """ """
  2046. raise RedisClusterException("method unwatch() is not implemented")
  2047. def script_load_for_pipeline(self, *args, **kwargs):
  2048. """ """
  2049. raise RedisClusterException(
  2050. "method script_load_for_pipeline() is not implemented"
  2051. )
  2052. def delete(self, *names):
  2053. """
  2054. "Delete a key specified by ``names``"
  2055. """
  2056. if len(names) != 1:
  2057. raise RedisClusterException(
  2058. "deleting multiple keys is not implemented in pipeline command"
  2059. )
  2060. return self.execute_command("DEL", names[0])
  2061. def unlink(self, *names):
  2062. """
  2063. "Unlink a key specified by ``names``"
  2064. """
  2065. if len(names) != 1:
  2066. raise RedisClusterException(
  2067. "unlinking multiple keys is not implemented in pipeline command"
  2068. )
  2069. return self.execute_command("UNLINK", names[0])
  2070. def block_pipeline_command(name: str) -> Callable[..., Any]:
  2071. """
  2072. Prints error because some pipelined commands should
  2073. be blocked when running in cluster-mode
  2074. """
  2075. def inner(*args, **kwargs):
  2076. raise RedisClusterException(
  2077. f"ERROR: Calling pipelined function {name} is blocked "
  2078. f"when running redis in cluster mode..."
  2079. )
  2080. return inner
  2081. # Blocked pipeline commands
  2082. PIPELINE_BLOCKED_COMMANDS = (
  2083. "BGREWRITEAOF",
  2084. "BGSAVE",
  2085. "BITOP",
  2086. "BRPOPLPUSH",
  2087. "CLIENT GETNAME",
  2088. "CLIENT KILL",
  2089. "CLIENT LIST",
  2090. "CLIENT SETNAME",
  2091. "CLIENT",
  2092. "CONFIG GET",
  2093. "CONFIG RESETSTAT",
  2094. "CONFIG REWRITE",
  2095. "CONFIG SET",
  2096. "CONFIG",
  2097. "DBSIZE",
  2098. "ECHO",
  2099. "EVALSHA",
  2100. "FLUSHALL",
  2101. "FLUSHDB",
  2102. "INFO",
  2103. "KEYS",
  2104. "LASTSAVE",
  2105. "MGET",
  2106. "MGET NONATOMIC",
  2107. "MOVE",
  2108. "MSET",
  2109. "MSET NONATOMIC",
  2110. "MSETNX",
  2111. "PFCOUNT",
  2112. "PFMERGE",
  2113. "PING",
  2114. "PUBLISH",
  2115. "RANDOMKEY",
  2116. "READONLY",
  2117. "READWRITE",
  2118. "RENAME",
  2119. "RENAMENX",
  2120. "RPOPLPUSH",
  2121. "SAVE",
  2122. "SCAN",
  2123. "SCRIPT EXISTS",
  2124. "SCRIPT FLUSH",
  2125. "SCRIPT KILL",
  2126. "SCRIPT LOAD",
  2127. "SCRIPT",
  2128. "SDIFF",
  2129. "SDIFFSTORE",
  2130. "SENTINEL GET MASTER ADDR BY NAME",
  2131. "SENTINEL MASTER",
  2132. "SENTINEL MASTERS",
  2133. "SENTINEL MONITOR",
  2134. "SENTINEL REMOVE",
  2135. "SENTINEL SENTINELS",
  2136. "SENTINEL SET",
  2137. "SENTINEL SLAVES",
  2138. "SENTINEL",
  2139. "SHUTDOWN",
  2140. "SINTER",
  2141. "SINTERSTORE",
  2142. "SLAVEOF",
  2143. "SLOWLOG GET",
  2144. "SLOWLOG LEN",
  2145. "SLOWLOG RESET",
  2146. "SLOWLOG",
  2147. "SMOVE",
  2148. "SORT",
  2149. "SUNION",
  2150. "SUNIONSTORE",
  2151. "TIME",
  2152. )
  2153. for command in PIPELINE_BLOCKED_COMMANDS:
  2154. command = command.replace(" ", "_").lower()
  2155. setattr(ClusterPipeline, command, block_pipeline_command(command))
  2156. class PipelineCommand:
  2157. """ """
  2158. def __init__(self, args, options=None, position=None):
  2159. self.args = args
  2160. if options is None:
  2161. options = {}
  2162. self.options = options
  2163. self.position = position
  2164. self.result = None
  2165. self.node = None
  2166. self.asking = False
  2167. class NodeCommands:
  2168. """ """
  2169. def __init__(self, parse_response, connection_pool, connection):
  2170. """ """
  2171. self.parse_response = parse_response
  2172. self.connection_pool = connection_pool
  2173. self.connection = connection
  2174. self.commands = []
  2175. def append(self, c):
  2176. """ """
  2177. self.commands.append(c)
  2178. def write(self):
  2179. """
  2180. Code borrowed from Redis so it can be fixed
  2181. """
  2182. connection = self.connection
  2183. commands = self.commands
  2184. # We are going to clobber the commands with the write, so go ahead
  2185. # and ensure that nothing is sitting there from a previous run.
  2186. for c in commands:
  2187. c.result = None
  2188. # build up all commands into a single request to increase network perf
  2189. # send all the commands and catch connection and timeout errors.
  2190. try:
  2191. connection.send_packed_command(
  2192. connection.pack_commands([c.args for c in commands])
  2193. )
  2194. except (ConnectionError, TimeoutError) as e:
  2195. for c in commands:
  2196. c.result = e
  2197. def read(self):
  2198. """ """
  2199. connection = self.connection
  2200. for c in self.commands:
  2201. # if there is a result on this command,
  2202. # it means we ran into an exception
  2203. # like a connection error. Trying to parse
  2204. # a response on a connection that
  2205. # is no longer open will result in a
  2206. # connection error raised by redis-py.
  2207. # but redis-py doesn't check in parse_response
  2208. # that the sock object is
  2209. # still set and if you try to
  2210. # read from a closed connection, it will
  2211. # result in an AttributeError because
  2212. # it will do a readline() call on None.
  2213. # This can have all kinds of nasty side-effects.
  2214. # Treating this case as a connection error
  2215. # is fine because it will dump
  2216. # the connection object back into the
  2217. # pool and on the next write, it will
  2218. # explicitly open the connection and all will be well.
  2219. if c.result is None:
  2220. try:
  2221. c.result = self.parse_response(connection, c.args[0], **c.options)
  2222. except (ConnectionError, TimeoutError) as e:
  2223. for c in self.commands:
  2224. c.result = e
  2225. return
  2226. except RedisError:
  2227. c.result = sys.exc_info()[1]