cluster.py 121 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352
  1. import random
  2. import socket
  3. import sys
  4. import threading
  5. import time
  6. from abc import ABC, abstractmethod
  7. from collections import OrderedDict
  8. from copy import copy
  9. from enum import Enum
  10. from itertools import chain
  11. from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
  12. from redis._parsers import CommandsParser, Encoder
  13. from redis._parsers.helpers import parse_scan
  14. from redis.backoff import ExponentialWithJitterBackoff, NoBackoff
  15. from redis.cache import CacheConfig, CacheFactory, CacheFactoryInterface, CacheInterface
  16. from redis.client import EMPTY_RESPONSE, CaseInsensitiveDict, PubSub, Redis
  17. from redis.commands import READ_COMMANDS, RedisClusterCommands
  18. from redis.commands.helpers import list_or_args
  19. from redis.connection import (
  20. Connection,
  21. ConnectionPool,
  22. parse_url,
  23. )
  24. from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
  25. from redis.event import (
  26. AfterPooledConnectionsInstantiationEvent,
  27. AfterPubSubConnectionInstantiationEvent,
  28. ClientType,
  29. EventDispatcher,
  30. )
  31. from redis.exceptions import (
  32. AskError,
  33. AuthenticationError,
  34. ClusterDownError,
  35. ClusterError,
  36. ConnectionError,
  37. CrossSlotTransactionError,
  38. DataError,
  39. ExecAbortError,
  40. InvalidPipelineStack,
  41. MovedError,
  42. RedisClusterException,
  43. RedisError,
  44. ResponseError,
  45. SlotNotCoveredError,
  46. TimeoutError,
  47. TryAgainError,
  48. WatchError,
  49. )
  50. from redis.lock import Lock
  51. from redis.retry import Retry
  52. from redis.utils import (
  53. deprecated_args,
  54. dict_merge,
  55. list_keys_to_dict,
  56. merge_result,
  57. safe_str,
  58. str_if_bytes,
  59. truncate_text,
  60. )
  61. def get_node_name(host: str, port: Union[str, int]) -> str:
  62. return f"{host}:{port}"
  63. @deprecated_args(
  64. allowed_args=["redis_node"],
  65. reason="Use get_connection(redis_node) instead",
  66. version="5.3.0",
  67. )
  68. def get_connection(redis_node: Redis, *args, **options) -> Connection:
  69. return redis_node.connection or redis_node.connection_pool.get_connection()
  70. def parse_scan_result(command, res, **options):
  71. cursors = {}
  72. ret = []
  73. for node_name, response in res.items():
  74. cursor, r = parse_scan(response, **options)
  75. cursors[node_name] = cursor
  76. ret += r
  77. return cursors, ret
  78. def parse_pubsub_numsub(command, res, **options):
  79. numsub_d = OrderedDict()
  80. for numsub_tups in res.values():
  81. for channel, numsubbed in numsub_tups:
  82. try:
  83. numsub_d[channel] += numsubbed
  84. except KeyError:
  85. numsub_d[channel] = numsubbed
  86. ret_numsub = [(channel, numsub) for channel, numsub in numsub_d.items()]
  87. return ret_numsub
  88. def parse_cluster_slots(
  89. resp: Any, **options: Any
  90. ) -> Dict[Tuple[int, int], Dict[str, Any]]:
  91. current_host = options.get("current_host", "")
  92. def fix_server(*args: Any) -> Tuple[str, Any]:
  93. return str_if_bytes(args[0]) or current_host, args[1]
  94. slots = {}
  95. for slot in resp:
  96. start, end, primary = slot[:3]
  97. replicas = slot[3:]
  98. slots[start, end] = {
  99. "primary": fix_server(*primary),
  100. "replicas": [fix_server(*replica) for replica in replicas],
  101. }
  102. return slots
  103. def parse_cluster_shards(resp, **options):
  104. """
  105. Parse CLUSTER SHARDS response.
  106. """
  107. if isinstance(resp[0], dict):
  108. return resp
  109. shards = []
  110. for x in resp:
  111. shard = {"slots": [], "nodes": []}
  112. for i in range(0, len(x[1]), 2):
  113. shard["slots"].append((x[1][i], (x[1][i + 1])))
  114. nodes = x[3]
  115. for node in nodes:
  116. dict_node = {}
  117. for i in range(0, len(node), 2):
  118. dict_node[node[i]] = node[i + 1]
  119. shard["nodes"].append(dict_node)
  120. shards.append(shard)
  121. return shards
  122. def parse_cluster_myshardid(resp, **options):
  123. """
  124. Parse CLUSTER MYSHARDID response.
  125. """
  126. return resp.decode("utf-8")
  127. PRIMARY = "primary"
  128. REPLICA = "replica"
  129. SLOT_ID = "slot-id"
  130. REDIS_ALLOWED_KEYS = (
  131. "connection_class",
  132. "connection_pool",
  133. "connection_pool_class",
  134. "client_name",
  135. "credential_provider",
  136. "db",
  137. "decode_responses",
  138. "encoding",
  139. "encoding_errors",
  140. "host",
  141. "lib_name",
  142. "lib_version",
  143. "max_connections",
  144. "nodes_flag",
  145. "redis_connect_func",
  146. "password",
  147. "port",
  148. "queue_class",
  149. "retry",
  150. "retry_on_timeout",
  151. "protocol",
  152. "socket_connect_timeout",
  153. "socket_keepalive",
  154. "socket_keepalive_options",
  155. "socket_timeout",
  156. "ssl",
  157. "ssl_ca_certs",
  158. "ssl_ca_data",
  159. "ssl_certfile",
  160. "ssl_cert_reqs",
  161. "ssl_keyfile",
  162. "ssl_password",
  163. "ssl_check_hostname",
  164. "unix_socket_path",
  165. "username",
  166. "cache",
  167. "cache_config",
  168. )
  169. KWARGS_DISABLED_KEYS = ("host", "port", "retry")
  170. def cleanup_kwargs(**kwargs):
  171. """
  172. Remove unsupported or disabled keys from kwargs
  173. """
  174. connection_kwargs = {
  175. k: v
  176. for k, v in kwargs.items()
  177. if k in REDIS_ALLOWED_KEYS and k not in KWARGS_DISABLED_KEYS
  178. }
  179. return connection_kwargs
  180. class AbstractRedisCluster:
  181. RedisClusterRequestTTL = 16
  182. PRIMARIES = "primaries"
  183. REPLICAS = "replicas"
  184. ALL_NODES = "all"
  185. RANDOM = "random"
  186. DEFAULT_NODE = "default-node"
  187. NODE_FLAGS = {PRIMARIES, REPLICAS, ALL_NODES, RANDOM, DEFAULT_NODE}
  188. COMMAND_FLAGS = dict_merge(
  189. list_keys_to_dict(
  190. [
  191. "ACL CAT",
  192. "ACL DELUSER",
  193. "ACL DRYRUN",
  194. "ACL GENPASS",
  195. "ACL GETUSER",
  196. "ACL HELP",
  197. "ACL LIST",
  198. "ACL LOG",
  199. "ACL LOAD",
  200. "ACL SAVE",
  201. "ACL SETUSER",
  202. "ACL USERS",
  203. "ACL WHOAMI",
  204. "AUTH",
  205. "CLIENT LIST",
  206. "CLIENT SETINFO",
  207. "CLIENT SETNAME",
  208. "CLIENT GETNAME",
  209. "CONFIG SET",
  210. "CONFIG REWRITE",
  211. "CONFIG RESETSTAT",
  212. "TIME",
  213. "PUBSUB CHANNELS",
  214. "PUBSUB NUMPAT",
  215. "PUBSUB NUMSUB",
  216. "PUBSUB SHARDCHANNELS",
  217. "PUBSUB SHARDNUMSUB",
  218. "PING",
  219. "INFO",
  220. "SHUTDOWN",
  221. "KEYS",
  222. "DBSIZE",
  223. "BGSAVE",
  224. "SLOWLOG GET",
  225. "SLOWLOG LEN",
  226. "SLOWLOG RESET",
  227. "WAIT",
  228. "WAITAOF",
  229. "SAVE",
  230. "MEMORY PURGE",
  231. "MEMORY MALLOC-STATS",
  232. "MEMORY STATS",
  233. "LASTSAVE",
  234. "CLIENT TRACKINGINFO",
  235. "CLIENT PAUSE",
  236. "CLIENT UNPAUSE",
  237. "CLIENT UNBLOCK",
  238. "CLIENT ID",
  239. "CLIENT REPLY",
  240. "CLIENT GETREDIR",
  241. "CLIENT INFO",
  242. "CLIENT KILL",
  243. "READONLY",
  244. "CLUSTER INFO",
  245. "CLUSTER MEET",
  246. "CLUSTER MYSHARDID",
  247. "CLUSTER NODES",
  248. "CLUSTER REPLICAS",
  249. "CLUSTER RESET",
  250. "CLUSTER SET-CONFIG-EPOCH",
  251. "CLUSTER SLOTS",
  252. "CLUSTER SHARDS",
  253. "CLUSTER COUNT-FAILURE-REPORTS",
  254. "CLUSTER KEYSLOT",
  255. "COMMAND",
  256. "COMMAND COUNT",
  257. "COMMAND LIST",
  258. "COMMAND GETKEYS",
  259. "CONFIG GET",
  260. "DEBUG",
  261. "RANDOMKEY",
  262. "READONLY",
  263. "READWRITE",
  264. "TIME",
  265. "TFUNCTION LOAD",
  266. "TFUNCTION DELETE",
  267. "TFUNCTION LIST",
  268. "TFCALL",
  269. "TFCALLASYNC",
  270. "LATENCY HISTORY",
  271. "LATENCY LATEST",
  272. "LATENCY RESET",
  273. "MODULE LIST",
  274. "MODULE LOAD",
  275. "MODULE UNLOAD",
  276. "MODULE LOADEX",
  277. ],
  278. DEFAULT_NODE,
  279. ),
  280. list_keys_to_dict(
  281. [
  282. "FLUSHALL",
  283. "FLUSHDB",
  284. "FUNCTION DELETE",
  285. "FUNCTION FLUSH",
  286. "FUNCTION LIST",
  287. "FUNCTION LOAD",
  288. "FUNCTION RESTORE",
  289. "SCAN",
  290. "SCRIPT EXISTS",
  291. "SCRIPT FLUSH",
  292. "SCRIPT LOAD",
  293. ],
  294. PRIMARIES,
  295. ),
  296. list_keys_to_dict(["FUNCTION DUMP"], RANDOM),
  297. list_keys_to_dict(
  298. [
  299. "CLUSTER COUNTKEYSINSLOT",
  300. "CLUSTER DELSLOTS",
  301. "CLUSTER DELSLOTSRANGE",
  302. "CLUSTER GETKEYSINSLOT",
  303. "CLUSTER SETSLOT",
  304. ],
  305. SLOT_ID,
  306. ),
  307. )
  308. SEARCH_COMMANDS = (
  309. [
  310. "FT.CREATE",
  311. "FT.SEARCH",
  312. "FT.AGGREGATE",
  313. "FT.EXPLAIN",
  314. "FT.EXPLAINCLI",
  315. "FT,PROFILE",
  316. "FT.ALTER",
  317. "FT.DROPINDEX",
  318. "FT.ALIASADD",
  319. "FT.ALIASUPDATE",
  320. "FT.ALIASDEL",
  321. "FT.TAGVALS",
  322. "FT.SUGADD",
  323. "FT.SUGGET",
  324. "FT.SUGDEL",
  325. "FT.SUGLEN",
  326. "FT.SYNUPDATE",
  327. "FT.SYNDUMP",
  328. "FT.SPELLCHECK",
  329. "FT.DICTADD",
  330. "FT.DICTDEL",
  331. "FT.DICTDUMP",
  332. "FT.INFO",
  333. "FT._LIST",
  334. "FT.CONFIG",
  335. "FT.ADD",
  336. "FT.DEL",
  337. "FT.DROP",
  338. "FT.GET",
  339. "FT.MGET",
  340. "FT.SYNADD",
  341. ],
  342. )
  343. CLUSTER_COMMANDS_RESPONSE_CALLBACKS = {
  344. "CLUSTER SLOTS": parse_cluster_slots,
  345. "CLUSTER SHARDS": parse_cluster_shards,
  346. "CLUSTER MYSHARDID": parse_cluster_myshardid,
  347. }
  348. RESULT_CALLBACKS = dict_merge(
  349. list_keys_to_dict(["PUBSUB NUMSUB", "PUBSUB SHARDNUMSUB"], parse_pubsub_numsub),
  350. list_keys_to_dict(
  351. ["PUBSUB NUMPAT"], lambda command, res: sum(list(res.values()))
  352. ),
  353. list_keys_to_dict(
  354. ["KEYS", "PUBSUB CHANNELS", "PUBSUB SHARDCHANNELS"], merge_result
  355. ),
  356. list_keys_to_dict(
  357. [
  358. "PING",
  359. "CONFIG SET",
  360. "CONFIG REWRITE",
  361. "CONFIG RESETSTAT",
  362. "CLIENT SETNAME",
  363. "BGSAVE",
  364. "SLOWLOG RESET",
  365. "SAVE",
  366. "MEMORY PURGE",
  367. "CLIENT PAUSE",
  368. "CLIENT UNPAUSE",
  369. ],
  370. lambda command, res: all(res.values()) if isinstance(res, dict) else res,
  371. ),
  372. list_keys_to_dict(
  373. ["DBSIZE", "WAIT"],
  374. lambda command, res: sum(res.values()) if isinstance(res, dict) else res,
  375. ),
  376. list_keys_to_dict(
  377. ["CLIENT UNBLOCK"], lambda command, res: 1 if sum(res.values()) > 0 else 0
  378. ),
  379. list_keys_to_dict(["SCAN"], parse_scan_result),
  380. list_keys_to_dict(
  381. ["SCRIPT LOAD"], lambda command, res: list(res.values()).pop()
  382. ),
  383. list_keys_to_dict(
  384. ["SCRIPT EXISTS"], lambda command, res: [all(k) for k in zip(*res.values())]
  385. ),
  386. list_keys_to_dict(["SCRIPT FLUSH"], lambda command, res: all(res.values())),
  387. )
  388. ERRORS_ALLOW_RETRY = (
  389. ConnectionError,
  390. TimeoutError,
  391. ClusterDownError,
  392. SlotNotCoveredError,
  393. )
  394. def replace_default_node(self, target_node: "ClusterNode" = None) -> None:
  395. """Replace the default cluster node.
  396. A random cluster node will be chosen if target_node isn't passed, and primaries
  397. will be prioritized. The default node will not be changed if there are no other
  398. nodes in the cluster.
  399. Args:
  400. target_node (ClusterNode, optional): Target node to replace the default
  401. node. Defaults to None.
  402. """
  403. if target_node:
  404. self.nodes_manager.default_node = target_node
  405. else:
  406. curr_node = self.get_default_node()
  407. primaries = [node for node in self.get_primaries() if node != curr_node]
  408. if primaries:
  409. # Choose a primary if the cluster contains different primaries
  410. self.nodes_manager.default_node = random.choice(primaries)
  411. else:
  412. # Otherwise, choose a primary if the cluster contains different primaries
  413. replicas = [node for node in self.get_replicas() if node != curr_node]
  414. if replicas:
  415. self.nodes_manager.default_node = random.choice(replicas)
  416. class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
  417. @classmethod
  418. def from_url(cls, url, **kwargs):
  419. """
  420. Return a Redis client object configured from the given URL
  421. For example::
  422. redis://[[username]:[password]]@localhost:6379/0
  423. rediss://[[username]:[password]]@localhost:6379/0
  424. unix://[username@]/path/to/socket.sock?db=0[&password=password]
  425. Three URL schemes are supported:
  426. - `redis://` creates a TCP socket connection. See more at:
  427. <https://www.iana.org/assignments/uri-schemes/prov/redis>
  428. - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
  429. <https://www.iana.org/assignments/uri-schemes/prov/rediss>
  430. - ``unix://``: creates a Unix Domain Socket connection.
  431. The username, password, hostname, path and all querystring values
  432. are passed through urllib.parse.unquote in order to replace any
  433. percent-encoded values with their corresponding characters.
  434. There are several ways to specify a database number. The first value
  435. found will be used:
  436. 1. A ``db`` querystring option, e.g. redis://localhost?db=0
  437. 2. If using the redis:// or rediss:// schemes, the path argument
  438. of the url, e.g. redis://localhost/0
  439. 3. A ``db`` keyword argument to this function.
  440. If none of these options are specified, the default db=0 is used.
  441. All querystring options are cast to their appropriate Python types.
  442. Boolean arguments can be specified with string values "True"/"False"
  443. or "Yes"/"No". Values that cannot be properly cast cause a
  444. ``ValueError`` to be raised. Once parsed, the querystring arguments
  445. and keyword arguments are passed to the ``ConnectionPool``'s
  446. class initializer. In the case of conflicting arguments, querystring
  447. arguments always win.
  448. """
  449. return cls(url=url, **kwargs)
  450. @deprecated_args(
  451. args_to_warn=["read_from_replicas"],
  452. reason="Please configure the 'load_balancing_strategy' instead",
  453. version="5.3.0",
  454. )
  455. @deprecated_args(
  456. args_to_warn=[
  457. "cluster_error_retry_attempts",
  458. ],
  459. reason="Please configure the 'retry' object instead",
  460. version="6.0.0",
  461. )
  462. def __init__(
  463. self,
  464. host: Optional[str] = None,
  465. port: int = 6379,
  466. startup_nodes: Optional[List["ClusterNode"]] = None,
  467. cluster_error_retry_attempts: int = 3,
  468. retry: Optional["Retry"] = None,
  469. require_full_coverage: bool = True,
  470. reinitialize_steps: int = 5,
  471. read_from_replicas: bool = False,
  472. load_balancing_strategy: Optional["LoadBalancingStrategy"] = None,
  473. dynamic_startup_nodes: bool = True,
  474. url: Optional[str] = None,
  475. address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
  476. cache: Optional[CacheInterface] = None,
  477. cache_config: Optional[CacheConfig] = None,
  478. event_dispatcher: Optional[EventDispatcher] = None,
  479. **kwargs,
  480. ):
  481. """
  482. Initialize a new RedisCluster client.
  483. :param startup_nodes:
  484. List of nodes from which initial bootstrapping can be done
  485. :param host:
  486. Can be used to point to a startup node
  487. :param port:
  488. Can be used to point to a startup node
  489. :param require_full_coverage:
  490. When set to False (default value): the client will not require a
  491. full coverage of the slots. However, if not all slots are covered,
  492. and at least one node has 'cluster-require-full-coverage' set to
  493. 'yes,' the server will throw a ClusterDownError for some key-based
  494. commands. See -
  495. https://redis.io/topics/cluster-tutorial#redis-cluster-configuration-parameters
  496. When set to True: all slots must be covered to construct the
  497. cluster client. If not all slots are covered, RedisClusterException
  498. will be thrown.
  499. :param read_from_replicas:
  500. @deprecated - please use load_balancing_strategy instead
  501. Enable read from replicas in READONLY mode. You can read possibly
  502. stale data.
  503. When set to true, read commands will be assigned between the
  504. primary and its replications in a Round-Robin manner.
  505. :param load_balancing_strategy:
  506. Enable read from replicas in READONLY mode and defines the load balancing
  507. strategy that will be used for cluster node selection.
  508. The data read from replicas is eventually consistent with the data in primary nodes.
  509. :param dynamic_startup_nodes:
  510. Set the RedisCluster's startup nodes to all of the discovered nodes.
  511. If true (default value), the cluster's discovered nodes will be used to
  512. determine the cluster nodes-slots mapping in the next topology refresh.
  513. It will remove the initial passed startup nodes if their endpoints aren't
  514. listed in the CLUSTER SLOTS output.
  515. If you use dynamic DNS endpoints for startup nodes but CLUSTER SLOTS lists
  516. specific IP addresses, it is best to set it to false.
  517. :param cluster_error_retry_attempts:
  518. @deprecated - Please configure the 'retry' object instead
  519. In case 'retry' object is set - this argument is ignored!
  520. Number of times to retry before raising an error when
  521. :class:`~.TimeoutError` or :class:`~.ConnectionError`, :class:`~.SlotNotCoveredError` or
  522. :class:`~.ClusterDownError` are encountered
  523. :param retry:
  524. A retry object that defines the retry strategy and the number of
  525. retries for the cluster client.
  526. In current implementation for the cluster client (starting form redis-py version 6.0.0)
  527. the retry object is not yet fully utilized, instead it is used just to determine
  528. the number of retries for the cluster client.
  529. In the future releases the retry object will be used to handle the cluster client retries!
  530. :param reinitialize_steps:
  531. Specifies the number of MOVED errors that need to occur before
  532. reinitializing the whole cluster topology. If a MOVED error occurs
  533. and the cluster does not need to be reinitialized on this current
  534. error handling, only the MOVED slot will be patched with the
  535. redirected node.
  536. To reinitialize the cluster on every MOVED error, set
  537. reinitialize_steps to 1.
  538. To avoid reinitializing the cluster on moved errors, set
  539. reinitialize_steps to 0.
  540. :param address_remap:
  541. An optional callable which, when provided with an internal network
  542. address of a node, e.g. a `(host, port)` tuple, will return the address
  543. where the node is reachable. This can be used to map the addresses at
  544. which the nodes _think_ they are, to addresses at which a client may
  545. reach them, such as when they sit behind a proxy.
  546. :**kwargs:
  547. Extra arguments that will be sent into Redis instance when created
  548. (See Official redis-py doc for supported kwargs - the only limitation
  549. is that you can't provide 'retry' object as part of kwargs.
  550. [https://github.com/andymccurdy/redis-py/blob/master/redis/client.py])
  551. Some kwargs are not supported and will raise a
  552. RedisClusterException:
  553. - db (Redis do not support database SELECT in cluster mode)
  554. """
  555. if startup_nodes is None:
  556. startup_nodes = []
  557. if "db" in kwargs:
  558. # Argument 'db' is not possible to use in cluster mode
  559. raise RedisClusterException(
  560. "Argument 'db' is not possible to use in cluster mode"
  561. )
  562. if "retry" in kwargs:
  563. # Argument 'retry' is not possible to be used in kwargs when in cluster mode
  564. # the kwargs are set to the lower level connections to the cluster nodes
  565. # and there we provide retry configuration without retries allowed.
  566. # The retries should be handled on cluster client level.
  567. raise RedisClusterException(
  568. "The 'retry' argument cannot be used in kwargs when running in cluster mode."
  569. )
  570. # Get the startup node/s
  571. from_url = False
  572. if url is not None:
  573. from_url = True
  574. url_options = parse_url(url)
  575. if "path" in url_options:
  576. raise RedisClusterException(
  577. "RedisCluster does not currently support Unix Domain "
  578. "Socket connections"
  579. )
  580. if "db" in url_options and url_options["db"] != 0:
  581. # Argument 'db' is not possible to use in cluster mode
  582. raise RedisClusterException(
  583. "A ``db`` querystring option can only be 0 in cluster mode"
  584. )
  585. kwargs.update(url_options)
  586. host = kwargs.get("host")
  587. port = kwargs.get("port", port)
  588. startup_nodes.append(ClusterNode(host, port))
  589. elif host is not None and port is not None:
  590. startup_nodes.append(ClusterNode(host, port))
  591. elif len(startup_nodes) == 0:
  592. # No startup node was provided
  593. raise RedisClusterException(
  594. "RedisCluster requires at least one node to discover the "
  595. "cluster. Please provide one of the followings:\n"
  596. "1. host and port, for example:\n"
  597. " RedisCluster(host='localhost', port=6379)\n"
  598. "2. list of startup nodes, for example:\n"
  599. " RedisCluster(startup_nodes=[ClusterNode('localhost', 6379),"
  600. " ClusterNode('localhost', 6378)])"
  601. )
  602. # Update the connection arguments
  603. # Whenever a new connection is established, RedisCluster's on_connect
  604. # method should be run
  605. # If the user passed on_connect function we'll save it and run it
  606. # inside the RedisCluster.on_connect() function
  607. self.user_on_connect_func = kwargs.pop("redis_connect_func", None)
  608. kwargs.update({"redis_connect_func": self.on_connect})
  609. kwargs = cleanup_kwargs(**kwargs)
  610. if retry:
  611. self.retry = retry
  612. else:
  613. self.retry = Retry(
  614. backoff=ExponentialWithJitterBackoff(base=1, cap=10),
  615. retries=cluster_error_retry_attempts,
  616. )
  617. self.encoder = Encoder(
  618. kwargs.get("encoding", "utf-8"),
  619. kwargs.get("encoding_errors", "strict"),
  620. kwargs.get("decode_responses", False),
  621. )
  622. protocol = kwargs.get("protocol", None)
  623. if (cache_config or cache) and protocol not in [3, "3"]:
  624. raise RedisError("Client caching is only supported with RESP version 3")
  625. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  626. self.node_flags = self.__class__.NODE_FLAGS.copy()
  627. self.read_from_replicas = read_from_replicas
  628. self.load_balancing_strategy = load_balancing_strategy
  629. self.reinitialize_counter = 0
  630. self.reinitialize_steps = reinitialize_steps
  631. if event_dispatcher is None:
  632. self._event_dispatcher = EventDispatcher()
  633. else:
  634. self._event_dispatcher = event_dispatcher
  635. self.nodes_manager = NodesManager(
  636. startup_nodes=startup_nodes,
  637. from_url=from_url,
  638. require_full_coverage=require_full_coverage,
  639. dynamic_startup_nodes=dynamic_startup_nodes,
  640. address_remap=address_remap,
  641. cache=cache,
  642. cache_config=cache_config,
  643. event_dispatcher=self._event_dispatcher,
  644. **kwargs,
  645. )
  646. self.cluster_response_callbacks = CaseInsensitiveDict(
  647. self.__class__.CLUSTER_COMMANDS_RESPONSE_CALLBACKS
  648. )
  649. self.result_callbacks = CaseInsensitiveDict(self.__class__.RESULT_CALLBACKS)
  650. self.commands_parser = CommandsParser(self)
  651. self._lock = threading.RLock()
  652. def __enter__(self):
  653. return self
  654. def __exit__(self, exc_type, exc_value, traceback):
  655. self.close()
  656. def __del__(self):
  657. try:
  658. self.close()
  659. except Exception:
  660. pass
  661. def disconnect_connection_pools(self):
  662. for node in self.get_nodes():
  663. if node.redis_connection:
  664. try:
  665. node.redis_connection.connection_pool.disconnect()
  666. except OSError:
  667. # Client was already disconnected. do nothing
  668. pass
  669. def on_connect(self, connection):
  670. """
  671. Initialize the connection, authenticate and select a database and send
  672. READONLY if it is set during object initialization.
  673. """
  674. connection.on_connect()
  675. if self.read_from_replicas or self.load_balancing_strategy:
  676. # Sending READONLY command to server to configure connection as
  677. # readonly. Since each cluster node may change its server type due
  678. # to a failover, we should establish a READONLY connection
  679. # regardless of the server type. If this is a primary connection,
  680. # READONLY would not affect executing write commands.
  681. connection.send_command("READONLY")
  682. if str_if_bytes(connection.read_response()) != "OK":
  683. raise ConnectionError("READONLY command failed")
  684. if self.user_on_connect_func is not None:
  685. self.user_on_connect_func(connection)
  686. def get_redis_connection(self, node: "ClusterNode") -> Redis:
  687. if not node.redis_connection:
  688. with self._lock:
  689. if not node.redis_connection:
  690. self.nodes_manager.create_redis_connections([node])
  691. return node.redis_connection
  692. def get_node(self, host=None, port=None, node_name=None):
  693. return self.nodes_manager.get_node(host, port, node_name)
  694. def get_primaries(self):
  695. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  696. def get_replicas(self):
  697. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  698. def get_random_node(self):
  699. return random.choice(list(self.nodes_manager.nodes_cache.values()))
  700. def get_nodes(self):
  701. return list(self.nodes_manager.nodes_cache.values())
  702. def get_node_from_key(self, key, replica=False):
  703. """
  704. Get the node that holds the key's slot.
  705. If replica set to True but the slot doesn't have any replicas, None is
  706. returned.
  707. """
  708. slot = self.keyslot(key)
  709. slot_cache = self.nodes_manager.slots_cache.get(slot)
  710. if slot_cache is None or len(slot_cache) == 0:
  711. raise SlotNotCoveredError(f'Slot "{slot}" is not covered by the cluster.')
  712. if replica and len(self.nodes_manager.slots_cache[slot]) < 2:
  713. return None
  714. elif replica:
  715. node_idx = 1
  716. else:
  717. # primary
  718. node_idx = 0
  719. return slot_cache[node_idx]
  720. def get_default_node(self):
  721. """
  722. Get the cluster's default node
  723. """
  724. return self.nodes_manager.default_node
  725. def set_default_node(self, node):
  726. """
  727. Set the default node of the cluster.
  728. :param node: 'ClusterNode'
  729. :return True if the default node was set, else False
  730. """
  731. if node is None or self.get_node(node_name=node.name) is None:
  732. return False
  733. self.nodes_manager.default_node = node
  734. return True
  735. def set_retry(self, retry: Retry) -> None:
  736. self.retry = retry
  737. def monitor(self, target_node=None):
  738. """
  739. Returns a Monitor object for the specified target node.
  740. The default cluster node will be selected if no target node was
  741. specified.
  742. Monitor is useful for handling the MONITOR command to the redis server.
  743. next_command() method returns one command from monitor
  744. listen() method yields commands from monitor.
  745. """
  746. if target_node is None:
  747. target_node = self.get_default_node()
  748. if target_node.redis_connection is None:
  749. raise RedisClusterException(
  750. f"Cluster Node {target_node.name} has no redis_connection"
  751. )
  752. return target_node.redis_connection.monitor()
  753. def pubsub(self, node=None, host=None, port=None, **kwargs):
  754. """
  755. Allows passing a ClusterNode, or host&port, to get a pubsub instance
  756. connected to the specified node
  757. """
  758. return ClusterPubSub(self, node=node, host=host, port=port, **kwargs)
  759. def pipeline(self, transaction=None, shard_hint=None):
  760. """
  761. Cluster impl:
  762. Pipelines do not work in cluster mode the same way they
  763. do in normal mode. Create a clone of this object so
  764. that simulating pipelines will work correctly. Each
  765. command will be called directly when used and
  766. when calling execute() will only return the result stack.
  767. """
  768. if shard_hint:
  769. raise RedisClusterException("shard_hint is deprecated in cluster mode")
  770. return ClusterPipeline(
  771. nodes_manager=self.nodes_manager,
  772. commands_parser=self.commands_parser,
  773. startup_nodes=self.nodes_manager.startup_nodes,
  774. result_callbacks=self.result_callbacks,
  775. cluster_response_callbacks=self.cluster_response_callbacks,
  776. cluster_error_retry_attempts=self.retry.get_retries(),
  777. read_from_replicas=self.read_from_replicas,
  778. load_balancing_strategy=self.load_balancing_strategy,
  779. reinitialize_steps=self.reinitialize_steps,
  780. retry=self.retry,
  781. lock=self._lock,
  782. transaction=transaction,
  783. )
  784. def lock(
  785. self,
  786. name,
  787. timeout=None,
  788. sleep=0.1,
  789. blocking=True,
  790. blocking_timeout=None,
  791. lock_class=None,
  792. thread_local=True,
  793. raise_on_release_error: bool = True,
  794. ):
  795. """
  796. Return a new Lock object using key ``name`` that mimics
  797. the behavior of threading.Lock.
  798. If specified, ``timeout`` indicates a maximum life for the lock.
  799. By default, it will remain locked until release() is called.
  800. ``sleep`` indicates the amount of time to sleep per loop iteration
  801. when the lock is in blocking mode and another client is currently
  802. holding the lock.
  803. ``blocking`` indicates whether calling ``acquire`` should block until
  804. the lock has been acquired or to fail immediately, causing ``acquire``
  805. to return False and the lock not being acquired. Defaults to True.
  806. Note this value can be overridden by passing a ``blocking``
  807. argument to ``acquire``.
  808. ``blocking_timeout`` indicates the maximum amount of time in seconds to
  809. spend trying to acquire the lock. A value of ``None`` indicates
  810. continue trying forever. ``blocking_timeout`` can be specified as a
  811. float or integer, both representing the number of seconds to wait.
  812. ``lock_class`` forces the specified lock implementation. Note that as
  813. of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
  814. a Lua-based lock). So, it's unlikely you'll need this parameter, unless
  815. you have created your own custom lock class.
  816. ``thread_local`` indicates whether the lock token is placed in
  817. thread-local storage. By default, the token is placed in thread local
  818. storage so that a thread only sees its token, not a token set by
  819. another thread. Consider the following timeline:
  820. time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
  821. thread-1 sets the token to "abc"
  822. time: 1, thread-2 blocks trying to acquire `my-lock` using the
  823. Lock instance.
  824. time: 5, thread-1 has not yet completed. redis expires the lock
  825. key.
  826. time: 5, thread-2 acquired `my-lock` now that it's available.
  827. thread-2 sets the token to "xyz"
  828. time: 6, thread-1 finishes its work and calls release(). if the
  829. token is *not* stored in thread local storage, then
  830. thread-1 would see the token value as "xyz" and would be
  831. able to successfully release the thread-2's lock.
  832. ``raise_on_release_error`` indicates whether to raise an exception when
  833. the lock is no longer owned when exiting the context manager. By default,
  834. this is True, meaning an exception will be raised. If False, the warning
  835. will be logged and the exception will be suppressed.
  836. In some use cases it's necessary to disable thread local storage. For
  837. example, if you have code where one thread acquires a lock and passes
  838. that lock instance to a worker thread to release later. If thread
  839. local storage isn't disabled in this case, the worker thread won't see
  840. the token set by the thread that acquired the lock. Our assumption
  841. is that these cases aren't common and as such default to using
  842. thread local storage."""
  843. if lock_class is None:
  844. lock_class = Lock
  845. return lock_class(
  846. self,
  847. name,
  848. timeout=timeout,
  849. sleep=sleep,
  850. blocking=blocking,
  851. blocking_timeout=blocking_timeout,
  852. thread_local=thread_local,
  853. raise_on_release_error=raise_on_release_error,
  854. )
  855. def set_response_callback(self, command, callback):
  856. """Set a custom Response Callback"""
  857. self.cluster_response_callbacks[command] = callback
  858. def _determine_nodes(self, *args, **kwargs) -> List["ClusterNode"]:
  859. # Determine which nodes should be executed the command on.
  860. # Returns a list of target nodes.
  861. command = args[0].upper()
  862. if len(args) >= 2 and f"{args[0]} {args[1]}".upper() in self.command_flags:
  863. command = f"{args[0]} {args[1]}".upper()
  864. nodes_flag = kwargs.pop("nodes_flag", None)
  865. if nodes_flag is not None:
  866. # nodes flag passed by the user
  867. command_flag = nodes_flag
  868. else:
  869. # get the nodes group for this command if it was predefined
  870. command_flag = self.command_flags.get(command)
  871. if command_flag == self.__class__.RANDOM:
  872. # return a random node
  873. return [self.get_random_node()]
  874. elif command_flag == self.__class__.PRIMARIES:
  875. # return all primaries
  876. return self.get_primaries()
  877. elif command_flag == self.__class__.REPLICAS:
  878. # return all replicas
  879. return self.get_replicas()
  880. elif command_flag == self.__class__.ALL_NODES:
  881. # return all nodes
  882. return self.get_nodes()
  883. elif command_flag == self.__class__.DEFAULT_NODE:
  884. # return the cluster's default node
  885. return [self.nodes_manager.default_node]
  886. elif command in self.__class__.SEARCH_COMMANDS[0]:
  887. return [self.nodes_manager.default_node]
  888. else:
  889. # get the node that holds the key's slot
  890. slot = self.determine_slot(*args)
  891. node = self.nodes_manager.get_node_from_slot(
  892. slot,
  893. self.read_from_replicas and command in READ_COMMANDS,
  894. self.load_balancing_strategy if command in READ_COMMANDS else None,
  895. )
  896. return [node]
  897. def _should_reinitialized(self):
  898. # To reinitialize the cluster on every MOVED error,
  899. # set reinitialize_steps to 1.
  900. # To avoid reinitializing the cluster on moved errors, set
  901. # reinitialize_steps to 0.
  902. if self.reinitialize_steps == 0:
  903. return False
  904. else:
  905. return self.reinitialize_counter % self.reinitialize_steps == 0
  906. def keyslot(self, key):
  907. """
  908. Calculate keyslot for a given key.
  909. See Keys distribution model in https://redis.io/topics/cluster-spec
  910. """
  911. k = self.encoder.encode(key)
  912. return key_slot(k)
  913. def _get_command_keys(self, *args):
  914. """
  915. Get the keys in the command. If the command has no keys in in, None is
  916. returned.
  917. NOTE: Due to a bug in redis<7.0, this function does not work properly
  918. for EVAL or EVALSHA when the `numkeys` arg is 0.
  919. - issue: https://github.com/redis/redis/issues/9493
  920. - fix: https://github.com/redis/redis/pull/9733
  921. So, don't use this function with EVAL or EVALSHA.
  922. """
  923. redis_conn = self.get_default_node().redis_connection
  924. return self.commands_parser.get_keys(redis_conn, *args)
  925. def determine_slot(self, *args) -> int:
  926. """
  927. Figure out what slot to use based on args.
  928. Raises a RedisClusterException if there's a missing key and we can't
  929. determine what slots to map the command to; or, if the keys don't
  930. all map to the same key slot.
  931. """
  932. command = args[0]
  933. if self.command_flags.get(command) == SLOT_ID:
  934. # The command contains the slot ID
  935. return args[1]
  936. # Get the keys in the command
  937. # EVAL and EVALSHA are common enough that it's wasteful to go to the
  938. # redis server to parse the keys. Besides, there is a bug in redis<7.0
  939. # where `self._get_command_keys()` fails anyway. So, we special case
  940. # EVAL/EVALSHA.
  941. if command.upper() in ("EVAL", "EVALSHA"):
  942. # command syntax: EVAL "script body" num_keys ...
  943. if len(args) <= 2:
  944. raise RedisClusterException(f"Invalid args in command: {args}")
  945. num_actual_keys = int(args[2])
  946. eval_keys = args[3 : 3 + num_actual_keys]
  947. # if there are 0 keys, that means the script can be run on any node
  948. # so we can just return a random slot
  949. if len(eval_keys) == 0:
  950. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  951. keys = eval_keys
  952. else:
  953. keys = self._get_command_keys(*args)
  954. if keys is None or len(keys) == 0:
  955. # FCALL can call a function with 0 keys, that means the function
  956. # can be run on any node so we can just return a random slot
  957. if command.upper() in ("FCALL", "FCALL_RO"):
  958. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  959. raise RedisClusterException(
  960. "No way to dispatch this command to Redis Cluster. "
  961. "Missing key.\nYou can execute the command by specifying "
  962. f"target nodes.\nCommand: {args}"
  963. )
  964. # single key command
  965. if len(keys) == 1:
  966. return self.keyslot(keys[0])
  967. # multi-key command; we need to make sure all keys are mapped to
  968. # the same slot
  969. slots = {self.keyslot(key) for key in keys}
  970. if len(slots) != 1:
  971. raise RedisClusterException(
  972. f"{command} - all keys must map to the same key slot"
  973. )
  974. return slots.pop()
  975. def get_encoder(self):
  976. """
  977. Get the connections' encoder
  978. """
  979. return self.encoder
  980. def get_connection_kwargs(self):
  981. """
  982. Get the connections' key-word arguments
  983. """
  984. return self.nodes_manager.connection_kwargs
  985. def _is_nodes_flag(self, target_nodes):
  986. return isinstance(target_nodes, str) and target_nodes in self.node_flags
  987. def _parse_target_nodes(self, target_nodes):
  988. if isinstance(target_nodes, list):
  989. nodes = target_nodes
  990. elif isinstance(target_nodes, ClusterNode):
  991. # Supports passing a single ClusterNode as a variable
  992. nodes = [target_nodes]
  993. elif isinstance(target_nodes, dict):
  994. # Supports dictionaries of the format {node_name: node}.
  995. # It enables to execute commands with multi nodes as follows:
  996. # rc.cluster_save_config(rc.get_primaries())
  997. nodes = target_nodes.values()
  998. else:
  999. raise TypeError(
  1000. "target_nodes type can be one of the following: "
  1001. "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
  1002. "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
  1003. f"The passed type is {type(target_nodes)}"
  1004. )
  1005. return nodes
  1006. def execute_command(self, *args, **kwargs):
  1007. return self._internal_execute_command(*args, **kwargs)
  1008. def _internal_execute_command(self, *args, **kwargs):
  1009. """
  1010. Wrapper for ERRORS_ALLOW_RETRY error handling.
  1011. It will try the number of times specified by the retries property from
  1012. config option "self.retry" which defaults to 3 unless manually
  1013. configured.
  1014. If it reaches the number of times, the command will raise the exception
  1015. Key argument :target_nodes: can be passed with the following types:
  1016. nodes_flag: PRIMARIES, REPLICAS, ALL_NODES, RANDOM
  1017. ClusterNode
  1018. list<ClusterNode>
  1019. dict<Any, ClusterNode>
  1020. """
  1021. target_nodes_specified = False
  1022. is_default_node = False
  1023. target_nodes = None
  1024. passed_targets = kwargs.pop("target_nodes", None)
  1025. if passed_targets is not None and not self._is_nodes_flag(passed_targets):
  1026. target_nodes = self._parse_target_nodes(passed_targets)
  1027. target_nodes_specified = True
  1028. # If an error that allows retrying was thrown, the nodes and slots
  1029. # cache were reinitialized. We will retry executing the command with
  1030. # the updated cluster setup only when the target nodes can be
  1031. # determined again with the new cache tables. Therefore, when target
  1032. # nodes were passed to this function, we cannot retry the command
  1033. # execution since the nodes may not be valid anymore after the tables
  1034. # were reinitialized. So in case of passed target nodes,
  1035. # retry_attempts will be set to 0.
  1036. retry_attempts = 0 if target_nodes_specified else self.retry.get_retries()
  1037. # Add one for the first execution
  1038. execute_attempts = 1 + retry_attempts
  1039. for _ in range(execute_attempts):
  1040. try:
  1041. res = {}
  1042. if not target_nodes_specified:
  1043. # Determine the nodes to execute the command on
  1044. target_nodes = self._determine_nodes(
  1045. *args, **kwargs, nodes_flag=passed_targets
  1046. )
  1047. if not target_nodes:
  1048. raise RedisClusterException(
  1049. f"No targets were found to execute {args} command on"
  1050. )
  1051. if (
  1052. len(target_nodes) == 1
  1053. and target_nodes[0] == self.get_default_node()
  1054. ):
  1055. is_default_node = True
  1056. for node in target_nodes:
  1057. res[node.name] = self._execute_command(node, *args, **kwargs)
  1058. # Return the processed result
  1059. return self._process_result(args[0], res, **kwargs)
  1060. except Exception as e:
  1061. if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  1062. if is_default_node:
  1063. # Replace the default cluster node
  1064. self.replace_default_node()
  1065. # The nodes and slots cache were reinitialized.
  1066. # Try again with the new cluster setup.
  1067. retry_attempts -= 1
  1068. continue
  1069. else:
  1070. # raise the exception
  1071. raise e
  1072. def _execute_command(self, target_node, *args, **kwargs):
  1073. """
  1074. Send a command to a node in the cluster
  1075. """
  1076. command = args[0]
  1077. redis_node = None
  1078. connection = None
  1079. redirect_addr = None
  1080. asking = False
  1081. moved = False
  1082. ttl = int(self.RedisClusterRequestTTL)
  1083. while ttl > 0:
  1084. ttl -= 1
  1085. try:
  1086. if asking:
  1087. target_node = self.get_node(node_name=redirect_addr)
  1088. elif moved:
  1089. # MOVED occurred and the slots cache was updated,
  1090. # refresh the target node
  1091. slot = self.determine_slot(*args)
  1092. target_node = self.nodes_manager.get_node_from_slot(
  1093. slot,
  1094. self.read_from_replicas and command in READ_COMMANDS,
  1095. self.load_balancing_strategy
  1096. if command in READ_COMMANDS
  1097. else None,
  1098. )
  1099. moved = False
  1100. redis_node = self.get_redis_connection(target_node)
  1101. connection = get_connection(redis_node)
  1102. if asking:
  1103. connection.send_command("ASKING")
  1104. redis_node.parse_response(connection, "ASKING", **kwargs)
  1105. asking = False
  1106. connection.send_command(*args, **kwargs)
  1107. response = redis_node.parse_response(connection, command, **kwargs)
  1108. # Remove keys entry, it needs only for cache.
  1109. kwargs.pop("keys", None)
  1110. if command in self.cluster_response_callbacks:
  1111. response = self.cluster_response_callbacks[command](
  1112. response, **kwargs
  1113. )
  1114. return response
  1115. except AuthenticationError:
  1116. raise
  1117. except (ConnectionError, TimeoutError) as e:
  1118. # ConnectionError can also be raised if we couldn't get a
  1119. # connection from the pool before timing out, so check that
  1120. # this is an actual connection before attempting to disconnect.
  1121. if connection is not None:
  1122. connection.disconnect()
  1123. # Remove the failed node from the startup nodes before we try
  1124. # to reinitialize the cluster
  1125. self.nodes_manager.startup_nodes.pop(target_node.name, None)
  1126. # Reset the cluster node's connection
  1127. target_node.redis_connection = None
  1128. self.nodes_manager.initialize()
  1129. raise e
  1130. except MovedError as e:
  1131. # First, we will try to patch the slots/nodes cache with the
  1132. # redirected node output and try again. If MovedError exceeds
  1133. # 'reinitialize_steps' number of times, we will force
  1134. # reinitializing the tables, and then try again.
  1135. # 'reinitialize_steps' counter will increase faster when
  1136. # the same client object is shared between multiple threads. To
  1137. # reduce the frequency you can set this variable in the
  1138. # RedisCluster constructor.
  1139. self.reinitialize_counter += 1
  1140. if self._should_reinitialized():
  1141. self.nodes_manager.initialize()
  1142. # Reset the counter
  1143. self.reinitialize_counter = 0
  1144. else:
  1145. self.nodes_manager.update_moved_exception(e)
  1146. moved = True
  1147. except TryAgainError:
  1148. if ttl < self.RedisClusterRequestTTL / 2:
  1149. time.sleep(0.05)
  1150. except AskError as e:
  1151. redirect_addr = get_node_name(host=e.host, port=e.port)
  1152. asking = True
  1153. except (ClusterDownError, SlotNotCoveredError):
  1154. # ClusterDownError can occur during a failover and to get
  1155. # self-healed, we will try to reinitialize the cluster layout
  1156. # and retry executing the command
  1157. # SlotNotCoveredError can occur when the cluster is not fully
  1158. # initialized or can be temporary issue.
  1159. # We will try to reinitialize the cluster topology
  1160. # and retry executing the command
  1161. time.sleep(0.25)
  1162. self.nodes_manager.initialize()
  1163. raise
  1164. except ResponseError:
  1165. raise
  1166. except Exception as e:
  1167. if connection:
  1168. connection.disconnect()
  1169. raise e
  1170. finally:
  1171. if connection is not None:
  1172. redis_node.connection_pool.release(connection)
  1173. raise ClusterError("TTL exhausted.")
  1174. def close(self) -> None:
  1175. try:
  1176. with self._lock:
  1177. if self.nodes_manager:
  1178. self.nodes_manager.close()
  1179. except AttributeError:
  1180. # RedisCluster's __init__ can fail before nodes_manager is set
  1181. pass
  1182. def _process_result(self, command, res, **kwargs):
  1183. """
  1184. Process the result of the executed command.
  1185. The function would return a dict or a single value.
  1186. :type command: str
  1187. :type res: dict
  1188. `res` should be in the following format:
  1189. Dict<node_name, command_result>
  1190. """
  1191. if command in self.result_callbacks:
  1192. return self.result_callbacks[command](command, res, **kwargs)
  1193. elif len(res) == 1:
  1194. # When we execute the command on a single node, we can
  1195. # remove the dictionary and return a single response
  1196. return list(res.values())[0]
  1197. else:
  1198. return res
  1199. def load_external_module(self, funcname, func):
  1200. """
  1201. This function can be used to add externally defined redis modules,
  1202. and their namespaces to the redis client.
  1203. ``funcname`` - A string containing the name of the function to create
  1204. ``func`` - The function, being added to this class.
  1205. """
  1206. setattr(self, funcname, func)
  1207. def transaction(self, func, *watches, **kwargs):
  1208. """
  1209. Convenience method for executing the callable `func` as a transaction
  1210. while watching all keys specified in `watches`. The 'func' callable
  1211. should expect a single argument which is a Pipeline object.
  1212. """
  1213. shard_hint = kwargs.pop("shard_hint", None)
  1214. value_from_callable = kwargs.pop("value_from_callable", False)
  1215. watch_delay = kwargs.pop("watch_delay", None)
  1216. with self.pipeline(True, shard_hint) as pipe:
  1217. while True:
  1218. try:
  1219. if watches:
  1220. pipe.watch(*watches)
  1221. func_value = func(pipe)
  1222. exec_value = pipe.execute()
  1223. return func_value if value_from_callable else exec_value
  1224. except WatchError:
  1225. if watch_delay is not None and watch_delay > 0:
  1226. time.sleep(watch_delay)
  1227. continue
  1228. class ClusterNode:
  1229. def __init__(self, host, port, server_type=None, redis_connection=None):
  1230. if host == "localhost":
  1231. host = socket.gethostbyname(host)
  1232. self.host = host
  1233. self.port = port
  1234. self.name = get_node_name(host, port)
  1235. self.server_type = server_type
  1236. self.redis_connection = redis_connection
  1237. def __repr__(self):
  1238. return (
  1239. f"[host={self.host},"
  1240. f"port={self.port},"
  1241. f"name={self.name},"
  1242. f"server_type={self.server_type},"
  1243. f"redis_connection={self.redis_connection}]"
  1244. )
  1245. def __eq__(self, obj):
  1246. return isinstance(obj, ClusterNode) and obj.name == self.name
  1247. def __del__(self):
  1248. try:
  1249. if self.redis_connection is not None:
  1250. self.redis_connection.close()
  1251. except Exception:
  1252. # Ignore errors when closing the connection
  1253. pass
  1254. class LoadBalancingStrategy(Enum):
  1255. ROUND_ROBIN = "round_robin"
  1256. ROUND_ROBIN_REPLICAS = "round_robin_replicas"
  1257. RANDOM_REPLICA = "random_replica"
  1258. class LoadBalancer:
  1259. """
  1260. Round-Robin Load Balancing
  1261. """
  1262. def __init__(self, start_index: int = 0) -> None:
  1263. self.primary_to_idx = {}
  1264. self.start_index = start_index
  1265. def get_server_index(
  1266. self,
  1267. primary: str,
  1268. list_size: int,
  1269. load_balancing_strategy: LoadBalancingStrategy = LoadBalancingStrategy.ROUND_ROBIN,
  1270. ) -> int:
  1271. if load_balancing_strategy == LoadBalancingStrategy.RANDOM_REPLICA:
  1272. return self._get_random_replica_index(list_size)
  1273. else:
  1274. return self._get_round_robin_index(
  1275. primary,
  1276. list_size,
  1277. load_balancing_strategy == LoadBalancingStrategy.ROUND_ROBIN_REPLICAS,
  1278. )
  1279. def reset(self) -> None:
  1280. self.primary_to_idx.clear()
  1281. def _get_random_replica_index(self, list_size: int) -> int:
  1282. return random.randint(1, list_size - 1)
  1283. def _get_round_robin_index(
  1284. self, primary: str, list_size: int, replicas_only: bool
  1285. ) -> int:
  1286. server_index = self.primary_to_idx.setdefault(primary, self.start_index)
  1287. if replicas_only and server_index == 0:
  1288. # skip the primary node index
  1289. server_index = 1
  1290. # Update the index for the next round
  1291. self.primary_to_idx[primary] = (server_index + 1) % list_size
  1292. return server_index
  1293. class NodesManager:
  1294. def __init__(
  1295. self,
  1296. startup_nodes,
  1297. from_url=False,
  1298. require_full_coverage=False,
  1299. lock=None,
  1300. dynamic_startup_nodes=True,
  1301. connection_pool_class=ConnectionPool,
  1302. address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
  1303. cache: Optional[CacheInterface] = None,
  1304. cache_config: Optional[CacheConfig] = None,
  1305. cache_factory: Optional[CacheFactoryInterface] = None,
  1306. event_dispatcher: Optional[EventDispatcher] = None,
  1307. **kwargs,
  1308. ):
  1309. self.nodes_cache: Dict[str, Redis] = {}
  1310. self.slots_cache = {}
  1311. self.startup_nodes = {}
  1312. self.default_node = None
  1313. self.populate_startup_nodes(startup_nodes)
  1314. self.from_url = from_url
  1315. self._require_full_coverage = require_full_coverage
  1316. self._dynamic_startup_nodes = dynamic_startup_nodes
  1317. self.connection_pool_class = connection_pool_class
  1318. self.address_remap = address_remap
  1319. self._cache = cache
  1320. self._cache_config = cache_config
  1321. self._cache_factory = cache_factory
  1322. self._moved_exception = None
  1323. self.connection_kwargs = kwargs
  1324. self.read_load_balancer = LoadBalancer()
  1325. if lock is None:
  1326. lock = threading.RLock()
  1327. self._lock = lock
  1328. if event_dispatcher is None:
  1329. self._event_dispatcher = EventDispatcher()
  1330. else:
  1331. self._event_dispatcher = event_dispatcher
  1332. self._credential_provider = self.connection_kwargs.get(
  1333. "credential_provider", None
  1334. )
  1335. self.initialize()
  1336. def get_node(self, host=None, port=None, node_name=None):
  1337. """
  1338. Get the requested node from the cluster's nodes.
  1339. nodes.
  1340. :return: ClusterNode if the node exists, else None
  1341. """
  1342. if host and port:
  1343. # the user passed host and port
  1344. if host == "localhost":
  1345. host = socket.gethostbyname(host)
  1346. return self.nodes_cache.get(get_node_name(host=host, port=port))
  1347. elif node_name:
  1348. return self.nodes_cache.get(node_name)
  1349. else:
  1350. return None
  1351. def update_moved_exception(self, exception):
  1352. self._moved_exception = exception
  1353. def _update_moved_slots(self):
  1354. """
  1355. Update the slot's node with the redirected one
  1356. """
  1357. e = self._moved_exception
  1358. redirected_node = self.get_node(host=e.host, port=e.port)
  1359. if redirected_node is not None:
  1360. # The node already exists
  1361. if redirected_node.server_type is not PRIMARY:
  1362. # Update the node's server type
  1363. redirected_node.server_type = PRIMARY
  1364. else:
  1365. # This is a new node, we will add it to the nodes cache
  1366. redirected_node = ClusterNode(e.host, e.port, PRIMARY)
  1367. self.nodes_cache[redirected_node.name] = redirected_node
  1368. if redirected_node in self.slots_cache[e.slot_id]:
  1369. # The MOVED error resulted from a failover, and the new slot owner
  1370. # had previously been a replica.
  1371. old_primary = self.slots_cache[e.slot_id][0]
  1372. # Update the old primary to be a replica and add it to the end of
  1373. # the slot's node list
  1374. old_primary.server_type = REPLICA
  1375. self.slots_cache[e.slot_id].append(old_primary)
  1376. # Remove the old replica, which is now a primary, from the slot's
  1377. # node list
  1378. self.slots_cache[e.slot_id].remove(redirected_node)
  1379. # Override the old primary with the new one
  1380. self.slots_cache[e.slot_id][0] = redirected_node
  1381. if self.default_node == old_primary:
  1382. # Update the default node with the new primary
  1383. self.default_node = redirected_node
  1384. else:
  1385. # The new slot owner is a new server, or a server from a different
  1386. # shard. We need to remove all current nodes from the slot's list
  1387. # (including replications) and add just the new node.
  1388. self.slots_cache[e.slot_id] = [redirected_node]
  1389. # Reset moved_exception
  1390. self._moved_exception = None
  1391. @deprecated_args(
  1392. args_to_warn=["server_type"],
  1393. reason=(
  1394. "In case you need select some load balancing strategy "
  1395. "that will use replicas, please set it through 'load_balancing_strategy'"
  1396. ),
  1397. version="5.3.0",
  1398. )
  1399. def get_node_from_slot(
  1400. self,
  1401. slot,
  1402. read_from_replicas=False,
  1403. load_balancing_strategy=None,
  1404. server_type=None,
  1405. ) -> ClusterNode:
  1406. """
  1407. Gets a node that servers this hash slot
  1408. """
  1409. if self._moved_exception:
  1410. with self._lock:
  1411. if self._moved_exception:
  1412. self._update_moved_slots()
  1413. if self.slots_cache.get(slot) is None or len(self.slots_cache[slot]) == 0:
  1414. raise SlotNotCoveredError(
  1415. f'Slot "{slot}" not covered by the cluster. '
  1416. f'"require_full_coverage={self._require_full_coverage}"'
  1417. )
  1418. if read_from_replicas is True and load_balancing_strategy is None:
  1419. load_balancing_strategy = LoadBalancingStrategy.ROUND_ROBIN
  1420. if len(self.slots_cache[slot]) > 1 and load_balancing_strategy:
  1421. # get the server index using the strategy defined in load_balancing_strategy
  1422. primary_name = self.slots_cache[slot][0].name
  1423. node_idx = self.read_load_balancer.get_server_index(
  1424. primary_name, len(self.slots_cache[slot]), load_balancing_strategy
  1425. )
  1426. elif (
  1427. server_type is None
  1428. or server_type == PRIMARY
  1429. or len(self.slots_cache[slot]) == 1
  1430. ):
  1431. # return a primary
  1432. node_idx = 0
  1433. else:
  1434. # return a replica
  1435. # randomly choose one of the replicas
  1436. node_idx = random.randint(1, len(self.slots_cache[slot]) - 1)
  1437. return self.slots_cache[slot][node_idx]
  1438. def get_nodes_by_server_type(self, server_type):
  1439. """
  1440. Get all nodes with the specified server type
  1441. :param server_type: 'primary' or 'replica'
  1442. :return: list of ClusterNode
  1443. """
  1444. return [
  1445. node
  1446. for node in self.nodes_cache.values()
  1447. if node.server_type == server_type
  1448. ]
  1449. def populate_startup_nodes(self, nodes):
  1450. """
  1451. Populate all startup nodes and filters out any duplicates
  1452. """
  1453. for n in nodes:
  1454. self.startup_nodes[n.name] = n
  1455. def check_slots_coverage(self, slots_cache):
  1456. # Validate if all slots are covered or if we should try next
  1457. # startup node
  1458. for i in range(0, REDIS_CLUSTER_HASH_SLOTS):
  1459. if i not in slots_cache:
  1460. return False
  1461. return True
  1462. def create_redis_connections(self, nodes):
  1463. """
  1464. This function will create a redis connection to all nodes in :nodes:
  1465. """
  1466. connection_pools = []
  1467. for node in nodes:
  1468. if node.redis_connection is None:
  1469. node.redis_connection = self.create_redis_node(
  1470. host=node.host, port=node.port, **self.connection_kwargs
  1471. )
  1472. connection_pools.append(node.redis_connection.connection_pool)
  1473. self._event_dispatcher.dispatch(
  1474. AfterPooledConnectionsInstantiationEvent(
  1475. connection_pools, ClientType.SYNC, self._credential_provider
  1476. )
  1477. )
  1478. def create_redis_node(self, host, port, **kwargs):
  1479. # We are configuring the connection pool not to retry
  1480. # connections on lower level clients to avoid retrying
  1481. # connections to nodes that are not reachable
  1482. # and to avoid blocking the connection pool.
  1483. # The only error that will have some handling in the lower
  1484. # level clients is ConnectionError which will trigger disconnection
  1485. # of the socket.
  1486. # The retries will be handled on cluster client level
  1487. # where we will have proper handling of the cluster topology
  1488. node_retry_config = Retry(
  1489. backoff=NoBackoff(), retries=0, supported_errors=(ConnectionError,)
  1490. )
  1491. if self.from_url:
  1492. # Create a redis node with a costumed connection pool
  1493. kwargs.update({"host": host})
  1494. kwargs.update({"port": port})
  1495. kwargs.update({"cache": self._cache})
  1496. kwargs.update({"retry": node_retry_config})
  1497. r = Redis(connection_pool=self.connection_pool_class(**kwargs))
  1498. else:
  1499. r = Redis(
  1500. host=host,
  1501. port=port,
  1502. cache=self._cache,
  1503. retry=node_retry_config,
  1504. **kwargs,
  1505. )
  1506. return r
  1507. def _get_or_create_cluster_node(self, host, port, role, tmp_nodes_cache):
  1508. node_name = get_node_name(host, port)
  1509. # check if we already have this node in the tmp_nodes_cache
  1510. target_node = tmp_nodes_cache.get(node_name)
  1511. if target_node is None:
  1512. # before creating a new cluster node, check if the cluster node already
  1513. # exists in the current nodes cache and has a valid connection so we can
  1514. # reuse it
  1515. target_node = self.nodes_cache.get(node_name)
  1516. if target_node is None or target_node.redis_connection is None:
  1517. # create new cluster node for this cluster
  1518. target_node = ClusterNode(host, port, role)
  1519. if target_node.server_type != role:
  1520. target_node.server_type = role
  1521. # add this node to the nodes cache
  1522. tmp_nodes_cache[target_node.name] = target_node
  1523. return target_node
  1524. def initialize(self):
  1525. """
  1526. Initializes the nodes cache, slots cache and redis connections.
  1527. :startup_nodes:
  1528. Responsible for discovering other nodes in the cluster
  1529. """
  1530. self.reset()
  1531. tmp_nodes_cache = {}
  1532. tmp_slots = {}
  1533. disagreements = []
  1534. startup_nodes_reachable = False
  1535. fully_covered = False
  1536. kwargs = self.connection_kwargs
  1537. exception = None
  1538. # Convert to tuple to prevent RuntimeError if self.startup_nodes
  1539. # is modified during iteration
  1540. for startup_node in tuple(self.startup_nodes.values()):
  1541. try:
  1542. if startup_node.redis_connection:
  1543. r = startup_node.redis_connection
  1544. else:
  1545. # Create a new Redis connection
  1546. r = self.create_redis_node(
  1547. startup_node.host, startup_node.port, **kwargs
  1548. )
  1549. self.startup_nodes[startup_node.name].redis_connection = r
  1550. # Make sure cluster mode is enabled on this node
  1551. try:
  1552. cluster_slots = str_if_bytes(r.execute_command("CLUSTER SLOTS"))
  1553. r.connection_pool.disconnect()
  1554. except ResponseError:
  1555. raise RedisClusterException(
  1556. "Cluster mode is not enabled on this node"
  1557. )
  1558. startup_nodes_reachable = True
  1559. except Exception as e:
  1560. # Try the next startup node.
  1561. # The exception is saved and raised only if we have no more nodes.
  1562. exception = e
  1563. continue
  1564. # CLUSTER SLOTS command results in the following output:
  1565. # [[slot_section[from_slot,to_slot,master,replica1,...,replicaN]]]
  1566. # where each node contains the following list: [IP, port, node_id]
  1567. # Therefore, cluster_slots[0][2][0] will be the IP address of the
  1568. # primary node of the first slot section.
  1569. # If there's only one server in the cluster, its ``host`` is ''
  1570. # Fix it to the host in startup_nodes
  1571. if (
  1572. len(cluster_slots) == 1
  1573. and len(cluster_slots[0][2][0]) == 0
  1574. and len(self.startup_nodes) == 1
  1575. ):
  1576. cluster_slots[0][2][0] = startup_node.host
  1577. for slot in cluster_slots:
  1578. primary_node = slot[2]
  1579. host = str_if_bytes(primary_node[0])
  1580. if host == "":
  1581. host = startup_node.host
  1582. port = int(primary_node[1])
  1583. host, port = self.remap_host_port(host, port)
  1584. nodes_for_slot = []
  1585. target_node = self._get_or_create_cluster_node(
  1586. host, port, PRIMARY, tmp_nodes_cache
  1587. )
  1588. nodes_for_slot.append(target_node)
  1589. replica_nodes = slot[3:]
  1590. for replica_node in replica_nodes:
  1591. host = str_if_bytes(replica_node[0])
  1592. port = int(replica_node[1])
  1593. host, port = self.remap_host_port(host, port)
  1594. target_replica_node = self._get_or_create_cluster_node(
  1595. host, port, REPLICA, tmp_nodes_cache
  1596. )
  1597. nodes_for_slot.append(target_replica_node)
  1598. for i in range(int(slot[0]), int(slot[1]) + 1):
  1599. if i not in tmp_slots:
  1600. tmp_slots[i] = nodes_for_slot
  1601. else:
  1602. # Validate that 2 nodes want to use the same slot cache
  1603. # setup
  1604. tmp_slot = tmp_slots[i][0]
  1605. if tmp_slot.name != target_node.name:
  1606. disagreements.append(
  1607. f"{tmp_slot.name} vs {target_node.name} on slot: {i}"
  1608. )
  1609. if len(disagreements) > 5:
  1610. raise RedisClusterException(
  1611. f"startup_nodes could not agree on a valid "
  1612. f"slots cache: {', '.join(disagreements)}"
  1613. )
  1614. fully_covered = self.check_slots_coverage(tmp_slots)
  1615. if fully_covered:
  1616. # Don't need to continue to the next startup node if all
  1617. # slots are covered
  1618. break
  1619. if not startup_nodes_reachable:
  1620. raise RedisClusterException(
  1621. f"Redis Cluster cannot be connected. Please provide at least "
  1622. f"one reachable node: {str(exception)}"
  1623. ) from exception
  1624. if self._cache is None and self._cache_config is not None:
  1625. if self._cache_factory is None:
  1626. self._cache = CacheFactory(self._cache_config).get_cache()
  1627. else:
  1628. self._cache = self._cache_factory.get_cache()
  1629. # Create Redis connections to all nodes
  1630. self.create_redis_connections(list(tmp_nodes_cache.values()))
  1631. # Check if the slots are not fully covered
  1632. if not fully_covered and self._require_full_coverage:
  1633. # Despite the requirement that the slots be covered, there
  1634. # isn't a full coverage
  1635. raise RedisClusterException(
  1636. f"All slots are not covered after query all startup_nodes. "
  1637. f"{len(tmp_slots)} of {REDIS_CLUSTER_HASH_SLOTS} "
  1638. f"covered..."
  1639. )
  1640. # Set the tmp variables to the real variables
  1641. self.nodes_cache = tmp_nodes_cache
  1642. self.slots_cache = tmp_slots
  1643. # Set the default node
  1644. self.default_node = self.get_nodes_by_server_type(PRIMARY)[0]
  1645. if self._dynamic_startup_nodes:
  1646. # Populate the startup nodes with all discovered nodes
  1647. self.startup_nodes = tmp_nodes_cache
  1648. # If initialize was called after a MovedError, clear it
  1649. self._moved_exception = None
  1650. def close(self) -> None:
  1651. self.default_node = None
  1652. for node in self.nodes_cache.values():
  1653. if node.redis_connection:
  1654. node.redis_connection.close()
  1655. def reset(self):
  1656. try:
  1657. self.read_load_balancer.reset()
  1658. except TypeError:
  1659. # The read_load_balancer is None, do nothing
  1660. pass
  1661. def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
  1662. """
  1663. Remap the host and port returned from the cluster to a different
  1664. internal value. Useful if the client is not connecting directly
  1665. to the cluster.
  1666. """
  1667. if self.address_remap:
  1668. return self.address_remap((host, port))
  1669. return host, port
  1670. def find_connection_owner(self, connection: Connection) -> Optional[Redis]:
  1671. node_name = get_node_name(connection.host, connection.port)
  1672. for node in tuple(self.nodes_cache.values()):
  1673. if node.redis_connection:
  1674. conn_args = node.redis_connection.connection_pool.connection_kwargs
  1675. if node_name == get_node_name(
  1676. conn_args.get("host"), conn_args.get("port")
  1677. ):
  1678. return node
  1679. class ClusterPubSub(PubSub):
  1680. """
  1681. Wrapper for PubSub class.
  1682. IMPORTANT: before using ClusterPubSub, read about the known limitations
  1683. with pubsub in Cluster mode and learn how to workaround them:
  1684. https://redis-py-cluster.readthedocs.io/en/stable/pubsub.html
  1685. """
  1686. def __init__(
  1687. self,
  1688. redis_cluster,
  1689. node=None,
  1690. host=None,
  1691. port=None,
  1692. push_handler_func=None,
  1693. event_dispatcher: Optional["EventDispatcher"] = None,
  1694. **kwargs,
  1695. ):
  1696. """
  1697. When a pubsub instance is created without specifying a node, a single
  1698. node will be transparently chosen for the pubsub connection on the
  1699. first command execution. The node will be determined by:
  1700. 1. Hashing the channel name in the request to find its keyslot
  1701. 2. Selecting a node that handles the keyslot: If read_from_replicas is
  1702. set to true or load_balancing_strategy is set, a replica can be selected.
  1703. :type redis_cluster: RedisCluster
  1704. :type node: ClusterNode
  1705. :type host: str
  1706. :type port: int
  1707. """
  1708. self.node = None
  1709. self.set_pubsub_node(redis_cluster, node, host, port)
  1710. connection_pool = (
  1711. None
  1712. if self.node is None
  1713. else redis_cluster.get_redis_connection(self.node).connection_pool
  1714. )
  1715. self.cluster = redis_cluster
  1716. self.node_pubsub_mapping = {}
  1717. self._pubsubs_generator = self._pubsubs_generator()
  1718. if event_dispatcher is None:
  1719. self._event_dispatcher = EventDispatcher()
  1720. else:
  1721. self._event_dispatcher = event_dispatcher
  1722. super().__init__(
  1723. connection_pool=connection_pool,
  1724. encoder=redis_cluster.encoder,
  1725. push_handler_func=push_handler_func,
  1726. event_dispatcher=self._event_dispatcher,
  1727. **kwargs,
  1728. )
  1729. def set_pubsub_node(self, cluster, node=None, host=None, port=None):
  1730. """
  1731. The pubsub node will be set according to the passed node, host and port
  1732. When none of the node, host, or port are specified - the node is set
  1733. to None and will be determined by the keyslot of the channel in the
  1734. first command to be executed.
  1735. RedisClusterException will be thrown if the passed node does not exist
  1736. in the cluster.
  1737. If host is passed without port, or vice versa, a DataError will be
  1738. thrown.
  1739. :type cluster: RedisCluster
  1740. :type node: ClusterNode
  1741. :type host: str
  1742. :type port: int
  1743. """
  1744. if node is not None:
  1745. # node is passed by the user
  1746. self._raise_on_invalid_node(cluster, node, node.host, node.port)
  1747. pubsub_node = node
  1748. elif host is not None and port is not None:
  1749. # host and port passed by the user
  1750. node = cluster.get_node(host=host, port=port)
  1751. self._raise_on_invalid_node(cluster, node, host, port)
  1752. pubsub_node = node
  1753. elif any([host, port]) is True:
  1754. # only 'host' or 'port' passed
  1755. raise DataError("Passing a host requires passing a port, and vice versa")
  1756. else:
  1757. # nothing passed by the user. set node to None
  1758. pubsub_node = None
  1759. self.node = pubsub_node
  1760. def get_pubsub_node(self):
  1761. """
  1762. Get the node that is being used as the pubsub connection
  1763. """
  1764. return self.node
  1765. def _raise_on_invalid_node(self, redis_cluster, node, host, port):
  1766. """
  1767. Raise a RedisClusterException if the node is None or doesn't exist in
  1768. the cluster.
  1769. """
  1770. if node is None or redis_cluster.get_node(node_name=node.name) is None:
  1771. raise RedisClusterException(
  1772. f"Node {host}:{port} doesn't exist in the cluster"
  1773. )
  1774. def execute_command(self, *args):
  1775. """
  1776. Execute a subscribe/unsubscribe command.
  1777. Taken code from redis-py and tweak to make it work within a cluster.
  1778. """
  1779. # NOTE: don't parse the response in this function -- it could pull a
  1780. # legitimate message off the stack if the connection is already
  1781. # subscribed to one or more channels
  1782. if self.connection is None:
  1783. if self.connection_pool is None:
  1784. if len(args) > 1:
  1785. # Hash the first channel and get one of the nodes holding
  1786. # this slot
  1787. channel = args[1]
  1788. slot = self.cluster.keyslot(channel)
  1789. node = self.cluster.nodes_manager.get_node_from_slot(
  1790. slot,
  1791. self.cluster.read_from_replicas,
  1792. self.cluster.load_balancing_strategy,
  1793. )
  1794. else:
  1795. # Get a random node
  1796. node = self.cluster.get_random_node()
  1797. self.node = node
  1798. redis_connection = self.cluster.get_redis_connection(node)
  1799. self.connection_pool = redis_connection.connection_pool
  1800. self.connection = self.connection_pool.get_connection()
  1801. # register a callback that re-subscribes to any channels we
  1802. # were listening to when we were disconnected
  1803. self.connection.register_connect_callback(self.on_connect)
  1804. if self.push_handler_func is not None:
  1805. self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
  1806. self._event_dispatcher.dispatch(
  1807. AfterPubSubConnectionInstantiationEvent(
  1808. self.connection, self.connection_pool, ClientType.SYNC, self._lock
  1809. )
  1810. )
  1811. connection = self.connection
  1812. self._execute(connection, connection.send_command, *args)
  1813. def _get_node_pubsub(self, node):
  1814. try:
  1815. return self.node_pubsub_mapping[node.name]
  1816. except KeyError:
  1817. pubsub = node.redis_connection.pubsub(
  1818. push_handler_func=self.push_handler_func
  1819. )
  1820. self.node_pubsub_mapping[node.name] = pubsub
  1821. return pubsub
  1822. def _sharded_message_generator(self):
  1823. for _ in range(len(self.node_pubsub_mapping)):
  1824. pubsub = next(self._pubsubs_generator)
  1825. message = pubsub.get_message()
  1826. if message is not None:
  1827. return message
  1828. return None
  1829. def _pubsubs_generator(self):
  1830. while True:
  1831. yield from self.node_pubsub_mapping.values()
  1832. def get_sharded_message(
  1833. self, ignore_subscribe_messages=False, timeout=0.0, target_node=None
  1834. ):
  1835. if target_node:
  1836. message = self.node_pubsub_mapping[target_node.name].get_message(
  1837. ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout
  1838. )
  1839. else:
  1840. message = self._sharded_message_generator()
  1841. if message is None:
  1842. return None
  1843. elif str_if_bytes(message["type"]) == "sunsubscribe":
  1844. if message["channel"] in self.pending_unsubscribe_shard_channels:
  1845. self.pending_unsubscribe_shard_channels.remove(message["channel"])
  1846. self.shard_channels.pop(message["channel"], None)
  1847. node = self.cluster.get_node_from_key(message["channel"])
  1848. if self.node_pubsub_mapping[node.name].subscribed is False:
  1849. self.node_pubsub_mapping.pop(node.name)
  1850. if not self.channels and not self.patterns and not self.shard_channels:
  1851. # There are no subscriptions anymore, set subscribed_event flag
  1852. # to false
  1853. self.subscribed_event.clear()
  1854. if self.ignore_subscribe_messages or ignore_subscribe_messages:
  1855. return None
  1856. return message
  1857. def ssubscribe(self, *args, **kwargs):
  1858. if args:
  1859. args = list_or_args(args[0], args[1:])
  1860. s_channels = dict.fromkeys(args)
  1861. s_channels.update(kwargs)
  1862. for s_channel, handler in s_channels.items():
  1863. node = self.cluster.get_node_from_key(s_channel)
  1864. pubsub = self._get_node_pubsub(node)
  1865. if handler:
  1866. pubsub.ssubscribe(**{s_channel: handler})
  1867. else:
  1868. pubsub.ssubscribe(s_channel)
  1869. self.shard_channels.update(pubsub.shard_channels)
  1870. self.pending_unsubscribe_shard_channels.difference_update(
  1871. self._normalize_keys({s_channel: None})
  1872. )
  1873. if pubsub.subscribed and not self.subscribed:
  1874. self.subscribed_event.set()
  1875. self.health_check_response_counter = 0
  1876. def sunsubscribe(self, *args):
  1877. if args:
  1878. args = list_or_args(args[0], args[1:])
  1879. else:
  1880. args = self.shard_channels
  1881. for s_channel in args:
  1882. node = self.cluster.get_node_from_key(s_channel)
  1883. p = self._get_node_pubsub(node)
  1884. p.sunsubscribe(s_channel)
  1885. self.pending_unsubscribe_shard_channels.update(
  1886. p.pending_unsubscribe_shard_channels
  1887. )
  1888. def get_redis_connection(self):
  1889. """
  1890. Get the Redis connection of the pubsub connected node.
  1891. """
  1892. if self.node is not None:
  1893. return self.node.redis_connection
  1894. def disconnect(self):
  1895. """
  1896. Disconnect the pubsub connection.
  1897. """
  1898. if self.connection:
  1899. self.connection.disconnect()
  1900. for pubsub in self.node_pubsub_mapping.values():
  1901. pubsub.connection.disconnect()
  1902. class ClusterPipeline(RedisCluster):
  1903. """
  1904. Support for Redis pipeline
  1905. in cluster mode
  1906. """
  1907. ERRORS_ALLOW_RETRY = (
  1908. ConnectionError,
  1909. TimeoutError,
  1910. MovedError,
  1911. AskError,
  1912. TryAgainError,
  1913. )
  1914. NO_SLOTS_COMMANDS = {"UNWATCH"}
  1915. IMMEDIATE_EXECUTE_COMMANDS = {"WATCH", "UNWATCH"}
  1916. UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
  1917. @deprecated_args(
  1918. args_to_warn=[
  1919. "cluster_error_retry_attempts",
  1920. ],
  1921. reason="Please configure the 'retry' object instead",
  1922. version="6.0.0",
  1923. )
  1924. def __init__(
  1925. self,
  1926. nodes_manager: "NodesManager",
  1927. commands_parser: "CommandsParser",
  1928. result_callbacks: Optional[Dict[str, Callable]] = None,
  1929. cluster_response_callbacks: Optional[Dict[str, Callable]] = None,
  1930. startup_nodes: Optional[List["ClusterNode"]] = None,
  1931. read_from_replicas: bool = False,
  1932. load_balancing_strategy: Optional[LoadBalancingStrategy] = None,
  1933. cluster_error_retry_attempts: int = 3,
  1934. reinitialize_steps: int = 5,
  1935. retry: Optional[Retry] = None,
  1936. lock=None,
  1937. transaction=False,
  1938. **kwargs,
  1939. ):
  1940. """ """
  1941. self.command_stack = []
  1942. self.nodes_manager = nodes_manager
  1943. self.commands_parser = commands_parser
  1944. self.refresh_table_asap = False
  1945. self.result_callbacks = (
  1946. result_callbacks or self.__class__.RESULT_CALLBACKS.copy()
  1947. )
  1948. self.startup_nodes = startup_nodes if startup_nodes else []
  1949. self.read_from_replicas = read_from_replicas
  1950. self.load_balancing_strategy = load_balancing_strategy
  1951. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  1952. self.cluster_response_callbacks = cluster_response_callbacks
  1953. self.reinitialize_counter = 0
  1954. self.reinitialize_steps = reinitialize_steps
  1955. if retry is not None:
  1956. self.retry = retry
  1957. else:
  1958. self.retry = Retry(
  1959. backoff=ExponentialWithJitterBackoff(base=1, cap=10),
  1960. retries=cluster_error_retry_attempts,
  1961. )
  1962. self.encoder = Encoder(
  1963. kwargs.get("encoding", "utf-8"),
  1964. kwargs.get("encoding_errors", "strict"),
  1965. kwargs.get("decode_responses", False),
  1966. )
  1967. if lock is None:
  1968. lock = threading.RLock()
  1969. self._lock = lock
  1970. self.parent_execute_command = super().execute_command
  1971. self._execution_strategy: ExecutionStrategy = (
  1972. PipelineStrategy(self) if not transaction else TransactionStrategy(self)
  1973. )
  1974. def __repr__(self):
  1975. """ """
  1976. return f"{type(self).__name__}"
  1977. def __enter__(self):
  1978. """ """
  1979. return self
  1980. def __exit__(self, exc_type, exc_value, traceback):
  1981. """ """
  1982. self.reset()
  1983. def __del__(self):
  1984. try:
  1985. self.reset()
  1986. except Exception:
  1987. pass
  1988. def __len__(self):
  1989. """ """
  1990. return len(self._execution_strategy.command_queue)
  1991. def __bool__(self):
  1992. "Pipeline instances should always evaluate to True on Python 3+"
  1993. return True
  1994. def execute_command(self, *args, **kwargs):
  1995. """
  1996. Wrapper function for pipeline_execute_command
  1997. """
  1998. return self._execution_strategy.execute_command(*args, **kwargs)
  1999. def pipeline_execute_command(self, *args, **options):
  2000. """
  2001. Stage a command to be executed when execute() is next called
  2002. Returns the current Pipeline object back so commands can be
  2003. chained together, such as:
  2004. pipe = pipe.set('foo', 'bar').incr('baz').decr('bang')
  2005. At some other point, you can then run: pipe.execute(),
  2006. which will execute all commands queued in the pipe.
  2007. """
  2008. return self._execution_strategy.execute_command(*args, **options)
  2009. def annotate_exception(self, exception, number, command):
  2010. """
  2011. Provides extra context to the exception prior to it being handled
  2012. """
  2013. self._execution_strategy.annotate_exception(exception, number, command)
  2014. def execute(self, raise_on_error: bool = True) -> List[Any]:
  2015. """
  2016. Execute all the commands in the current pipeline
  2017. """
  2018. try:
  2019. return self._execution_strategy.execute(raise_on_error)
  2020. finally:
  2021. self.reset()
  2022. def reset(self):
  2023. """
  2024. Reset back to empty pipeline.
  2025. """
  2026. self._execution_strategy.reset()
  2027. def send_cluster_commands(
  2028. self, stack, raise_on_error=True, allow_redirections=True
  2029. ):
  2030. return self._execution_strategy.send_cluster_commands(
  2031. stack, raise_on_error=raise_on_error, allow_redirections=allow_redirections
  2032. )
  2033. def exists(self, *keys):
  2034. return self._execution_strategy.exists(*keys)
  2035. def eval(self):
  2036. """ """
  2037. return self._execution_strategy.eval()
  2038. def multi(self):
  2039. """
  2040. Start a transactional block of the pipeline after WATCH commands
  2041. are issued. End the transactional block with `execute`.
  2042. """
  2043. self._execution_strategy.multi()
  2044. def load_scripts(self):
  2045. """ """
  2046. self._execution_strategy.load_scripts()
  2047. def discard(self):
  2048. """ """
  2049. self._execution_strategy.discard()
  2050. def watch(self, *names):
  2051. """Watches the values at keys ``names``"""
  2052. self._execution_strategy.watch(*names)
  2053. def unwatch(self):
  2054. """Unwatches all previously specified keys"""
  2055. self._execution_strategy.unwatch()
  2056. def script_load_for_pipeline(self, *args, **kwargs):
  2057. self._execution_strategy.script_load_for_pipeline(*args, **kwargs)
  2058. def delete(self, *names):
  2059. self._execution_strategy.delete(*names)
  2060. def unlink(self, *names):
  2061. self._execution_strategy.unlink(*names)
  2062. def block_pipeline_command(name: str) -> Callable[..., Any]:
  2063. """
  2064. Prints error because some pipelined commands should
  2065. be blocked when running in cluster-mode
  2066. """
  2067. def inner(*args, **kwargs):
  2068. raise RedisClusterException(
  2069. f"ERROR: Calling pipelined function {name} is blocked "
  2070. f"when running redis in cluster mode..."
  2071. )
  2072. return inner
  2073. # Blocked pipeline commands
  2074. PIPELINE_BLOCKED_COMMANDS = (
  2075. "BGREWRITEAOF",
  2076. "BGSAVE",
  2077. "BITOP",
  2078. "BRPOPLPUSH",
  2079. "CLIENT GETNAME",
  2080. "CLIENT KILL",
  2081. "CLIENT LIST",
  2082. "CLIENT SETNAME",
  2083. "CLIENT",
  2084. "CONFIG GET",
  2085. "CONFIG RESETSTAT",
  2086. "CONFIG REWRITE",
  2087. "CONFIG SET",
  2088. "CONFIG",
  2089. "DBSIZE",
  2090. "ECHO",
  2091. "EVALSHA",
  2092. "FLUSHALL",
  2093. "FLUSHDB",
  2094. "INFO",
  2095. "KEYS",
  2096. "LASTSAVE",
  2097. "MGET",
  2098. "MGET NONATOMIC",
  2099. "MOVE",
  2100. "MSET",
  2101. "MSET NONATOMIC",
  2102. "MSETNX",
  2103. "PFCOUNT",
  2104. "PFMERGE",
  2105. "PING",
  2106. "PUBLISH",
  2107. "RANDOMKEY",
  2108. "READONLY",
  2109. "READWRITE",
  2110. "RENAME",
  2111. "RENAMENX",
  2112. "RPOPLPUSH",
  2113. "SAVE",
  2114. "SCAN",
  2115. "SCRIPT EXISTS",
  2116. "SCRIPT FLUSH",
  2117. "SCRIPT KILL",
  2118. "SCRIPT LOAD",
  2119. "SCRIPT",
  2120. "SDIFF",
  2121. "SDIFFSTORE",
  2122. "SENTINEL GET MASTER ADDR BY NAME",
  2123. "SENTINEL MASTER",
  2124. "SENTINEL MASTERS",
  2125. "SENTINEL MONITOR",
  2126. "SENTINEL REMOVE",
  2127. "SENTINEL SENTINELS",
  2128. "SENTINEL SET",
  2129. "SENTINEL SLAVES",
  2130. "SENTINEL",
  2131. "SHUTDOWN",
  2132. "SINTER",
  2133. "SINTERSTORE",
  2134. "SLAVEOF",
  2135. "SLOWLOG GET",
  2136. "SLOWLOG LEN",
  2137. "SLOWLOG RESET",
  2138. "SLOWLOG",
  2139. "SMOVE",
  2140. "SORT",
  2141. "SUNION",
  2142. "SUNIONSTORE",
  2143. "TIME",
  2144. )
  2145. for command in PIPELINE_BLOCKED_COMMANDS:
  2146. command = command.replace(" ", "_").lower()
  2147. setattr(ClusterPipeline, command, block_pipeline_command(command))
  2148. class PipelineCommand:
  2149. """ """
  2150. def __init__(self, args, options=None, position=None):
  2151. self.args = args
  2152. if options is None:
  2153. options = {}
  2154. self.options = options
  2155. self.position = position
  2156. self.result = None
  2157. self.node = None
  2158. self.asking = False
  2159. class NodeCommands:
  2160. """ """
  2161. def __init__(self, parse_response, connection_pool, connection):
  2162. """ """
  2163. self.parse_response = parse_response
  2164. self.connection_pool = connection_pool
  2165. self.connection = connection
  2166. self.commands = []
  2167. def append(self, c):
  2168. """ """
  2169. self.commands.append(c)
  2170. def write(self):
  2171. """
  2172. Code borrowed from Redis so it can be fixed
  2173. """
  2174. connection = self.connection
  2175. commands = self.commands
  2176. # We are going to clobber the commands with the write, so go ahead
  2177. # and ensure that nothing is sitting there from a previous run.
  2178. for c in commands:
  2179. c.result = None
  2180. # build up all commands into a single request to increase network perf
  2181. # send all the commands and catch connection and timeout errors.
  2182. try:
  2183. connection.send_packed_command(
  2184. connection.pack_commands([c.args for c in commands])
  2185. )
  2186. except (ConnectionError, TimeoutError) as e:
  2187. for c in commands:
  2188. c.result = e
  2189. def read(self):
  2190. """ """
  2191. connection = self.connection
  2192. for c in self.commands:
  2193. # if there is a result on this command,
  2194. # it means we ran into an exception
  2195. # like a connection error. Trying to parse
  2196. # a response on a connection that
  2197. # is no longer open will result in a
  2198. # connection error raised by redis-py.
  2199. # but redis-py doesn't check in parse_response
  2200. # that the sock object is
  2201. # still set and if you try to
  2202. # read from a closed connection, it will
  2203. # result in an AttributeError because
  2204. # it will do a readline() call on None.
  2205. # This can have all kinds of nasty side-effects.
  2206. # Treating this case as a connection error
  2207. # is fine because it will dump
  2208. # the connection object back into the
  2209. # pool and on the next write, it will
  2210. # explicitly open the connection and all will be well.
  2211. if c.result is None:
  2212. try:
  2213. c.result = self.parse_response(connection, c.args[0], **c.options)
  2214. except (ConnectionError, TimeoutError) as e:
  2215. for c in self.commands:
  2216. c.result = e
  2217. return
  2218. except RedisError:
  2219. c.result = sys.exc_info()[1]
  2220. class ExecutionStrategy(ABC):
  2221. @property
  2222. @abstractmethod
  2223. def command_queue(self):
  2224. pass
  2225. @abstractmethod
  2226. def execute_command(self, *args, **kwargs):
  2227. """
  2228. Execution flow for current execution strategy.
  2229. See: ClusterPipeline.execute_command()
  2230. """
  2231. pass
  2232. @abstractmethod
  2233. def annotate_exception(self, exception, number, command):
  2234. """
  2235. Annotate exception according to current execution strategy.
  2236. See: ClusterPipeline.annotate_exception()
  2237. """
  2238. pass
  2239. @abstractmethod
  2240. def pipeline_execute_command(self, *args, **options):
  2241. """
  2242. Pipeline execution flow for current execution strategy.
  2243. See: ClusterPipeline.pipeline_execute_command()
  2244. """
  2245. pass
  2246. @abstractmethod
  2247. def execute(self, raise_on_error: bool = True) -> List[Any]:
  2248. """
  2249. Executes current execution strategy.
  2250. See: ClusterPipeline.execute()
  2251. """
  2252. pass
  2253. @abstractmethod
  2254. def send_cluster_commands(
  2255. self, stack, raise_on_error=True, allow_redirections=True
  2256. ):
  2257. """
  2258. Sends commands according to current execution strategy.
  2259. See: ClusterPipeline.send_cluster_commands()
  2260. """
  2261. pass
  2262. @abstractmethod
  2263. def reset(self):
  2264. """
  2265. Resets current execution strategy.
  2266. See: ClusterPipeline.reset()
  2267. """
  2268. pass
  2269. @abstractmethod
  2270. def exists(self, *keys):
  2271. pass
  2272. @abstractmethod
  2273. def eval(self):
  2274. pass
  2275. @abstractmethod
  2276. def multi(self):
  2277. """
  2278. Starts transactional context.
  2279. See: ClusterPipeline.multi()
  2280. """
  2281. pass
  2282. @abstractmethod
  2283. def load_scripts(self):
  2284. pass
  2285. @abstractmethod
  2286. def watch(self, *names):
  2287. pass
  2288. @abstractmethod
  2289. def unwatch(self):
  2290. """
  2291. Unwatches all previously specified keys
  2292. See: ClusterPipeline.unwatch()
  2293. """
  2294. pass
  2295. @abstractmethod
  2296. def script_load_for_pipeline(self, *args, **kwargs):
  2297. pass
  2298. @abstractmethod
  2299. def delete(self, *names):
  2300. """
  2301. "Delete a key specified by ``names``"
  2302. See: ClusterPipeline.delete()
  2303. """
  2304. pass
  2305. @abstractmethod
  2306. def unlink(self, *names):
  2307. """
  2308. "Unlink a key specified by ``names``"
  2309. See: ClusterPipeline.unlink()
  2310. """
  2311. pass
  2312. @abstractmethod
  2313. def discard(self):
  2314. pass
  2315. class AbstractStrategy(ExecutionStrategy):
  2316. def __init__(
  2317. self,
  2318. pipe: ClusterPipeline,
  2319. ):
  2320. self._command_queue: List[PipelineCommand] = []
  2321. self._pipe = pipe
  2322. self._nodes_manager = self._pipe.nodes_manager
  2323. @property
  2324. def command_queue(self):
  2325. return self._command_queue
  2326. @command_queue.setter
  2327. def command_queue(self, queue: List[PipelineCommand]):
  2328. self._command_queue = queue
  2329. @abstractmethod
  2330. def execute_command(self, *args, **kwargs):
  2331. pass
  2332. def pipeline_execute_command(self, *args, **options):
  2333. self._command_queue.append(
  2334. PipelineCommand(args, options, len(self._command_queue))
  2335. )
  2336. return self._pipe
  2337. @abstractmethod
  2338. def execute(self, raise_on_error: bool = True) -> List[Any]:
  2339. pass
  2340. @abstractmethod
  2341. def send_cluster_commands(
  2342. self, stack, raise_on_error=True, allow_redirections=True
  2343. ):
  2344. pass
  2345. @abstractmethod
  2346. def reset(self):
  2347. pass
  2348. def exists(self, *keys):
  2349. return self.execute_command("EXISTS", *keys)
  2350. def eval(self):
  2351. """ """
  2352. raise RedisClusterException("method eval() is not implemented")
  2353. def load_scripts(self):
  2354. """ """
  2355. raise RedisClusterException("method load_scripts() is not implemented")
  2356. def script_load_for_pipeline(self, *args, **kwargs):
  2357. """ """
  2358. raise RedisClusterException(
  2359. "method script_load_for_pipeline() is not implemented"
  2360. )
  2361. def annotate_exception(self, exception, number, command):
  2362. """
  2363. Provides extra context to the exception prior to it being handled
  2364. """
  2365. cmd = " ".join(map(safe_str, command))
  2366. msg = (
  2367. f"Command # {number} ({truncate_text(cmd)}) of pipeline "
  2368. f"caused error: {exception.args[0]}"
  2369. )
  2370. exception.args = (msg,) + exception.args[1:]
  2371. class PipelineStrategy(AbstractStrategy):
  2372. def __init__(self, pipe: ClusterPipeline):
  2373. super().__init__(pipe)
  2374. self.command_flags = pipe.command_flags
  2375. def execute_command(self, *args, **kwargs):
  2376. return self.pipeline_execute_command(*args, **kwargs)
  2377. def _raise_first_error(self, stack):
  2378. """
  2379. Raise the first exception on the stack
  2380. """
  2381. for c in stack:
  2382. r = c.result
  2383. if isinstance(r, Exception):
  2384. self.annotate_exception(r, c.position + 1, c.args)
  2385. raise r
  2386. def execute(self, raise_on_error: bool = True) -> List[Any]:
  2387. stack = self._command_queue
  2388. if not stack:
  2389. return []
  2390. try:
  2391. return self.send_cluster_commands(stack, raise_on_error)
  2392. finally:
  2393. self.reset()
  2394. def reset(self):
  2395. """
  2396. Reset back to empty pipeline.
  2397. """
  2398. self._command_queue = []
  2399. def send_cluster_commands(
  2400. self, stack, raise_on_error=True, allow_redirections=True
  2401. ):
  2402. """
  2403. Wrapper for RedisCluster.ERRORS_ALLOW_RETRY errors handling.
  2404. If one of the retryable exceptions has been thrown we assume that:
  2405. - connection_pool was disconnected
  2406. - connection_pool was reseted
  2407. - refereh_table_asap set to True
  2408. It will try the number of times specified by
  2409. the retries in config option "self.retry"
  2410. which defaults to 3 unless manually configured.
  2411. If it reaches the number of times, the command will
  2412. raises ClusterDownException.
  2413. """
  2414. if not stack:
  2415. return []
  2416. retry_attempts = self._pipe.retry.get_retries()
  2417. while True:
  2418. try:
  2419. return self._send_cluster_commands(
  2420. stack,
  2421. raise_on_error=raise_on_error,
  2422. allow_redirections=allow_redirections,
  2423. )
  2424. except RedisCluster.ERRORS_ALLOW_RETRY as e:
  2425. if retry_attempts > 0:
  2426. # Try again with the new cluster setup. All other errors
  2427. # should be raised.
  2428. retry_attempts -= 1
  2429. pass
  2430. else:
  2431. raise e
  2432. def _send_cluster_commands(
  2433. self, stack, raise_on_error=True, allow_redirections=True
  2434. ):
  2435. """
  2436. Send a bunch of cluster commands to the redis cluster.
  2437. `allow_redirections` If the pipeline should follow
  2438. `ASK` & `MOVED` responses automatically. If set
  2439. to false it will raise RedisClusterException.
  2440. """
  2441. # the first time sending the commands we send all of
  2442. # the commands that were queued up.
  2443. # if we have to run through it again, we only retry
  2444. # the commands that failed.
  2445. attempt = sorted(stack, key=lambda x: x.position)
  2446. is_default_node = False
  2447. # build a list of node objects based on node names we need to
  2448. nodes = {}
  2449. # as we move through each command that still needs to be processed,
  2450. # we figure out the slot number that command maps to, then from
  2451. # the slot determine the node.
  2452. for c in attempt:
  2453. while True:
  2454. # refer to our internal node -> slot table that
  2455. # tells us where a given command should route to.
  2456. # (it might be possible we have a cached node that no longer
  2457. # exists in the cluster, which is why we do this in a loop)
  2458. passed_targets = c.options.pop("target_nodes", None)
  2459. if passed_targets and not self._is_nodes_flag(passed_targets):
  2460. target_nodes = self._parse_target_nodes(passed_targets)
  2461. else:
  2462. target_nodes = self._determine_nodes(
  2463. *c.args, node_flag=passed_targets
  2464. )
  2465. if not target_nodes:
  2466. raise RedisClusterException(
  2467. f"No targets were found to execute {c.args} command on"
  2468. )
  2469. if len(target_nodes) > 1:
  2470. raise RedisClusterException(
  2471. f"Too many targets for command {c.args}"
  2472. )
  2473. node = target_nodes[0]
  2474. if node == self._pipe.get_default_node():
  2475. is_default_node = True
  2476. # now that we know the name of the node
  2477. # ( it's just a string in the form of host:port )
  2478. # we can build a list of commands for each node.
  2479. node_name = node.name
  2480. if node_name not in nodes:
  2481. redis_node = self._pipe.get_redis_connection(node)
  2482. try:
  2483. connection = get_connection(redis_node)
  2484. except (ConnectionError, TimeoutError):
  2485. for n in nodes.values():
  2486. n.connection_pool.release(n.connection)
  2487. # Connection retries are being handled in the node's
  2488. # Retry object. Reinitialize the node -> slot table.
  2489. self._nodes_manager.initialize()
  2490. if is_default_node:
  2491. self._pipe.replace_default_node()
  2492. raise
  2493. nodes[node_name] = NodeCommands(
  2494. redis_node.parse_response,
  2495. redis_node.connection_pool,
  2496. connection,
  2497. )
  2498. nodes[node_name].append(c)
  2499. break
  2500. # send the commands in sequence.
  2501. # we write to all the open sockets for each node first,
  2502. # before reading anything
  2503. # this allows us to flush all the requests out across the
  2504. # network
  2505. # so that we can read them from different sockets as they come back.
  2506. # we dont' multiplex on the sockets as they come available,
  2507. # but that shouldn't make too much difference.
  2508. try:
  2509. node_commands = nodes.values()
  2510. for n in node_commands:
  2511. n.write()
  2512. for n in node_commands:
  2513. n.read()
  2514. finally:
  2515. # release all of the redis connections we allocated earlier
  2516. # back into the connection pool.
  2517. # we used to do this step as part of a try/finally block,
  2518. # but it is really dangerous to
  2519. # release connections back into the pool if for some
  2520. # reason the socket has data still left in it
  2521. # from a previous operation. The write and
  2522. # read operations already have try/catch around them for
  2523. # all known types of errors including connection
  2524. # and socket level errors.
  2525. # So if we hit an exception, something really bad
  2526. # happened and putting any oF
  2527. # these connections back into the pool is a very bad idea.
  2528. # the socket might have unread buffer still sitting in it,
  2529. # and then the next time we read from it we pass the
  2530. # buffered result back from a previous command and
  2531. # every single request after to that connection will always get
  2532. # a mismatched result.
  2533. for n in nodes.values():
  2534. n.connection_pool.release(n.connection)
  2535. # if the response isn't an exception it is a
  2536. # valid response from the node
  2537. # we're all done with that command, YAY!
  2538. # if we have more commands to attempt, we've run into problems.
  2539. # collect all the commands we are allowed to retry.
  2540. # (MOVED, ASK, or connection errors or timeout errors)
  2541. attempt = sorted(
  2542. (
  2543. c
  2544. for c in attempt
  2545. if isinstance(c.result, ClusterPipeline.ERRORS_ALLOW_RETRY)
  2546. ),
  2547. key=lambda x: x.position,
  2548. )
  2549. if attempt and allow_redirections:
  2550. # RETRY MAGIC HAPPENS HERE!
  2551. # send these remaining commands one at a time using `execute_command`
  2552. # in the main client. This keeps our retry logic
  2553. # in one place mostly,
  2554. # and allows us to be more confident in correctness of behavior.
  2555. # at this point any speed gains from pipelining have been lost
  2556. # anyway, so we might as well make the best
  2557. # attempt to get the correct behavior.
  2558. #
  2559. # The client command will handle retries for each
  2560. # individual command sequentially as we pass each
  2561. # one into `execute_command`. Any exceptions
  2562. # that bubble out should only appear once all
  2563. # retries have been exhausted.
  2564. #
  2565. # If a lot of commands have failed, we'll be setting the
  2566. # flag to rebuild the slots table from scratch.
  2567. # So MOVED errors should correct themselves fairly quickly.
  2568. self._pipe.reinitialize_counter += 1
  2569. if self._pipe._should_reinitialized():
  2570. self._nodes_manager.initialize()
  2571. if is_default_node:
  2572. self._pipe.replace_default_node()
  2573. for c in attempt:
  2574. try:
  2575. # send each command individually like we
  2576. # do in the main client.
  2577. c.result = self._pipe.parent_execute_command(*c.args, **c.options)
  2578. except RedisError as e:
  2579. c.result = e
  2580. # turn the response back into a simple flat array that corresponds
  2581. # to the sequence of commands issued in the stack in pipeline.execute()
  2582. response = []
  2583. for c in sorted(stack, key=lambda x: x.position):
  2584. if c.args[0] in self._pipe.cluster_response_callbacks:
  2585. # Remove keys entry, it needs only for cache.
  2586. c.options.pop("keys", None)
  2587. c.result = self._pipe.cluster_response_callbacks[c.args[0]](
  2588. c.result, **c.options
  2589. )
  2590. response.append(c.result)
  2591. if raise_on_error:
  2592. self._raise_first_error(stack)
  2593. return response
  2594. def _is_nodes_flag(self, target_nodes):
  2595. return isinstance(target_nodes, str) and target_nodes in self._pipe.node_flags
  2596. def _parse_target_nodes(self, target_nodes):
  2597. if isinstance(target_nodes, list):
  2598. nodes = target_nodes
  2599. elif isinstance(target_nodes, ClusterNode):
  2600. # Supports passing a single ClusterNode as a variable
  2601. nodes = [target_nodes]
  2602. elif isinstance(target_nodes, dict):
  2603. # Supports dictionaries of the format {node_name: node}.
  2604. # It enables to execute commands with multi nodes as follows:
  2605. # rc.cluster_save_config(rc.get_primaries())
  2606. nodes = target_nodes.values()
  2607. else:
  2608. raise TypeError(
  2609. "target_nodes type can be one of the following: "
  2610. "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
  2611. "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
  2612. f"The passed type is {type(target_nodes)}"
  2613. )
  2614. return nodes
  2615. def _determine_nodes(self, *args, **kwargs) -> List["ClusterNode"]:
  2616. # Determine which nodes should be executed the command on.
  2617. # Returns a list of target nodes.
  2618. command = args[0].upper()
  2619. if (
  2620. len(args) >= 2
  2621. and f"{args[0]} {args[1]}".upper() in self._pipe.command_flags
  2622. ):
  2623. command = f"{args[0]} {args[1]}".upper()
  2624. nodes_flag = kwargs.pop("nodes_flag", None)
  2625. if nodes_flag is not None:
  2626. # nodes flag passed by the user
  2627. command_flag = nodes_flag
  2628. else:
  2629. # get the nodes group for this command if it was predefined
  2630. command_flag = self._pipe.command_flags.get(command)
  2631. if command_flag == self._pipe.RANDOM:
  2632. # return a random node
  2633. return [self._pipe.get_random_node()]
  2634. elif command_flag == self._pipe.PRIMARIES:
  2635. # return all primaries
  2636. return self._pipe.get_primaries()
  2637. elif command_flag == self._pipe.REPLICAS:
  2638. # return all replicas
  2639. return self._pipe.get_replicas()
  2640. elif command_flag == self._pipe.ALL_NODES:
  2641. # return all nodes
  2642. return self._pipe.get_nodes()
  2643. elif command_flag == self._pipe.DEFAULT_NODE:
  2644. # return the cluster's default node
  2645. return [self._nodes_manager.default_node]
  2646. elif command in self._pipe.SEARCH_COMMANDS[0]:
  2647. return [self._nodes_manager.default_node]
  2648. else:
  2649. # get the node that holds the key's slot
  2650. slot = self._pipe.determine_slot(*args)
  2651. node = self._nodes_manager.get_node_from_slot(
  2652. slot,
  2653. self._pipe.read_from_replicas and command in READ_COMMANDS,
  2654. self._pipe.load_balancing_strategy
  2655. if command in READ_COMMANDS
  2656. else None,
  2657. )
  2658. return [node]
  2659. def multi(self):
  2660. raise RedisClusterException(
  2661. "method multi() is not supported outside of transactional context"
  2662. )
  2663. def discard(self):
  2664. raise RedisClusterException(
  2665. "method discard() is not supported outside of transactional context"
  2666. )
  2667. def watch(self, *names):
  2668. raise RedisClusterException(
  2669. "method watch() is not supported outside of transactional context"
  2670. )
  2671. def unwatch(self, *names):
  2672. raise RedisClusterException(
  2673. "method unwatch() is not supported outside of transactional context"
  2674. )
  2675. def delete(self, *names):
  2676. if len(names) != 1:
  2677. raise RedisClusterException(
  2678. "deleting multiple keys is not implemented in pipeline command"
  2679. )
  2680. return self.execute_command("DEL", names[0])
  2681. def unlink(self, *names):
  2682. if len(names) != 1:
  2683. raise RedisClusterException(
  2684. "unlinking multiple keys is not implemented in pipeline command"
  2685. )
  2686. return self.execute_command("UNLINK", names[0])
  2687. class TransactionStrategy(AbstractStrategy):
  2688. NO_SLOTS_COMMANDS = {"UNWATCH"}
  2689. IMMEDIATE_EXECUTE_COMMANDS = {"WATCH", "UNWATCH"}
  2690. UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
  2691. SLOT_REDIRECT_ERRORS = (AskError, MovedError)
  2692. CONNECTION_ERRORS = (
  2693. ConnectionError,
  2694. OSError,
  2695. ClusterDownError,
  2696. SlotNotCoveredError,
  2697. )
  2698. def __init__(self, pipe: ClusterPipeline):
  2699. super().__init__(pipe)
  2700. self._explicit_transaction = False
  2701. self._watching = False
  2702. self._pipeline_slots: Set[int] = set()
  2703. self._transaction_connection: Optional[Connection] = None
  2704. self._executing = False
  2705. self._retry = copy(self._pipe.retry)
  2706. self._retry.update_supported_errors(
  2707. RedisCluster.ERRORS_ALLOW_RETRY + self.SLOT_REDIRECT_ERRORS
  2708. )
  2709. def _get_client_and_connection_for_transaction(self) -> Tuple[Redis, Connection]:
  2710. """
  2711. Find a connection for a pipeline transaction.
  2712. For running an atomic transaction, watch keys ensure that contents have not been
  2713. altered as long as the watch commands for those keys were sent over the same
  2714. connection. So once we start watching a key, we fetch a connection to the
  2715. node that owns that slot and reuse it.
  2716. """
  2717. if not self._pipeline_slots:
  2718. raise RedisClusterException(
  2719. "At least a command with a key is needed to identify a node"
  2720. )
  2721. node: ClusterNode = self._nodes_manager.get_node_from_slot(
  2722. list(self._pipeline_slots)[0], False
  2723. )
  2724. redis_node: Redis = self._pipe.get_redis_connection(node)
  2725. if self._transaction_connection:
  2726. if not redis_node.connection_pool.owns_connection(
  2727. self._transaction_connection
  2728. ):
  2729. previous_node = self._nodes_manager.find_connection_owner(
  2730. self._transaction_connection
  2731. )
  2732. previous_node.connection_pool.release(self._transaction_connection)
  2733. self._transaction_connection = None
  2734. if not self._transaction_connection:
  2735. self._transaction_connection = get_connection(redis_node)
  2736. return redis_node, self._transaction_connection
  2737. def execute_command(self, *args, **kwargs):
  2738. slot_number: Optional[int] = None
  2739. if args[0] not in ClusterPipeline.NO_SLOTS_COMMANDS:
  2740. slot_number = self._pipe.determine_slot(*args)
  2741. if (
  2742. self._watching or args[0] in self.IMMEDIATE_EXECUTE_COMMANDS
  2743. ) and not self._explicit_transaction:
  2744. if args[0] == "WATCH":
  2745. self._validate_watch()
  2746. if slot_number is not None:
  2747. if self._pipeline_slots and slot_number not in self._pipeline_slots:
  2748. raise CrossSlotTransactionError(
  2749. "Cannot watch or send commands on different slots"
  2750. )
  2751. self._pipeline_slots.add(slot_number)
  2752. elif args[0] not in self.NO_SLOTS_COMMANDS:
  2753. raise RedisClusterException(
  2754. f"Cannot identify slot number for command: {args[0]},"
  2755. "it cannot be triggered in a transaction"
  2756. )
  2757. return self._immediate_execute_command(*args, **kwargs)
  2758. else:
  2759. if slot_number is not None:
  2760. self._pipeline_slots.add(slot_number)
  2761. return self.pipeline_execute_command(*args, **kwargs)
  2762. def _validate_watch(self):
  2763. if self._explicit_transaction:
  2764. raise RedisError("Cannot issue a WATCH after a MULTI")
  2765. self._watching = True
  2766. def _immediate_execute_command(self, *args, **options):
  2767. return self._retry.call_with_retry(
  2768. lambda: self._get_connection_and_send_command(*args, **options),
  2769. self._reinitialize_on_error,
  2770. )
  2771. def _get_connection_and_send_command(self, *args, **options):
  2772. redis_node, connection = self._get_client_and_connection_for_transaction()
  2773. return self._send_command_parse_response(
  2774. connection, redis_node, args[0], *args, **options
  2775. )
  2776. def _send_command_parse_response(
  2777. self, conn, redis_node: Redis, command_name, *args, **options
  2778. ):
  2779. """
  2780. Send a command and parse the response
  2781. """
  2782. conn.send_command(*args)
  2783. output = redis_node.parse_response(conn, command_name, **options)
  2784. if command_name in self.UNWATCH_COMMANDS:
  2785. self._watching = False
  2786. return output
  2787. def _reinitialize_on_error(self, error):
  2788. if self._watching:
  2789. if type(error) in self.SLOT_REDIRECT_ERRORS and self._executing:
  2790. raise WatchError("Slot rebalancing occurred while watching keys")
  2791. if (
  2792. type(error) in self.SLOT_REDIRECT_ERRORS
  2793. or type(error) in self.CONNECTION_ERRORS
  2794. ):
  2795. if self._transaction_connection:
  2796. self._transaction_connection = None
  2797. self._pipe.reinitialize_counter += 1
  2798. if self._pipe._should_reinitialized():
  2799. self._nodes_manager.initialize()
  2800. self.reinitialize_counter = 0
  2801. else:
  2802. self._nodes_manager.update_moved_exception(error)
  2803. self._executing = False
  2804. def _raise_first_error(self, responses, stack):
  2805. """
  2806. Raise the first exception on the stack
  2807. """
  2808. for r, cmd in zip(responses, stack):
  2809. if isinstance(r, Exception):
  2810. self.annotate_exception(r, cmd.position + 1, cmd.args)
  2811. raise r
  2812. def execute(self, raise_on_error: bool = True) -> List[Any]:
  2813. stack = self._command_queue
  2814. if not stack and (not self._watching or not self._pipeline_slots):
  2815. return []
  2816. return self._execute_transaction_with_retries(stack, raise_on_error)
  2817. def _execute_transaction_with_retries(
  2818. self, stack: List["PipelineCommand"], raise_on_error: bool
  2819. ):
  2820. return self._retry.call_with_retry(
  2821. lambda: self._execute_transaction(stack, raise_on_error),
  2822. self._reinitialize_on_error,
  2823. )
  2824. def _execute_transaction(
  2825. self, stack: List["PipelineCommand"], raise_on_error: bool
  2826. ):
  2827. if len(self._pipeline_slots) > 1:
  2828. raise CrossSlotTransactionError(
  2829. "All keys involved in a cluster transaction must map to the same slot"
  2830. )
  2831. self._executing = True
  2832. redis_node, connection = self._get_client_and_connection_for_transaction()
  2833. stack = chain(
  2834. [PipelineCommand(("MULTI",))],
  2835. stack,
  2836. [PipelineCommand(("EXEC",))],
  2837. )
  2838. commands = [c.args for c in stack if EMPTY_RESPONSE not in c.options]
  2839. packed_commands = connection.pack_commands(commands)
  2840. connection.send_packed_command(packed_commands)
  2841. errors = []
  2842. # parse off the response for MULTI
  2843. # NOTE: we need to handle ResponseErrors here and continue
  2844. # so that we read all the additional command messages from
  2845. # the socket
  2846. try:
  2847. redis_node.parse_response(connection, "MULTI")
  2848. except ResponseError as e:
  2849. self.annotate_exception(e, 0, "MULTI")
  2850. errors.append(e)
  2851. except self.CONNECTION_ERRORS as cluster_error:
  2852. self.annotate_exception(cluster_error, 0, "MULTI")
  2853. raise
  2854. # and all the other commands
  2855. for i, command in enumerate(self._command_queue):
  2856. if EMPTY_RESPONSE in command.options:
  2857. errors.append((i, command.options[EMPTY_RESPONSE]))
  2858. else:
  2859. try:
  2860. _ = redis_node.parse_response(connection, "_")
  2861. except self.SLOT_REDIRECT_ERRORS as slot_error:
  2862. self.annotate_exception(slot_error, i + 1, command.args)
  2863. errors.append(slot_error)
  2864. except self.CONNECTION_ERRORS as cluster_error:
  2865. self.annotate_exception(cluster_error, i + 1, command.args)
  2866. raise
  2867. except ResponseError as e:
  2868. self.annotate_exception(e, i + 1, command.args)
  2869. errors.append(e)
  2870. response = None
  2871. # parse the EXEC.
  2872. try:
  2873. response = redis_node.parse_response(connection, "EXEC")
  2874. except ExecAbortError:
  2875. if errors:
  2876. raise errors[0]
  2877. raise
  2878. self._executing = False
  2879. # EXEC clears any watched keys
  2880. self._watching = False
  2881. if response is None:
  2882. raise WatchError("Watched variable changed.")
  2883. # put any parse errors into the response
  2884. for i, e in errors:
  2885. response.insert(i, e)
  2886. if len(response) != len(self._command_queue):
  2887. raise InvalidPipelineStack(
  2888. "Unexpected response length for cluster pipeline EXEC."
  2889. " Command stack was {} but response had length {}".format(
  2890. [c.args[0] for c in self._command_queue], len(response)
  2891. )
  2892. )
  2893. # find any errors in the response and raise if necessary
  2894. if raise_on_error or len(errors) > 0:
  2895. self._raise_first_error(
  2896. response,
  2897. self._command_queue,
  2898. )
  2899. # We have to run response callbacks manually
  2900. data = []
  2901. for r, cmd in zip(response, self._command_queue):
  2902. if not isinstance(r, Exception):
  2903. command_name = cmd.args[0]
  2904. if command_name in self._pipe.cluster_response_callbacks:
  2905. r = self._pipe.cluster_response_callbacks[command_name](
  2906. r, **cmd.options
  2907. )
  2908. data.append(r)
  2909. return data
  2910. def reset(self):
  2911. self._command_queue = []
  2912. # make sure to reset the connection state in the event that we were
  2913. # watching something
  2914. if self._transaction_connection:
  2915. try:
  2916. # call this manually since our unwatch or
  2917. # immediate_execute_command methods can call reset()
  2918. self._transaction_connection.send_command("UNWATCH")
  2919. self._transaction_connection.read_response()
  2920. # we can safely return the connection to the pool here since we're
  2921. # sure we're no longer WATCHing anything
  2922. node = self._nodes_manager.find_connection_owner(
  2923. self._transaction_connection
  2924. )
  2925. node.redis_connection.connection_pool.release(
  2926. self._transaction_connection
  2927. )
  2928. self._transaction_connection = None
  2929. except self.CONNECTION_ERRORS:
  2930. # disconnect will also remove any previous WATCHes
  2931. if self._transaction_connection:
  2932. self._transaction_connection.disconnect()
  2933. # clean up the other instance attributes
  2934. self._watching = False
  2935. self._explicit_transaction = False
  2936. self._pipeline_slots = set()
  2937. self._executing = False
  2938. def send_cluster_commands(
  2939. self, stack, raise_on_error=True, allow_redirections=True
  2940. ):
  2941. raise NotImplementedError(
  2942. "send_cluster_commands cannot be executed in transactional context."
  2943. )
  2944. def multi(self):
  2945. if self._explicit_transaction:
  2946. raise RedisError("Cannot issue nested calls to MULTI")
  2947. if self._command_queue:
  2948. raise RedisError(
  2949. "Commands without an initial WATCH have already been issued"
  2950. )
  2951. self._explicit_transaction = True
  2952. def watch(self, *names):
  2953. if self._explicit_transaction:
  2954. raise RedisError("Cannot issue a WATCH after a MULTI")
  2955. return self.execute_command("WATCH", *names)
  2956. def unwatch(self):
  2957. if self._watching:
  2958. return self.execute_command("UNWATCH")
  2959. return True
  2960. def discard(self):
  2961. self.reset()
  2962. def delete(self, *names):
  2963. return self.execute_command("DEL", *names)
  2964. def unlink(self, *names):
  2965. return self.execute_command("UNLINK", *names)