cluster.py 88 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398
  1. import asyncio
  2. import collections
  3. import random
  4. import socket
  5. import threading
  6. import time
  7. import warnings
  8. from abc import ABC, abstractmethod
  9. from copy import copy
  10. from itertools import chain
  11. from typing import (
  12. Any,
  13. Callable,
  14. Coroutine,
  15. Deque,
  16. Dict,
  17. Generator,
  18. List,
  19. Mapping,
  20. Optional,
  21. Set,
  22. Tuple,
  23. Type,
  24. TypeVar,
  25. Union,
  26. )
  27. from redis._parsers import AsyncCommandsParser, Encoder
  28. from redis._parsers.helpers import (
  29. _RedisCallbacks,
  30. _RedisCallbacksRESP2,
  31. _RedisCallbacksRESP3,
  32. )
  33. from redis.asyncio.client import ResponseCallbackT
  34. from redis.asyncio.connection import Connection, SSLConnection, parse_url
  35. from redis.asyncio.lock import Lock
  36. from redis.asyncio.retry import Retry
  37. from redis.auth.token import TokenInterface
  38. from redis.backoff import ExponentialWithJitterBackoff, NoBackoff
  39. from redis.client import EMPTY_RESPONSE, NEVER_DECODE, AbstractRedis
  40. from redis.cluster import (
  41. PIPELINE_BLOCKED_COMMANDS,
  42. PRIMARY,
  43. REPLICA,
  44. SLOT_ID,
  45. AbstractRedisCluster,
  46. LoadBalancer,
  47. LoadBalancingStrategy,
  48. block_pipeline_command,
  49. get_node_name,
  50. parse_cluster_slots,
  51. )
  52. from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
  53. from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
  54. from redis.credentials import CredentialProvider
  55. from redis.event import AfterAsyncClusterInstantiationEvent, EventDispatcher
  56. from redis.exceptions import (
  57. AskError,
  58. BusyLoadingError,
  59. ClusterDownError,
  60. ClusterError,
  61. ConnectionError,
  62. CrossSlotTransactionError,
  63. DataError,
  64. ExecAbortError,
  65. InvalidPipelineStack,
  66. MaxConnectionsError,
  67. MovedError,
  68. RedisClusterException,
  69. RedisError,
  70. ResponseError,
  71. SlotNotCoveredError,
  72. TimeoutError,
  73. TryAgainError,
  74. WatchError,
  75. )
  76. from redis.typing import AnyKeyT, EncodableT, KeyT
  77. from redis.utils import (
  78. SSL_AVAILABLE,
  79. deprecated_args,
  80. deprecated_function,
  81. get_lib_version,
  82. safe_str,
  83. str_if_bytes,
  84. truncate_text,
  85. )
  86. if SSL_AVAILABLE:
  87. from ssl import TLSVersion, VerifyMode
  88. else:
  89. TLSVersion = None
  90. VerifyMode = None
  91. TargetNodesT = TypeVar(
  92. "TargetNodesT", str, "ClusterNode", List["ClusterNode"], Dict[Any, "ClusterNode"]
  93. )
  94. class RedisCluster(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
  95. """
  96. Create a new RedisCluster client.
  97. Pass one of parameters:
  98. - `host` & `port`
  99. - `startup_nodes`
  100. | Use ``await`` :meth:`initialize` to find cluster nodes & create connections.
  101. | Use ``await`` :meth:`close` to disconnect connections & close client.
  102. Many commands support the target_nodes kwarg. It can be one of the
  103. :attr:`NODE_FLAGS`:
  104. - :attr:`PRIMARIES`
  105. - :attr:`REPLICAS`
  106. - :attr:`ALL_NODES`
  107. - :attr:`RANDOM`
  108. - :attr:`DEFAULT_NODE`
  109. Note: This client is not thread/process/fork safe.
  110. :param host:
  111. | Can be used to point to a startup node
  112. :param port:
  113. | Port used if **host** is provided
  114. :param startup_nodes:
  115. | :class:`~.ClusterNode` to used as a startup node
  116. :param require_full_coverage:
  117. | When set to ``False``: the client will not require a full coverage of
  118. the slots. However, if not all slots are covered, and at least one node
  119. has ``cluster-require-full-coverage`` set to ``yes``, the server will throw
  120. a :class:`~.ClusterDownError` for some key-based commands.
  121. | When set to ``True``: all slots must be covered to construct the cluster
  122. client. If not all slots are covered, :class:`~.RedisClusterException` will be
  123. thrown.
  124. | See:
  125. https://redis.io/docs/manual/scaling/#redis-cluster-configuration-parameters
  126. :param read_from_replicas:
  127. | @deprecated - please use load_balancing_strategy instead
  128. | Enable read from replicas in READONLY mode.
  129. When set to true, read commands will be assigned between the primary and
  130. its replications in a Round-Robin manner.
  131. The data read from replicas is eventually consistent with the data in primary nodes.
  132. :param load_balancing_strategy:
  133. | Enable read from replicas in READONLY mode and defines the load balancing
  134. strategy that will be used for cluster node selection.
  135. The data read from replicas is eventually consistent with the data in primary nodes.
  136. :param dynamic_startup_nodes:
  137. | Set the RedisCluster's startup nodes to all the discovered nodes.
  138. If true (default value), the cluster's discovered nodes will be used to
  139. determine the cluster nodes-slots mapping in the next topology refresh.
  140. It will remove the initial passed startup nodes if their endpoints aren't
  141. listed in the CLUSTER SLOTS output.
  142. If you use dynamic DNS endpoints for startup nodes but CLUSTER SLOTS lists
  143. specific IP addresses, it is best to set it to false.
  144. :param reinitialize_steps:
  145. | Specifies the number of MOVED errors that need to occur before reinitializing
  146. the whole cluster topology. If a MOVED error occurs and the cluster does not
  147. need to be reinitialized on this current error handling, only the MOVED slot
  148. will be patched with the redirected node.
  149. To reinitialize the cluster on every MOVED error, set reinitialize_steps to 1.
  150. To avoid reinitializing the cluster on moved errors, set reinitialize_steps to
  151. 0.
  152. :param cluster_error_retry_attempts:
  153. | @deprecated - Please configure the 'retry' object instead
  154. In case 'retry' object is set - this argument is ignored!
  155. Number of times to retry before raising an error when :class:`~.TimeoutError`,
  156. :class:`~.ConnectionError`, :class:`~.SlotNotCoveredError`
  157. or :class:`~.ClusterDownError` are encountered
  158. :param retry:
  159. | A retry object that defines the retry strategy and the number of
  160. retries for the cluster client.
  161. In current implementation for the cluster client (starting form redis-py version 6.0.0)
  162. the retry object is not yet fully utilized, instead it is used just to determine
  163. the number of retries for the cluster client.
  164. In the future releases the retry object will be used to handle the cluster client retries!
  165. :param max_connections:
  166. | Maximum number of connections per node. If there are no free connections & the
  167. maximum number of connections are already created, a
  168. :class:`~.MaxConnectionsError` is raised.
  169. :param address_remap:
  170. | An optional callable which, when provided with an internal network
  171. address of a node, e.g. a `(host, port)` tuple, will return the address
  172. where the node is reachable. This can be used to map the addresses at
  173. which the nodes _think_ they are, to addresses at which a client may
  174. reach them, such as when they sit behind a proxy.
  175. | Rest of the arguments will be passed to the
  176. :class:`~redis.asyncio.connection.Connection` instances when created
  177. :raises RedisClusterException:
  178. if any arguments are invalid or unknown. Eg:
  179. - `db` != 0 or None
  180. - `path` argument for unix socket connection
  181. - none of the `host`/`port` & `startup_nodes` were provided
  182. """
  183. @classmethod
  184. def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
  185. """
  186. Return a Redis client object configured from the given URL.
  187. For example::
  188. redis://[[username]:[password]]@localhost:6379/0
  189. rediss://[[username]:[password]]@localhost:6379/0
  190. Three URL schemes are supported:
  191. - `redis://` creates a TCP socket connection. See more at:
  192. <https://www.iana.org/assignments/uri-schemes/prov/redis>
  193. - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
  194. <https://www.iana.org/assignments/uri-schemes/prov/rediss>
  195. The username, password, hostname, path and all querystring values are passed
  196. through ``urllib.parse.unquote`` in order to replace any percent-encoded values
  197. with their corresponding characters.
  198. All querystring options are cast to their appropriate Python types. Boolean
  199. arguments can be specified with string values "True"/"False" or "Yes"/"No".
  200. Values that cannot be properly cast cause a ``ValueError`` to be raised. Once
  201. parsed, the querystring arguments and keyword arguments are passed to
  202. :class:`~redis.asyncio.connection.Connection` when created.
  203. In the case of conflicting arguments, querystring arguments are used.
  204. """
  205. kwargs.update(parse_url(url))
  206. if kwargs.pop("connection_class", None) is SSLConnection:
  207. kwargs["ssl"] = True
  208. return cls(**kwargs)
  209. __slots__ = (
  210. "_initialize",
  211. "_lock",
  212. "retry",
  213. "command_flags",
  214. "commands_parser",
  215. "connection_kwargs",
  216. "encoder",
  217. "node_flags",
  218. "nodes_manager",
  219. "read_from_replicas",
  220. "reinitialize_counter",
  221. "reinitialize_steps",
  222. "response_callbacks",
  223. "result_callbacks",
  224. )
  225. @deprecated_args(
  226. args_to_warn=["read_from_replicas"],
  227. reason="Please configure the 'load_balancing_strategy' instead",
  228. version="5.3.0",
  229. )
  230. @deprecated_args(
  231. args_to_warn=[
  232. "cluster_error_retry_attempts",
  233. ],
  234. reason="Please configure the 'retry' object instead",
  235. version="6.0.0",
  236. )
  237. def __init__(
  238. self,
  239. host: Optional[str] = None,
  240. port: Union[str, int] = 6379,
  241. # Cluster related kwargs
  242. startup_nodes: Optional[List["ClusterNode"]] = None,
  243. require_full_coverage: bool = True,
  244. read_from_replicas: bool = False,
  245. load_balancing_strategy: Optional[LoadBalancingStrategy] = None,
  246. dynamic_startup_nodes: bool = True,
  247. reinitialize_steps: int = 5,
  248. cluster_error_retry_attempts: int = 3,
  249. max_connections: int = 2**31,
  250. retry: Optional["Retry"] = None,
  251. retry_on_error: Optional[List[Type[Exception]]] = None,
  252. # Client related kwargs
  253. db: Union[str, int] = 0,
  254. path: Optional[str] = None,
  255. credential_provider: Optional[CredentialProvider] = None,
  256. username: Optional[str] = None,
  257. password: Optional[str] = None,
  258. client_name: Optional[str] = None,
  259. lib_name: Optional[str] = "redis-py",
  260. lib_version: Optional[str] = get_lib_version(),
  261. # Encoding related kwargs
  262. encoding: str = "utf-8",
  263. encoding_errors: str = "strict",
  264. decode_responses: bool = False,
  265. # Connection related kwargs
  266. health_check_interval: float = 0,
  267. socket_connect_timeout: Optional[float] = None,
  268. socket_keepalive: bool = False,
  269. socket_keepalive_options: Optional[Mapping[int, Union[int, bytes]]] = None,
  270. socket_timeout: Optional[float] = None,
  271. # SSL related kwargs
  272. ssl: bool = False,
  273. ssl_ca_certs: Optional[str] = None,
  274. ssl_ca_data: Optional[str] = None,
  275. ssl_cert_reqs: Union[str, VerifyMode] = "required",
  276. ssl_certfile: Optional[str] = None,
  277. ssl_check_hostname: bool = True,
  278. ssl_keyfile: Optional[str] = None,
  279. ssl_min_version: Optional[TLSVersion] = None,
  280. ssl_ciphers: Optional[str] = None,
  281. protocol: Optional[int] = 2,
  282. address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
  283. event_dispatcher: Optional[EventDispatcher] = None,
  284. ) -> None:
  285. if db:
  286. raise RedisClusterException(
  287. "Argument 'db' must be 0 or None in cluster mode"
  288. )
  289. if path:
  290. raise RedisClusterException(
  291. "Unix domain socket is not supported in cluster mode"
  292. )
  293. if (not host or not port) and not startup_nodes:
  294. raise RedisClusterException(
  295. "RedisCluster requires at least one node to discover the cluster.\n"
  296. "Please provide one of the following or use RedisCluster.from_url:\n"
  297. ' - host and port: RedisCluster(host="localhost", port=6379)\n'
  298. " - startup_nodes: RedisCluster(startup_nodes=["
  299. 'ClusterNode("localhost", 6379), ClusterNode("localhost", 6380)])'
  300. )
  301. kwargs: Dict[str, Any] = {
  302. "max_connections": max_connections,
  303. "connection_class": Connection,
  304. # Client related kwargs
  305. "credential_provider": credential_provider,
  306. "username": username,
  307. "password": password,
  308. "client_name": client_name,
  309. "lib_name": lib_name,
  310. "lib_version": lib_version,
  311. # Encoding related kwargs
  312. "encoding": encoding,
  313. "encoding_errors": encoding_errors,
  314. "decode_responses": decode_responses,
  315. # Connection related kwargs
  316. "health_check_interval": health_check_interval,
  317. "socket_connect_timeout": socket_connect_timeout,
  318. "socket_keepalive": socket_keepalive,
  319. "socket_keepalive_options": socket_keepalive_options,
  320. "socket_timeout": socket_timeout,
  321. "protocol": protocol,
  322. }
  323. if ssl:
  324. # SSL related kwargs
  325. kwargs.update(
  326. {
  327. "connection_class": SSLConnection,
  328. "ssl_ca_certs": ssl_ca_certs,
  329. "ssl_ca_data": ssl_ca_data,
  330. "ssl_cert_reqs": ssl_cert_reqs,
  331. "ssl_certfile": ssl_certfile,
  332. "ssl_check_hostname": ssl_check_hostname,
  333. "ssl_keyfile": ssl_keyfile,
  334. "ssl_min_version": ssl_min_version,
  335. "ssl_ciphers": ssl_ciphers,
  336. }
  337. )
  338. if read_from_replicas or load_balancing_strategy:
  339. # Call our on_connect function to configure READONLY mode
  340. kwargs["redis_connect_func"] = self.on_connect
  341. if retry:
  342. self.retry = retry
  343. else:
  344. self.retry = Retry(
  345. backoff=ExponentialWithJitterBackoff(base=1, cap=10),
  346. retries=cluster_error_retry_attempts,
  347. )
  348. if retry_on_error:
  349. self.retry.update_supported_errors(retry_on_error)
  350. kwargs["response_callbacks"] = _RedisCallbacks.copy()
  351. if kwargs.get("protocol") in ["3", 3]:
  352. kwargs["response_callbacks"].update(_RedisCallbacksRESP3)
  353. else:
  354. kwargs["response_callbacks"].update(_RedisCallbacksRESP2)
  355. self.connection_kwargs = kwargs
  356. if startup_nodes:
  357. passed_nodes = []
  358. for node in startup_nodes:
  359. passed_nodes.append(
  360. ClusterNode(node.host, node.port, **self.connection_kwargs)
  361. )
  362. startup_nodes = passed_nodes
  363. else:
  364. startup_nodes = []
  365. if host and port:
  366. startup_nodes.append(ClusterNode(host, port, **self.connection_kwargs))
  367. if event_dispatcher is None:
  368. self._event_dispatcher = EventDispatcher()
  369. else:
  370. self._event_dispatcher = event_dispatcher
  371. self.nodes_manager = NodesManager(
  372. startup_nodes,
  373. require_full_coverage,
  374. kwargs,
  375. dynamic_startup_nodes=dynamic_startup_nodes,
  376. address_remap=address_remap,
  377. event_dispatcher=self._event_dispatcher,
  378. )
  379. self.encoder = Encoder(encoding, encoding_errors, decode_responses)
  380. self.read_from_replicas = read_from_replicas
  381. self.load_balancing_strategy = load_balancing_strategy
  382. self.reinitialize_steps = reinitialize_steps
  383. self.reinitialize_counter = 0
  384. self.commands_parser = AsyncCommandsParser()
  385. self.node_flags = self.__class__.NODE_FLAGS.copy()
  386. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  387. self.response_callbacks = kwargs["response_callbacks"]
  388. self.result_callbacks = self.__class__.RESULT_CALLBACKS.copy()
  389. self.result_callbacks["CLUSTER SLOTS"] = (
  390. lambda cmd, res, **kwargs: parse_cluster_slots(
  391. list(res.values())[0], **kwargs
  392. )
  393. )
  394. self._initialize = True
  395. self._lock: Optional[asyncio.Lock] = None
  396. async def initialize(self) -> "RedisCluster":
  397. """Get all nodes from startup nodes & creates connections if not initialized."""
  398. if self._initialize:
  399. if not self._lock:
  400. self._lock = asyncio.Lock()
  401. async with self._lock:
  402. if self._initialize:
  403. try:
  404. await self.nodes_manager.initialize()
  405. await self.commands_parser.initialize(
  406. self.nodes_manager.default_node
  407. )
  408. self._initialize = False
  409. except BaseException:
  410. await self.nodes_manager.aclose()
  411. await self.nodes_manager.aclose("startup_nodes")
  412. raise
  413. return self
  414. async def aclose(self) -> None:
  415. """Close all connections & client if initialized."""
  416. if not self._initialize:
  417. if not self._lock:
  418. self._lock = asyncio.Lock()
  419. async with self._lock:
  420. if not self._initialize:
  421. self._initialize = True
  422. await self.nodes_manager.aclose()
  423. await self.nodes_manager.aclose("startup_nodes")
  424. @deprecated_function(version="5.0.0", reason="Use aclose() instead", name="close")
  425. async def close(self) -> None:
  426. """alias for aclose() for backwards compatibility"""
  427. await self.aclose()
  428. async def __aenter__(self) -> "RedisCluster":
  429. return await self.initialize()
  430. async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  431. await self.aclose()
  432. def __await__(self) -> Generator[Any, None, "RedisCluster"]:
  433. return self.initialize().__await__()
  434. _DEL_MESSAGE = "Unclosed RedisCluster client"
  435. def __del__(
  436. self,
  437. _warn: Any = warnings.warn,
  438. _grl: Any = asyncio.get_running_loop,
  439. ) -> None:
  440. if hasattr(self, "_initialize") and not self._initialize:
  441. _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
  442. try:
  443. context = {"client": self, "message": self._DEL_MESSAGE}
  444. _grl().call_exception_handler(context)
  445. except RuntimeError:
  446. pass
  447. async def on_connect(self, connection: Connection) -> None:
  448. await connection.on_connect()
  449. # Sending READONLY command to server to configure connection as
  450. # readonly. Since each cluster node may change its server type due
  451. # to a failover, we should establish a READONLY connection
  452. # regardless of the server type. If this is a primary connection,
  453. # READONLY would not affect executing write commands.
  454. await connection.send_command("READONLY")
  455. if str_if_bytes(await connection.read_response()) != "OK":
  456. raise ConnectionError("READONLY command failed")
  457. def get_nodes(self) -> List["ClusterNode"]:
  458. """Get all nodes of the cluster."""
  459. return list(self.nodes_manager.nodes_cache.values())
  460. def get_primaries(self) -> List["ClusterNode"]:
  461. """Get the primary nodes of the cluster."""
  462. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  463. def get_replicas(self) -> List["ClusterNode"]:
  464. """Get the replica nodes of the cluster."""
  465. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  466. def get_random_node(self) -> "ClusterNode":
  467. """Get a random node of the cluster."""
  468. return random.choice(list(self.nodes_manager.nodes_cache.values()))
  469. def get_default_node(self) -> "ClusterNode":
  470. """Get the default node of the client."""
  471. return self.nodes_manager.default_node
  472. def set_default_node(self, node: "ClusterNode") -> None:
  473. """
  474. Set the default node of the client.
  475. :raises DataError: if None is passed or node does not exist in cluster.
  476. """
  477. if not node or not self.get_node(node_name=node.name):
  478. raise DataError("The requested node does not exist in the cluster.")
  479. self.nodes_manager.default_node = node
  480. def get_node(
  481. self,
  482. host: Optional[str] = None,
  483. port: Optional[int] = None,
  484. node_name: Optional[str] = None,
  485. ) -> Optional["ClusterNode"]:
  486. """Get node by (host, port) or node_name."""
  487. return self.nodes_manager.get_node(host, port, node_name)
  488. def get_node_from_key(
  489. self, key: str, replica: bool = False
  490. ) -> Optional["ClusterNode"]:
  491. """
  492. Get the cluster node corresponding to the provided key.
  493. :param key:
  494. :param replica:
  495. | Indicates if a replica should be returned
  496. |
  497. None will returned if no replica holds this key
  498. :raises SlotNotCoveredError: if the key is not covered by any slot.
  499. """
  500. slot = self.keyslot(key)
  501. slot_cache = self.nodes_manager.slots_cache.get(slot)
  502. if not slot_cache:
  503. raise SlotNotCoveredError(f'Slot "{slot}" is not covered by the cluster.')
  504. if replica:
  505. if len(self.nodes_manager.slots_cache[slot]) < 2:
  506. return None
  507. node_idx = 1
  508. else:
  509. node_idx = 0
  510. return slot_cache[node_idx]
  511. def keyslot(self, key: EncodableT) -> int:
  512. """
  513. Find the keyslot for a given key.
  514. See: https://redis.io/docs/manual/scaling/#redis-cluster-data-sharding
  515. """
  516. return key_slot(self.encoder.encode(key))
  517. def get_encoder(self) -> Encoder:
  518. """Get the encoder object of the client."""
  519. return self.encoder
  520. def get_connection_kwargs(self) -> Dict[str, Optional[Any]]:
  521. """Get the kwargs passed to :class:`~redis.asyncio.connection.Connection`."""
  522. return self.connection_kwargs
  523. def set_retry(self, retry: Retry) -> None:
  524. self.retry = retry
  525. def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
  526. """Set a custom response callback."""
  527. self.response_callbacks[command] = callback
  528. async def _determine_nodes(
  529. self, command: str, *args: Any, node_flag: Optional[str] = None
  530. ) -> List["ClusterNode"]:
  531. # Determine which nodes should be executed the command on.
  532. # Returns a list of target nodes.
  533. if not node_flag:
  534. # get the nodes group for this command if it was predefined
  535. node_flag = self.command_flags.get(command)
  536. if node_flag in self.node_flags:
  537. if node_flag == self.__class__.DEFAULT_NODE:
  538. # return the cluster's default node
  539. return [self.nodes_manager.default_node]
  540. if node_flag == self.__class__.PRIMARIES:
  541. # return all primaries
  542. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  543. if node_flag == self.__class__.REPLICAS:
  544. # return all replicas
  545. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  546. if node_flag == self.__class__.ALL_NODES:
  547. # return all nodes
  548. return list(self.nodes_manager.nodes_cache.values())
  549. if node_flag == self.__class__.RANDOM:
  550. # return a random node
  551. return [random.choice(list(self.nodes_manager.nodes_cache.values()))]
  552. # get the node that holds the key's slot
  553. return [
  554. self.nodes_manager.get_node_from_slot(
  555. await self._determine_slot(command, *args),
  556. self.read_from_replicas and command in READ_COMMANDS,
  557. self.load_balancing_strategy if command in READ_COMMANDS else None,
  558. )
  559. ]
  560. async def _determine_slot(self, command: str, *args: Any) -> int:
  561. if self.command_flags.get(command) == SLOT_ID:
  562. # The command contains the slot ID
  563. return int(args[0])
  564. # Get the keys in the command
  565. # EVAL and EVALSHA are common enough that it's wasteful to go to the
  566. # redis server to parse the keys. Besides, there is a bug in redis<7.0
  567. # where `self._get_command_keys()` fails anyway. So, we special case
  568. # EVAL/EVALSHA.
  569. # - issue: https://github.com/redis/redis/issues/9493
  570. # - fix: https://github.com/redis/redis/pull/9733
  571. if command.upper() in ("EVAL", "EVALSHA"):
  572. # command syntax: EVAL "script body" num_keys ...
  573. if len(args) < 2:
  574. raise RedisClusterException(
  575. f"Invalid args in command: {command, *args}"
  576. )
  577. keys = args[2 : 2 + int(args[1])]
  578. # if there are 0 keys, that means the script can be run on any node
  579. # so we can just return a random slot
  580. if not keys:
  581. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  582. else:
  583. keys = await self.commands_parser.get_keys(command, *args)
  584. if not keys:
  585. # FCALL can call a function with 0 keys, that means the function
  586. # can be run on any node so we can just return a random slot
  587. if command.upper() in ("FCALL", "FCALL_RO"):
  588. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  589. raise RedisClusterException(
  590. "No way to dispatch this command to Redis Cluster. "
  591. "Missing key.\nYou can execute the command by specifying "
  592. f"target nodes.\nCommand: {args}"
  593. )
  594. # single key command
  595. if len(keys) == 1:
  596. return self.keyslot(keys[0])
  597. # multi-key command; we need to make sure all keys are mapped to
  598. # the same slot
  599. slots = {self.keyslot(key) for key in keys}
  600. if len(slots) != 1:
  601. raise RedisClusterException(
  602. f"{command} - all keys must map to the same key slot"
  603. )
  604. return slots.pop()
  605. def _is_node_flag(self, target_nodes: Any) -> bool:
  606. return isinstance(target_nodes, str) and target_nodes in self.node_flags
  607. def _parse_target_nodes(self, target_nodes: Any) -> List["ClusterNode"]:
  608. if isinstance(target_nodes, list):
  609. nodes = target_nodes
  610. elif isinstance(target_nodes, ClusterNode):
  611. # Supports passing a single ClusterNode as a variable
  612. nodes = [target_nodes]
  613. elif isinstance(target_nodes, dict):
  614. # Supports dictionaries of the format {node_name: node}.
  615. # It enables to execute commands with multi nodes as follows:
  616. # rc.cluster_save_config(rc.get_primaries())
  617. nodes = list(target_nodes.values())
  618. else:
  619. raise TypeError(
  620. "target_nodes type can be one of the following: "
  621. "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
  622. "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
  623. f"The passed type is {type(target_nodes)}"
  624. )
  625. return nodes
  626. async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  627. """
  628. Execute a raw command on the appropriate cluster node or target_nodes.
  629. It will retry the command as specified by the retries property of
  630. the :attr:`retry` & then raise an exception.
  631. :param args:
  632. | Raw command args
  633. :param kwargs:
  634. - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
  635. or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
  636. - Rest of the kwargs are passed to the Redis connection
  637. :raises RedisClusterException: if target_nodes is not provided & the command
  638. can't be mapped to a slot
  639. """
  640. command = args[0]
  641. target_nodes = []
  642. target_nodes_specified = False
  643. retry_attempts = self.retry.get_retries()
  644. passed_targets = kwargs.pop("target_nodes", None)
  645. if passed_targets and not self._is_node_flag(passed_targets):
  646. target_nodes = self._parse_target_nodes(passed_targets)
  647. target_nodes_specified = True
  648. retry_attempts = 0
  649. # Add one for the first execution
  650. execute_attempts = 1 + retry_attempts
  651. for _ in range(execute_attempts):
  652. if self._initialize:
  653. await self.initialize()
  654. if (
  655. len(target_nodes) == 1
  656. and target_nodes[0] == self.get_default_node()
  657. ):
  658. # Replace the default cluster node
  659. self.replace_default_node()
  660. try:
  661. if not target_nodes_specified:
  662. # Determine the nodes to execute the command on
  663. target_nodes = await self._determine_nodes(
  664. *args, node_flag=passed_targets
  665. )
  666. if not target_nodes:
  667. raise RedisClusterException(
  668. f"No targets were found to execute {args} command on"
  669. )
  670. if len(target_nodes) == 1:
  671. # Return the processed result
  672. ret = await self._execute_command(target_nodes[0], *args, **kwargs)
  673. if command in self.result_callbacks:
  674. return self.result_callbacks[command](
  675. command, {target_nodes[0].name: ret}, **kwargs
  676. )
  677. return ret
  678. else:
  679. keys = [node.name for node in target_nodes]
  680. values = await asyncio.gather(
  681. *(
  682. asyncio.create_task(
  683. self._execute_command(node, *args, **kwargs)
  684. )
  685. for node in target_nodes
  686. )
  687. )
  688. if command in self.result_callbacks:
  689. return self.result_callbacks[command](
  690. command, dict(zip(keys, values)), **kwargs
  691. )
  692. return dict(zip(keys, values))
  693. except Exception as e:
  694. if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  695. # The nodes and slots cache were should be reinitialized.
  696. # Try again with the new cluster setup.
  697. retry_attempts -= 1
  698. continue
  699. else:
  700. # raise the exception
  701. raise e
  702. async def _execute_command(
  703. self, target_node: "ClusterNode", *args: Union[KeyT, EncodableT], **kwargs: Any
  704. ) -> Any:
  705. asking = moved = False
  706. redirect_addr = None
  707. ttl = self.RedisClusterRequestTTL
  708. while ttl > 0:
  709. ttl -= 1
  710. try:
  711. if asking:
  712. target_node = self.get_node(node_name=redirect_addr)
  713. await target_node.execute_command("ASKING")
  714. asking = False
  715. elif moved:
  716. # MOVED occurred and the slots cache was updated,
  717. # refresh the target node
  718. slot = await self._determine_slot(*args)
  719. target_node = self.nodes_manager.get_node_from_slot(
  720. slot,
  721. self.read_from_replicas and args[0] in READ_COMMANDS,
  722. self.load_balancing_strategy
  723. if args[0] in READ_COMMANDS
  724. else None,
  725. )
  726. moved = False
  727. return await target_node.execute_command(*args, **kwargs)
  728. except (BusyLoadingError, MaxConnectionsError):
  729. raise
  730. except (ConnectionError, TimeoutError):
  731. # Connection retries are being handled in the node's
  732. # Retry object.
  733. # Remove the failed node from the startup nodes before we try
  734. # to reinitialize the cluster
  735. self.nodes_manager.startup_nodes.pop(target_node.name, None)
  736. # Hard force of reinitialize of the node/slots setup
  737. # and try again with the new setup
  738. await self.aclose()
  739. raise
  740. except (ClusterDownError, SlotNotCoveredError):
  741. # ClusterDownError can occur during a failover and to get
  742. # self-healed, we will try to reinitialize the cluster layout
  743. # and retry executing the command
  744. # SlotNotCoveredError can occur when the cluster is not fully
  745. # initialized or can be temporary issue.
  746. # We will try to reinitialize the cluster topology
  747. # and retry executing the command
  748. await self.aclose()
  749. await asyncio.sleep(0.25)
  750. raise
  751. except MovedError as e:
  752. # First, we will try to patch the slots/nodes cache with the
  753. # redirected node output and try again. If MovedError exceeds
  754. # 'reinitialize_steps' number of times, we will force
  755. # reinitializing the tables, and then try again.
  756. # 'reinitialize_steps' counter will increase faster when
  757. # the same client object is shared between multiple threads. To
  758. # reduce the frequency you can set this variable in the
  759. # RedisCluster constructor.
  760. self.reinitialize_counter += 1
  761. if (
  762. self.reinitialize_steps
  763. and self.reinitialize_counter % self.reinitialize_steps == 0
  764. ):
  765. await self.aclose()
  766. # Reset the counter
  767. self.reinitialize_counter = 0
  768. else:
  769. self.nodes_manager._moved_exception = e
  770. moved = True
  771. except AskError as e:
  772. redirect_addr = get_node_name(host=e.host, port=e.port)
  773. asking = True
  774. except TryAgainError:
  775. if ttl < self.RedisClusterRequestTTL / 2:
  776. await asyncio.sleep(0.05)
  777. raise ClusterError("TTL exhausted.")
  778. def pipeline(
  779. self, transaction: Optional[Any] = None, shard_hint: Optional[Any] = None
  780. ) -> "ClusterPipeline":
  781. """
  782. Create & return a new :class:`~.ClusterPipeline` object.
  783. Cluster implementation of pipeline does not support transaction or shard_hint.
  784. :raises RedisClusterException: if transaction or shard_hint are truthy values
  785. """
  786. if shard_hint:
  787. raise RedisClusterException("shard_hint is deprecated in cluster mode")
  788. return ClusterPipeline(self, transaction)
  789. def lock(
  790. self,
  791. name: KeyT,
  792. timeout: Optional[float] = None,
  793. sleep: float = 0.1,
  794. blocking: bool = True,
  795. blocking_timeout: Optional[float] = None,
  796. lock_class: Optional[Type[Lock]] = None,
  797. thread_local: bool = True,
  798. raise_on_release_error: bool = True,
  799. ) -> Lock:
  800. """
  801. Return a new Lock object using key ``name`` that mimics
  802. the behavior of threading.Lock.
  803. If specified, ``timeout`` indicates a maximum life for the lock.
  804. By default, it will remain locked until release() is called.
  805. ``sleep`` indicates the amount of time to sleep per loop iteration
  806. when the lock is in blocking mode and another client is currently
  807. holding the lock.
  808. ``blocking`` indicates whether calling ``acquire`` should block until
  809. the lock has been acquired or to fail immediately, causing ``acquire``
  810. to return False and the lock not being acquired. Defaults to True.
  811. Note this value can be overridden by passing a ``blocking``
  812. argument to ``acquire``.
  813. ``blocking_timeout`` indicates the maximum amount of time in seconds to
  814. spend trying to acquire the lock. A value of ``None`` indicates
  815. continue trying forever. ``blocking_timeout`` can be specified as a
  816. float or integer, both representing the number of seconds to wait.
  817. ``lock_class`` forces the specified lock implementation. Note that as
  818. of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
  819. a Lua-based lock). So, it's unlikely you'll need this parameter, unless
  820. you have created your own custom lock class.
  821. ``thread_local`` indicates whether the lock token is placed in
  822. thread-local storage. By default, the token is placed in thread local
  823. storage so that a thread only sees its token, not a token set by
  824. another thread. Consider the following timeline:
  825. time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
  826. thread-1 sets the token to "abc"
  827. time: 1, thread-2 blocks trying to acquire `my-lock` using the
  828. Lock instance.
  829. time: 5, thread-1 has not yet completed. redis expires the lock
  830. key.
  831. time: 5, thread-2 acquired `my-lock` now that it's available.
  832. thread-2 sets the token to "xyz"
  833. time: 6, thread-1 finishes its work and calls release(). if the
  834. token is *not* stored in thread local storage, then
  835. thread-1 would see the token value as "xyz" and would be
  836. able to successfully release the thread-2's lock.
  837. ``raise_on_release_error`` indicates whether to raise an exception when
  838. the lock is no longer owned when exiting the context manager. By default,
  839. this is True, meaning an exception will be raised. If False, the warning
  840. will be logged and the exception will be suppressed.
  841. In some use cases it's necessary to disable thread local storage. For
  842. example, if you have code where one thread acquires a lock and passes
  843. that lock instance to a worker thread to release later. If thread
  844. local storage isn't disabled in this case, the worker thread won't see
  845. the token set by the thread that acquired the lock. Our assumption
  846. is that these cases aren't common and as such default to using
  847. thread local storage."""
  848. if lock_class is None:
  849. lock_class = Lock
  850. return lock_class(
  851. self,
  852. name,
  853. timeout=timeout,
  854. sleep=sleep,
  855. blocking=blocking,
  856. blocking_timeout=blocking_timeout,
  857. thread_local=thread_local,
  858. raise_on_release_error=raise_on_release_error,
  859. )
  860. async def transaction(
  861. self, func: Coroutine[None, "ClusterPipeline", Any], *watches, **kwargs
  862. ):
  863. """
  864. Convenience method for executing the callable `func` as a transaction
  865. while watching all keys specified in `watches`. The 'func' callable
  866. should expect a single argument which is a Pipeline object.
  867. """
  868. shard_hint = kwargs.pop("shard_hint", None)
  869. value_from_callable = kwargs.pop("value_from_callable", False)
  870. watch_delay = kwargs.pop("watch_delay", None)
  871. async with self.pipeline(True, shard_hint) as pipe:
  872. while True:
  873. try:
  874. if watches:
  875. await pipe.watch(*watches)
  876. func_value = await func(pipe)
  877. exec_value = await pipe.execute()
  878. return func_value if value_from_callable else exec_value
  879. except WatchError:
  880. if watch_delay is not None and watch_delay > 0:
  881. time.sleep(watch_delay)
  882. continue
  883. class ClusterNode:
  884. """
  885. Create a new ClusterNode.
  886. Each ClusterNode manages multiple :class:`~redis.asyncio.connection.Connection`
  887. objects for the (host, port).
  888. """
  889. __slots__ = (
  890. "_connections",
  891. "_free",
  892. "_lock",
  893. "_event_dispatcher",
  894. "connection_class",
  895. "connection_kwargs",
  896. "host",
  897. "max_connections",
  898. "name",
  899. "port",
  900. "response_callbacks",
  901. "server_type",
  902. )
  903. def __init__(
  904. self,
  905. host: str,
  906. port: Union[str, int],
  907. server_type: Optional[str] = None,
  908. *,
  909. max_connections: int = 2**31,
  910. connection_class: Type[Connection] = Connection,
  911. **connection_kwargs: Any,
  912. ) -> None:
  913. if host == "localhost":
  914. host = socket.gethostbyname(host)
  915. connection_kwargs["host"] = host
  916. connection_kwargs["port"] = port
  917. self.host = host
  918. self.port = port
  919. self.name = get_node_name(host, port)
  920. self.server_type = server_type
  921. self.max_connections = max_connections
  922. self.connection_class = connection_class
  923. self.connection_kwargs = connection_kwargs
  924. self.response_callbacks = connection_kwargs.pop("response_callbacks", {})
  925. self._connections: List[Connection] = []
  926. self._free: Deque[Connection] = collections.deque(maxlen=self.max_connections)
  927. self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
  928. if self._event_dispatcher is None:
  929. self._event_dispatcher = EventDispatcher()
  930. def __repr__(self) -> str:
  931. return (
  932. f"[host={self.host}, port={self.port}, "
  933. f"name={self.name}, server_type={self.server_type}]"
  934. )
  935. def __eq__(self, obj: Any) -> bool:
  936. return isinstance(obj, ClusterNode) and obj.name == self.name
  937. _DEL_MESSAGE = "Unclosed ClusterNode object"
  938. def __del__(
  939. self,
  940. _warn: Any = warnings.warn,
  941. _grl: Any = asyncio.get_running_loop,
  942. ) -> None:
  943. for connection in self._connections:
  944. if connection.is_connected:
  945. _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
  946. try:
  947. context = {"client": self, "message": self._DEL_MESSAGE}
  948. _grl().call_exception_handler(context)
  949. except RuntimeError:
  950. pass
  951. break
  952. async def disconnect(self) -> None:
  953. ret = await asyncio.gather(
  954. *(
  955. asyncio.create_task(connection.disconnect())
  956. for connection in self._connections
  957. ),
  958. return_exceptions=True,
  959. )
  960. exc = next((res for res in ret if isinstance(res, Exception)), None)
  961. if exc:
  962. raise exc
  963. def acquire_connection(self) -> Connection:
  964. try:
  965. return self._free.popleft()
  966. except IndexError:
  967. if len(self._connections) < self.max_connections:
  968. # We are configuring the connection pool not to retry
  969. # connections on lower level clients to avoid retrying
  970. # connections to nodes that are not reachable
  971. # and to avoid blocking the connection pool.
  972. # The only error that will have some handling in the lower
  973. # level clients is ConnectionError which will trigger disconnection
  974. # of the socket.
  975. # The retries will be handled on cluster client level
  976. # where we will have proper handling of the cluster topology
  977. retry = Retry(
  978. backoff=NoBackoff(),
  979. retries=0,
  980. supported_errors=(ConnectionError,),
  981. )
  982. connection_kwargs = self.connection_kwargs.copy()
  983. connection_kwargs["retry"] = retry
  984. connection = self.connection_class(**connection_kwargs)
  985. self._connections.append(connection)
  986. return connection
  987. raise MaxConnectionsError()
  988. def release(self, connection: Connection) -> None:
  989. """
  990. Release connection back to free queue.
  991. """
  992. self._free.append(connection)
  993. async def parse_response(
  994. self, connection: Connection, command: str, **kwargs: Any
  995. ) -> Any:
  996. try:
  997. if NEVER_DECODE in kwargs:
  998. response = await connection.read_response(disable_decoding=True)
  999. kwargs.pop(NEVER_DECODE)
  1000. else:
  1001. response = await connection.read_response()
  1002. except ResponseError:
  1003. if EMPTY_RESPONSE in kwargs:
  1004. return kwargs[EMPTY_RESPONSE]
  1005. raise
  1006. if EMPTY_RESPONSE in kwargs:
  1007. kwargs.pop(EMPTY_RESPONSE)
  1008. # Remove keys entry, it needs only for cache.
  1009. kwargs.pop("keys", None)
  1010. # Return response
  1011. if command in self.response_callbacks:
  1012. return self.response_callbacks[command](response, **kwargs)
  1013. return response
  1014. async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
  1015. # Acquire connection
  1016. connection = self.acquire_connection()
  1017. # Execute command
  1018. await connection.send_packed_command(connection.pack_command(*args), False)
  1019. # Read response
  1020. try:
  1021. return await self.parse_response(connection, args[0], **kwargs)
  1022. finally:
  1023. # Release connection
  1024. self._free.append(connection)
  1025. async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
  1026. # Acquire connection
  1027. connection = self.acquire_connection()
  1028. # Execute command
  1029. await connection.send_packed_command(
  1030. connection.pack_commands(cmd.args for cmd in commands), False
  1031. )
  1032. # Read responses
  1033. ret = False
  1034. for cmd in commands:
  1035. try:
  1036. cmd.result = await self.parse_response(
  1037. connection, cmd.args[0], **cmd.kwargs
  1038. )
  1039. except Exception as e:
  1040. cmd.result = e
  1041. ret = True
  1042. # Release connection
  1043. self._free.append(connection)
  1044. return ret
  1045. async def re_auth_callback(self, token: TokenInterface):
  1046. tmp_queue = collections.deque()
  1047. while self._free:
  1048. conn = self._free.popleft()
  1049. await conn.retry.call_with_retry(
  1050. lambda: conn.send_command(
  1051. "AUTH", token.try_get("oid"), token.get_value()
  1052. ),
  1053. lambda error: self._mock(error),
  1054. )
  1055. await conn.retry.call_with_retry(
  1056. lambda: conn.read_response(), lambda error: self._mock(error)
  1057. )
  1058. tmp_queue.append(conn)
  1059. while tmp_queue:
  1060. conn = tmp_queue.popleft()
  1061. self._free.append(conn)
  1062. async def _mock(self, error: RedisError):
  1063. """
  1064. Dummy functions, needs to be passed as error callback to retry object.
  1065. :param error:
  1066. :return:
  1067. """
  1068. pass
  1069. class NodesManager:
  1070. __slots__ = (
  1071. "_dynamic_startup_nodes",
  1072. "_moved_exception",
  1073. "_event_dispatcher",
  1074. "connection_kwargs",
  1075. "default_node",
  1076. "nodes_cache",
  1077. "read_load_balancer",
  1078. "require_full_coverage",
  1079. "slots_cache",
  1080. "startup_nodes",
  1081. "address_remap",
  1082. )
  1083. def __init__(
  1084. self,
  1085. startup_nodes: List["ClusterNode"],
  1086. require_full_coverage: bool,
  1087. connection_kwargs: Dict[str, Any],
  1088. dynamic_startup_nodes: bool = True,
  1089. address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
  1090. event_dispatcher: Optional[EventDispatcher] = None,
  1091. ) -> None:
  1092. self.startup_nodes = {node.name: node for node in startup_nodes}
  1093. self.require_full_coverage = require_full_coverage
  1094. self.connection_kwargs = connection_kwargs
  1095. self.address_remap = address_remap
  1096. self.default_node: "ClusterNode" = None
  1097. self.nodes_cache: Dict[str, "ClusterNode"] = {}
  1098. self.slots_cache: Dict[int, List["ClusterNode"]] = {}
  1099. self.read_load_balancer = LoadBalancer()
  1100. self._dynamic_startup_nodes: bool = dynamic_startup_nodes
  1101. self._moved_exception: MovedError = None
  1102. if event_dispatcher is None:
  1103. self._event_dispatcher = EventDispatcher()
  1104. else:
  1105. self._event_dispatcher = event_dispatcher
  1106. def get_node(
  1107. self,
  1108. host: Optional[str] = None,
  1109. port: Optional[int] = None,
  1110. node_name: Optional[str] = None,
  1111. ) -> Optional["ClusterNode"]:
  1112. if host and port:
  1113. # the user passed host and port
  1114. if host == "localhost":
  1115. host = socket.gethostbyname(host)
  1116. return self.nodes_cache.get(get_node_name(host=host, port=port))
  1117. elif node_name:
  1118. return self.nodes_cache.get(node_name)
  1119. else:
  1120. raise DataError(
  1121. "get_node requires one of the following: 1. node name 2. host and port"
  1122. )
  1123. def set_nodes(
  1124. self,
  1125. old: Dict[str, "ClusterNode"],
  1126. new: Dict[str, "ClusterNode"],
  1127. remove_old: bool = False,
  1128. ) -> None:
  1129. if remove_old:
  1130. for name in list(old.keys()):
  1131. if name not in new:
  1132. task = asyncio.create_task(old.pop(name).disconnect()) # noqa
  1133. for name, node in new.items():
  1134. if name in old:
  1135. if old[name] is node:
  1136. continue
  1137. task = asyncio.create_task(old[name].disconnect()) # noqa
  1138. old[name] = node
  1139. def update_moved_exception(self, exception):
  1140. self._moved_exception = exception
  1141. def _update_moved_slots(self) -> None:
  1142. e = self._moved_exception
  1143. redirected_node = self.get_node(host=e.host, port=e.port)
  1144. if redirected_node:
  1145. # The node already exists
  1146. if redirected_node.server_type != PRIMARY:
  1147. # Update the node's server type
  1148. redirected_node.server_type = PRIMARY
  1149. else:
  1150. # This is a new node, we will add it to the nodes cache
  1151. redirected_node = ClusterNode(
  1152. e.host, e.port, PRIMARY, **self.connection_kwargs
  1153. )
  1154. self.set_nodes(self.nodes_cache, {redirected_node.name: redirected_node})
  1155. if redirected_node in self.slots_cache[e.slot_id]:
  1156. # The MOVED error resulted from a failover, and the new slot owner
  1157. # had previously been a replica.
  1158. old_primary = self.slots_cache[e.slot_id][0]
  1159. # Update the old primary to be a replica and add it to the end of
  1160. # the slot's node list
  1161. old_primary.server_type = REPLICA
  1162. self.slots_cache[e.slot_id].append(old_primary)
  1163. # Remove the old replica, which is now a primary, from the slot's
  1164. # node list
  1165. self.slots_cache[e.slot_id].remove(redirected_node)
  1166. # Override the old primary with the new one
  1167. self.slots_cache[e.slot_id][0] = redirected_node
  1168. if self.default_node == old_primary:
  1169. # Update the default node with the new primary
  1170. self.default_node = redirected_node
  1171. else:
  1172. # The new slot owner is a new server, or a server from a different
  1173. # shard. We need to remove all current nodes from the slot's list
  1174. # (including replications) and add just the new node.
  1175. self.slots_cache[e.slot_id] = [redirected_node]
  1176. # Reset moved_exception
  1177. self._moved_exception = None
  1178. def get_node_from_slot(
  1179. self,
  1180. slot: int,
  1181. read_from_replicas: bool = False,
  1182. load_balancing_strategy=None,
  1183. ) -> "ClusterNode":
  1184. if self._moved_exception:
  1185. self._update_moved_slots()
  1186. if read_from_replicas is True and load_balancing_strategy is None:
  1187. load_balancing_strategy = LoadBalancingStrategy.ROUND_ROBIN
  1188. try:
  1189. if len(self.slots_cache[slot]) > 1 and load_balancing_strategy:
  1190. # get the server index using the strategy defined in load_balancing_strategy
  1191. primary_name = self.slots_cache[slot][0].name
  1192. node_idx = self.read_load_balancer.get_server_index(
  1193. primary_name, len(self.slots_cache[slot]), load_balancing_strategy
  1194. )
  1195. return self.slots_cache[slot][node_idx]
  1196. return self.slots_cache[slot][0]
  1197. except (IndexError, TypeError):
  1198. raise SlotNotCoveredError(
  1199. f'Slot "{slot}" not covered by the cluster. '
  1200. f'"require_full_coverage={self.require_full_coverage}"'
  1201. )
  1202. def get_nodes_by_server_type(self, server_type: str) -> List["ClusterNode"]:
  1203. return [
  1204. node
  1205. for node in self.nodes_cache.values()
  1206. if node.server_type == server_type
  1207. ]
  1208. async def initialize(self) -> None:
  1209. self.read_load_balancer.reset()
  1210. tmp_nodes_cache: Dict[str, "ClusterNode"] = {}
  1211. tmp_slots: Dict[int, List["ClusterNode"]] = {}
  1212. disagreements = []
  1213. startup_nodes_reachable = False
  1214. fully_covered = False
  1215. exception = None
  1216. # Convert to tuple to prevent RuntimeError if self.startup_nodes
  1217. # is modified during iteration
  1218. for startup_node in tuple(self.startup_nodes.values()):
  1219. try:
  1220. # Make sure cluster mode is enabled on this node
  1221. try:
  1222. self._event_dispatcher.dispatch(
  1223. AfterAsyncClusterInstantiationEvent(
  1224. self.nodes_cache,
  1225. self.connection_kwargs.get("credential_provider", None),
  1226. )
  1227. )
  1228. cluster_slots = await startup_node.execute_command("CLUSTER SLOTS")
  1229. except ResponseError:
  1230. raise RedisClusterException(
  1231. "Cluster mode is not enabled on this node"
  1232. )
  1233. startup_nodes_reachable = True
  1234. except Exception as e:
  1235. # Try the next startup node.
  1236. # The exception is saved and raised only if we have no more nodes.
  1237. exception = e
  1238. continue
  1239. # CLUSTER SLOTS command results in the following output:
  1240. # [[slot_section[from_slot,to_slot,master,replica1,...,replicaN]]]
  1241. # where each node contains the following list: [IP, port, node_id]
  1242. # Therefore, cluster_slots[0][2][0] will be the IP address of the
  1243. # primary node of the first slot section.
  1244. # If there's only one server in the cluster, its ``host`` is ''
  1245. # Fix it to the host in startup_nodes
  1246. if (
  1247. len(cluster_slots) == 1
  1248. and not cluster_slots[0][2][0]
  1249. and len(self.startup_nodes) == 1
  1250. ):
  1251. cluster_slots[0][2][0] = startup_node.host
  1252. for slot in cluster_slots:
  1253. for i in range(2, len(slot)):
  1254. slot[i] = [str_if_bytes(val) for val in slot[i]]
  1255. primary_node = slot[2]
  1256. host = primary_node[0]
  1257. if host == "":
  1258. host = startup_node.host
  1259. port = int(primary_node[1])
  1260. host, port = self.remap_host_port(host, port)
  1261. nodes_for_slot = []
  1262. target_node = tmp_nodes_cache.get(get_node_name(host, port))
  1263. if not target_node:
  1264. target_node = ClusterNode(
  1265. host, port, PRIMARY, **self.connection_kwargs
  1266. )
  1267. # add this node to the nodes cache
  1268. tmp_nodes_cache[target_node.name] = target_node
  1269. nodes_for_slot.append(target_node)
  1270. replica_nodes = slot[3:]
  1271. for replica_node in replica_nodes:
  1272. host = replica_node[0]
  1273. port = replica_node[1]
  1274. host, port = self.remap_host_port(host, port)
  1275. target_replica_node = tmp_nodes_cache.get(get_node_name(host, port))
  1276. if not target_replica_node:
  1277. target_replica_node = ClusterNode(
  1278. host, port, REPLICA, **self.connection_kwargs
  1279. )
  1280. # add this node to the nodes cache
  1281. tmp_nodes_cache[target_replica_node.name] = target_replica_node
  1282. nodes_for_slot.append(target_replica_node)
  1283. for i in range(int(slot[0]), int(slot[1]) + 1):
  1284. if i not in tmp_slots:
  1285. tmp_slots[i] = nodes_for_slot
  1286. else:
  1287. # Validate that 2 nodes want to use the same slot cache
  1288. # setup
  1289. tmp_slot = tmp_slots[i][0]
  1290. if tmp_slot.name != target_node.name:
  1291. disagreements.append(
  1292. f"{tmp_slot.name} vs {target_node.name} on slot: {i}"
  1293. )
  1294. if len(disagreements) > 5:
  1295. raise RedisClusterException(
  1296. f"startup_nodes could not agree on a valid "
  1297. f"slots cache: {', '.join(disagreements)}"
  1298. )
  1299. # Validate if all slots are covered or if we should try next startup node
  1300. fully_covered = True
  1301. for i in range(REDIS_CLUSTER_HASH_SLOTS):
  1302. if i not in tmp_slots:
  1303. fully_covered = False
  1304. break
  1305. if fully_covered:
  1306. break
  1307. if not startup_nodes_reachable:
  1308. raise RedisClusterException(
  1309. f"Redis Cluster cannot be connected. Please provide at least "
  1310. f"one reachable node: {str(exception)}"
  1311. ) from exception
  1312. # Check if the slots are not fully covered
  1313. if not fully_covered and self.require_full_coverage:
  1314. # Despite the requirement that the slots be covered, there
  1315. # isn't a full coverage
  1316. raise RedisClusterException(
  1317. f"All slots are not covered after query all startup_nodes. "
  1318. f"{len(tmp_slots)} of {REDIS_CLUSTER_HASH_SLOTS} "
  1319. f"covered..."
  1320. )
  1321. # Set the tmp variables to the real variables
  1322. self.slots_cache = tmp_slots
  1323. self.set_nodes(self.nodes_cache, tmp_nodes_cache, remove_old=True)
  1324. if self._dynamic_startup_nodes:
  1325. # Populate the startup nodes with all discovered nodes
  1326. self.set_nodes(self.startup_nodes, self.nodes_cache, remove_old=True)
  1327. # Set the default node
  1328. self.default_node = self.get_nodes_by_server_type(PRIMARY)[0]
  1329. # If initialize was called after a MovedError, clear it
  1330. self._moved_exception = None
  1331. async def aclose(self, attr: str = "nodes_cache") -> None:
  1332. self.default_node = None
  1333. await asyncio.gather(
  1334. *(
  1335. asyncio.create_task(node.disconnect())
  1336. for node in getattr(self, attr).values()
  1337. )
  1338. )
  1339. def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
  1340. """
  1341. Remap the host and port returned from the cluster to a different
  1342. internal value. Useful if the client is not connecting directly
  1343. to the cluster.
  1344. """
  1345. if self.address_remap:
  1346. return self.address_remap((host, port))
  1347. return host, port
  1348. class ClusterPipeline(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
  1349. """
  1350. Create a new ClusterPipeline object.
  1351. Usage::
  1352. result = await (
  1353. rc.pipeline()
  1354. .set("A", 1)
  1355. .get("A")
  1356. .hset("K", "F", "V")
  1357. .hgetall("K")
  1358. .mset_nonatomic({"A": 2, "B": 3})
  1359. .get("A")
  1360. .get("B")
  1361. .delete("A", "B", "K")
  1362. .execute()
  1363. )
  1364. # result = [True, "1", 1, {"F": "V"}, True, True, "2", "3", 1, 1, 1]
  1365. Note: For commands `DELETE`, `EXISTS`, `TOUCH`, `UNLINK`, `mset_nonatomic`, which
  1366. are split across multiple nodes, you'll get multiple results for them in the array.
  1367. Retryable errors:
  1368. - :class:`~.ClusterDownError`
  1369. - :class:`~.ConnectionError`
  1370. - :class:`~.TimeoutError`
  1371. Redirection errors:
  1372. - :class:`~.TryAgainError`
  1373. - :class:`~.MovedError`
  1374. - :class:`~.AskError`
  1375. :param client:
  1376. | Existing :class:`~.RedisCluster` client
  1377. """
  1378. __slots__ = ("cluster_client", "_transaction", "_execution_strategy")
  1379. def __init__(
  1380. self, client: RedisCluster, transaction: Optional[bool] = None
  1381. ) -> None:
  1382. self.cluster_client = client
  1383. self._transaction = transaction
  1384. self._execution_strategy: ExecutionStrategy = (
  1385. PipelineStrategy(self)
  1386. if not self._transaction
  1387. else TransactionStrategy(self)
  1388. )
  1389. async def initialize(self) -> "ClusterPipeline":
  1390. await self._execution_strategy.initialize()
  1391. return self
  1392. async def __aenter__(self) -> "ClusterPipeline":
  1393. return await self.initialize()
  1394. async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  1395. await self.reset()
  1396. def __await__(self) -> Generator[Any, None, "ClusterPipeline"]:
  1397. return self.initialize().__await__()
  1398. def __enter__(self) -> "ClusterPipeline":
  1399. # TODO: Remove this method before 7.0.0
  1400. self._execution_strategy._command_queue = []
  1401. return self
  1402. def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  1403. # TODO: Remove this method before 7.0.0
  1404. self._execution_strategy._command_queue = []
  1405. def __bool__(self) -> bool:
  1406. "Pipeline instances should always evaluate to True on Python 3+"
  1407. return True
  1408. def __len__(self) -> int:
  1409. return len(self._execution_strategy)
  1410. def execute_command(
  1411. self, *args: Union[KeyT, EncodableT], **kwargs: Any
  1412. ) -> "ClusterPipeline":
  1413. """
  1414. Append a raw command to the pipeline.
  1415. :param args:
  1416. | Raw command args
  1417. :param kwargs:
  1418. - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
  1419. or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
  1420. - Rest of the kwargs are passed to the Redis connection
  1421. """
  1422. return self._execution_strategy.execute_command(*args, **kwargs)
  1423. async def execute(
  1424. self, raise_on_error: bool = True, allow_redirections: bool = True
  1425. ) -> List[Any]:
  1426. """
  1427. Execute the pipeline.
  1428. It will retry the commands as specified by retries specified in :attr:`retry`
  1429. & then raise an exception.
  1430. :param raise_on_error:
  1431. | Raise the first error if there are any errors
  1432. :param allow_redirections:
  1433. | Whether to retry each failed command individually in case of redirection
  1434. errors
  1435. :raises RedisClusterException: if target_nodes is not provided & the command
  1436. can't be mapped to a slot
  1437. """
  1438. try:
  1439. return await self._execution_strategy.execute(
  1440. raise_on_error, allow_redirections
  1441. )
  1442. finally:
  1443. await self.reset()
  1444. def _split_command_across_slots(
  1445. self, command: str, *keys: KeyT
  1446. ) -> "ClusterPipeline":
  1447. for slot_keys in self.cluster_client._partition_keys_by_slot(keys).values():
  1448. self.execute_command(command, *slot_keys)
  1449. return self
  1450. async def reset(self):
  1451. """
  1452. Reset back to empty pipeline.
  1453. """
  1454. await self._execution_strategy.reset()
  1455. def multi(self):
  1456. """
  1457. Start a transactional block of the pipeline after WATCH commands
  1458. are issued. End the transactional block with `execute`.
  1459. """
  1460. self._execution_strategy.multi()
  1461. async def discard(self):
  1462. """ """
  1463. await self._execution_strategy.discard()
  1464. async def watch(self, *names):
  1465. """Watches the values at keys ``names``"""
  1466. await self._execution_strategy.watch(*names)
  1467. async def unwatch(self):
  1468. """Unwatches all previously specified keys"""
  1469. await self._execution_strategy.unwatch()
  1470. async def unlink(self, *names):
  1471. await self._execution_strategy.unlink(*names)
  1472. def mset_nonatomic(
  1473. self, mapping: Mapping[AnyKeyT, EncodableT]
  1474. ) -> "ClusterPipeline":
  1475. return self._execution_strategy.mset_nonatomic(mapping)
  1476. for command in PIPELINE_BLOCKED_COMMANDS:
  1477. command = command.replace(" ", "_").lower()
  1478. if command == "mset_nonatomic":
  1479. continue
  1480. setattr(ClusterPipeline, command, block_pipeline_command(command))
  1481. class PipelineCommand:
  1482. def __init__(self, position: int, *args: Any, **kwargs: Any) -> None:
  1483. self.args = args
  1484. self.kwargs = kwargs
  1485. self.position = position
  1486. self.result: Union[Any, Exception] = None
  1487. def __repr__(self) -> str:
  1488. return f"[{self.position}] {self.args} ({self.kwargs})"
  1489. class ExecutionStrategy(ABC):
  1490. @abstractmethod
  1491. async def initialize(self) -> "ClusterPipeline":
  1492. """
  1493. Initialize the execution strategy.
  1494. See ClusterPipeline.initialize()
  1495. """
  1496. pass
  1497. @abstractmethod
  1498. def execute_command(
  1499. self, *args: Union[KeyT, EncodableT], **kwargs: Any
  1500. ) -> "ClusterPipeline":
  1501. """
  1502. Append a raw command to the pipeline.
  1503. See ClusterPipeline.execute_command()
  1504. """
  1505. pass
  1506. @abstractmethod
  1507. async def execute(
  1508. self, raise_on_error: bool = True, allow_redirections: bool = True
  1509. ) -> List[Any]:
  1510. """
  1511. Execute the pipeline.
  1512. It will retry the commands as specified by retries specified in :attr:`retry`
  1513. & then raise an exception.
  1514. See ClusterPipeline.execute()
  1515. """
  1516. pass
  1517. @abstractmethod
  1518. def mset_nonatomic(
  1519. self, mapping: Mapping[AnyKeyT, EncodableT]
  1520. ) -> "ClusterPipeline":
  1521. """
  1522. Executes multiple MSET commands according to the provided slot/pairs mapping.
  1523. See ClusterPipeline.mset_nonatomic()
  1524. """
  1525. pass
  1526. @abstractmethod
  1527. async def reset(self):
  1528. """
  1529. Resets current execution strategy.
  1530. See: ClusterPipeline.reset()
  1531. """
  1532. pass
  1533. @abstractmethod
  1534. def multi(self):
  1535. """
  1536. Starts transactional context.
  1537. See: ClusterPipeline.multi()
  1538. """
  1539. pass
  1540. @abstractmethod
  1541. async def watch(self, *names):
  1542. """
  1543. Watch given keys.
  1544. See: ClusterPipeline.watch()
  1545. """
  1546. pass
  1547. @abstractmethod
  1548. async def unwatch(self):
  1549. """
  1550. Unwatches all previously specified keys
  1551. See: ClusterPipeline.unwatch()
  1552. """
  1553. pass
  1554. @abstractmethod
  1555. async def discard(self):
  1556. pass
  1557. @abstractmethod
  1558. async def unlink(self, *names):
  1559. """
  1560. "Unlink a key specified by ``names``"
  1561. See: ClusterPipeline.unlink()
  1562. """
  1563. pass
  1564. @abstractmethod
  1565. def __len__(self) -> int:
  1566. pass
  1567. class AbstractStrategy(ExecutionStrategy):
  1568. def __init__(self, pipe: ClusterPipeline) -> None:
  1569. self._pipe: ClusterPipeline = pipe
  1570. self._command_queue: List["PipelineCommand"] = []
  1571. async def initialize(self) -> "ClusterPipeline":
  1572. if self._pipe.cluster_client._initialize:
  1573. await self._pipe.cluster_client.initialize()
  1574. self._command_queue = []
  1575. return self._pipe
  1576. def execute_command(
  1577. self, *args: Union[KeyT, EncodableT], **kwargs: Any
  1578. ) -> "ClusterPipeline":
  1579. self._command_queue.append(
  1580. PipelineCommand(len(self._command_queue), *args, **kwargs)
  1581. )
  1582. return self._pipe
  1583. def _annotate_exception(self, exception, number, command):
  1584. """
  1585. Provides extra context to the exception prior to it being handled
  1586. """
  1587. cmd = " ".join(map(safe_str, command))
  1588. msg = (
  1589. f"Command # {number} ({truncate_text(cmd)}) of pipeline "
  1590. f"caused error: {exception.args[0]}"
  1591. )
  1592. exception.args = (msg,) + exception.args[1:]
  1593. @abstractmethod
  1594. def mset_nonatomic(
  1595. self, mapping: Mapping[AnyKeyT, EncodableT]
  1596. ) -> "ClusterPipeline":
  1597. pass
  1598. @abstractmethod
  1599. async def execute(
  1600. self, raise_on_error: bool = True, allow_redirections: bool = True
  1601. ) -> List[Any]:
  1602. pass
  1603. @abstractmethod
  1604. async def reset(self):
  1605. pass
  1606. @abstractmethod
  1607. def multi(self):
  1608. pass
  1609. @abstractmethod
  1610. async def watch(self, *names):
  1611. pass
  1612. @abstractmethod
  1613. async def unwatch(self):
  1614. pass
  1615. @abstractmethod
  1616. async def discard(self):
  1617. pass
  1618. @abstractmethod
  1619. async def unlink(self, *names):
  1620. pass
  1621. def __len__(self) -> int:
  1622. return len(self._command_queue)
  1623. class PipelineStrategy(AbstractStrategy):
  1624. def __init__(self, pipe: ClusterPipeline) -> None:
  1625. super().__init__(pipe)
  1626. def mset_nonatomic(
  1627. self, mapping: Mapping[AnyKeyT, EncodableT]
  1628. ) -> "ClusterPipeline":
  1629. encoder = self._pipe.cluster_client.encoder
  1630. slots_pairs = {}
  1631. for pair in mapping.items():
  1632. slot = key_slot(encoder.encode(pair[0]))
  1633. slots_pairs.setdefault(slot, []).extend(pair)
  1634. for pairs in slots_pairs.values():
  1635. self.execute_command("MSET", *pairs)
  1636. return self._pipe
  1637. async def execute(
  1638. self, raise_on_error: bool = True, allow_redirections: bool = True
  1639. ) -> List[Any]:
  1640. if not self._command_queue:
  1641. return []
  1642. try:
  1643. retry_attempts = self._pipe.cluster_client.retry.get_retries()
  1644. while True:
  1645. try:
  1646. if self._pipe.cluster_client._initialize:
  1647. await self._pipe.cluster_client.initialize()
  1648. return await self._execute(
  1649. self._pipe.cluster_client,
  1650. self._command_queue,
  1651. raise_on_error=raise_on_error,
  1652. allow_redirections=allow_redirections,
  1653. )
  1654. except RedisCluster.ERRORS_ALLOW_RETRY as e:
  1655. if retry_attempts > 0:
  1656. # Try again with the new cluster setup. All other errors
  1657. # should be raised.
  1658. retry_attempts -= 1
  1659. await self._pipe.cluster_client.aclose()
  1660. await asyncio.sleep(0.25)
  1661. else:
  1662. # All other errors should be raised.
  1663. raise e
  1664. finally:
  1665. await self.reset()
  1666. async def _execute(
  1667. self,
  1668. client: "RedisCluster",
  1669. stack: List["PipelineCommand"],
  1670. raise_on_error: bool = True,
  1671. allow_redirections: bool = True,
  1672. ) -> List[Any]:
  1673. todo = [
  1674. cmd for cmd in stack if not cmd.result or isinstance(cmd.result, Exception)
  1675. ]
  1676. nodes = {}
  1677. for cmd in todo:
  1678. passed_targets = cmd.kwargs.pop("target_nodes", None)
  1679. if passed_targets and not client._is_node_flag(passed_targets):
  1680. target_nodes = client._parse_target_nodes(passed_targets)
  1681. else:
  1682. target_nodes = await client._determine_nodes(
  1683. *cmd.args, node_flag=passed_targets
  1684. )
  1685. if not target_nodes:
  1686. raise RedisClusterException(
  1687. f"No targets were found to execute {cmd.args} command on"
  1688. )
  1689. if len(target_nodes) > 1:
  1690. raise RedisClusterException(f"Too many targets for command {cmd.args}")
  1691. node = target_nodes[0]
  1692. if node.name not in nodes:
  1693. nodes[node.name] = (node, [])
  1694. nodes[node.name][1].append(cmd)
  1695. errors = await asyncio.gather(
  1696. *(
  1697. asyncio.create_task(node[0].execute_pipeline(node[1]))
  1698. for node in nodes.values()
  1699. )
  1700. )
  1701. if any(errors):
  1702. if allow_redirections:
  1703. # send each errored command individually
  1704. for cmd in todo:
  1705. if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
  1706. try:
  1707. cmd.result = await client.execute_command(
  1708. *cmd.args, **cmd.kwargs
  1709. )
  1710. except Exception as e:
  1711. cmd.result = e
  1712. if raise_on_error:
  1713. for cmd in todo:
  1714. result = cmd.result
  1715. if isinstance(result, Exception):
  1716. command = " ".join(map(safe_str, cmd.args))
  1717. msg = (
  1718. f"Command # {cmd.position + 1} "
  1719. f"({truncate_text(command)}) "
  1720. f"of pipeline caused error: {result.args}"
  1721. )
  1722. result.args = (msg,) + result.args[1:]
  1723. raise result
  1724. default_cluster_node = client.get_default_node()
  1725. # Check whether the default node was used. In some cases,
  1726. # 'client.get_default_node()' may return None. The check below
  1727. # prevents a potential AttributeError.
  1728. if default_cluster_node is not None:
  1729. default_node = nodes.get(default_cluster_node.name)
  1730. if default_node is not None:
  1731. # This pipeline execution used the default node, check if we need
  1732. # to replace it.
  1733. # Note: when the error is raised we'll reset the default node in the
  1734. # caller function.
  1735. for cmd in default_node[1]:
  1736. # Check if it has a command that failed with a relevant
  1737. # exception
  1738. if type(cmd.result) in RedisCluster.ERRORS_ALLOW_RETRY:
  1739. client.replace_default_node()
  1740. break
  1741. return [cmd.result for cmd in stack]
  1742. async def reset(self):
  1743. """
  1744. Reset back to empty pipeline.
  1745. """
  1746. self._command_queue = []
  1747. def multi(self):
  1748. raise RedisClusterException(
  1749. "method multi() is not supported outside of transactional context"
  1750. )
  1751. async def watch(self, *names):
  1752. raise RedisClusterException(
  1753. "method watch() is not supported outside of transactional context"
  1754. )
  1755. async def unwatch(self):
  1756. raise RedisClusterException(
  1757. "method unwatch() is not supported outside of transactional context"
  1758. )
  1759. async def discard(self):
  1760. raise RedisClusterException(
  1761. "method discard() is not supported outside of transactional context"
  1762. )
  1763. async def unlink(self, *names):
  1764. if len(names) != 1:
  1765. raise RedisClusterException(
  1766. "unlinking multiple keys is not implemented in pipeline command"
  1767. )
  1768. return self.execute_command("UNLINK", names[0])
  1769. class TransactionStrategy(AbstractStrategy):
  1770. NO_SLOTS_COMMANDS = {"UNWATCH"}
  1771. IMMEDIATE_EXECUTE_COMMANDS = {"WATCH", "UNWATCH"}
  1772. UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
  1773. SLOT_REDIRECT_ERRORS = (AskError, MovedError)
  1774. CONNECTION_ERRORS = (
  1775. ConnectionError,
  1776. OSError,
  1777. ClusterDownError,
  1778. SlotNotCoveredError,
  1779. )
  1780. def __init__(self, pipe: ClusterPipeline) -> None:
  1781. super().__init__(pipe)
  1782. self._explicit_transaction = False
  1783. self._watching = False
  1784. self._pipeline_slots: Set[int] = set()
  1785. self._transaction_node: Optional[ClusterNode] = None
  1786. self._transaction_connection: Optional[Connection] = None
  1787. self._executing = False
  1788. self._retry = copy(self._pipe.cluster_client.retry)
  1789. self._retry.update_supported_errors(
  1790. RedisCluster.ERRORS_ALLOW_RETRY + self.SLOT_REDIRECT_ERRORS
  1791. )
  1792. def _get_client_and_connection_for_transaction(
  1793. self,
  1794. ) -> Tuple[ClusterNode, Connection]:
  1795. """
  1796. Find a connection for a pipeline transaction.
  1797. For running an atomic transaction, watch keys ensure that contents have not been
  1798. altered as long as the watch commands for those keys were sent over the same
  1799. connection. So once we start watching a key, we fetch a connection to the
  1800. node that owns that slot and reuse it.
  1801. """
  1802. if not self._pipeline_slots:
  1803. raise RedisClusterException(
  1804. "At least a command with a key is needed to identify a node"
  1805. )
  1806. node: ClusterNode = self._pipe.cluster_client.nodes_manager.get_node_from_slot(
  1807. list(self._pipeline_slots)[0], False
  1808. )
  1809. self._transaction_node = node
  1810. if not self._transaction_connection:
  1811. connection: Connection = self._transaction_node.acquire_connection()
  1812. self._transaction_connection = connection
  1813. return self._transaction_node, self._transaction_connection
  1814. def execute_command(self, *args: Union[KeyT, EncodableT], **kwargs: Any) -> "Any":
  1815. # Given the limitation of ClusterPipeline sync API, we have to run it in thread.
  1816. response = None
  1817. error = None
  1818. def runner():
  1819. nonlocal response
  1820. nonlocal error
  1821. try:
  1822. response = asyncio.run(self._execute_command(*args, **kwargs))
  1823. except Exception as e:
  1824. error = e
  1825. thread = threading.Thread(target=runner)
  1826. thread.start()
  1827. thread.join()
  1828. if error:
  1829. raise error
  1830. return response
  1831. async def _execute_command(
  1832. self, *args: Union[KeyT, EncodableT], **kwargs: Any
  1833. ) -> Any:
  1834. if self._pipe.cluster_client._initialize:
  1835. await self._pipe.cluster_client.initialize()
  1836. slot_number: Optional[int] = None
  1837. if args[0] not in self.NO_SLOTS_COMMANDS:
  1838. slot_number = await self._pipe.cluster_client._determine_slot(*args)
  1839. if (
  1840. self._watching or args[0] in self.IMMEDIATE_EXECUTE_COMMANDS
  1841. ) and not self._explicit_transaction:
  1842. if args[0] == "WATCH":
  1843. self._validate_watch()
  1844. if slot_number is not None:
  1845. if self._pipeline_slots and slot_number not in self._pipeline_slots:
  1846. raise CrossSlotTransactionError(
  1847. "Cannot watch or send commands on different slots"
  1848. )
  1849. self._pipeline_slots.add(slot_number)
  1850. elif args[0] not in self.NO_SLOTS_COMMANDS:
  1851. raise RedisClusterException(
  1852. f"Cannot identify slot number for command: {args[0]},"
  1853. "it cannot be triggered in a transaction"
  1854. )
  1855. return self._immediate_execute_command(*args, **kwargs)
  1856. else:
  1857. if slot_number is not None:
  1858. self._pipeline_slots.add(slot_number)
  1859. return super().execute_command(*args, **kwargs)
  1860. def _validate_watch(self):
  1861. if self._explicit_transaction:
  1862. raise RedisError("Cannot issue a WATCH after a MULTI")
  1863. self._watching = True
  1864. async def _immediate_execute_command(self, *args, **options):
  1865. return await self._retry.call_with_retry(
  1866. lambda: self._get_connection_and_send_command(*args, **options),
  1867. self._reinitialize_on_error,
  1868. )
  1869. async def _get_connection_and_send_command(self, *args, **options):
  1870. redis_node, connection = self._get_client_and_connection_for_transaction()
  1871. return await self._send_command_parse_response(
  1872. connection, redis_node, args[0], *args, **options
  1873. )
  1874. async def _send_command_parse_response(
  1875. self,
  1876. connection: Connection,
  1877. redis_node: ClusterNode,
  1878. command_name,
  1879. *args,
  1880. **options,
  1881. ):
  1882. """
  1883. Send a command and parse the response
  1884. """
  1885. await connection.send_command(*args)
  1886. output = await redis_node.parse_response(connection, command_name, **options)
  1887. if command_name in self.UNWATCH_COMMANDS:
  1888. self._watching = False
  1889. return output
  1890. async def _reinitialize_on_error(self, error):
  1891. if self._watching:
  1892. if type(error) in self.SLOT_REDIRECT_ERRORS and self._executing:
  1893. raise WatchError("Slot rebalancing occurred while watching keys")
  1894. if (
  1895. type(error) in self.SLOT_REDIRECT_ERRORS
  1896. or type(error) in self.CONNECTION_ERRORS
  1897. ):
  1898. if self._transaction_connection:
  1899. self._transaction_connection = None
  1900. self._pipe.cluster_client.reinitialize_counter += 1
  1901. if (
  1902. self._pipe.cluster_client.reinitialize_steps
  1903. and self._pipe.cluster_client.reinitialize_counter
  1904. % self._pipe.cluster_client.reinitialize_steps
  1905. == 0
  1906. ):
  1907. await self._pipe.cluster_client.nodes_manager.initialize()
  1908. self.reinitialize_counter = 0
  1909. else:
  1910. self._pipe.cluster_client.nodes_manager.update_moved_exception(error)
  1911. self._executing = False
  1912. def _raise_first_error(self, responses, stack):
  1913. """
  1914. Raise the first exception on the stack
  1915. """
  1916. for r, cmd in zip(responses, stack):
  1917. if isinstance(r, Exception):
  1918. self._annotate_exception(r, cmd.position + 1, cmd.args)
  1919. raise r
  1920. def mset_nonatomic(
  1921. self, mapping: Mapping[AnyKeyT, EncodableT]
  1922. ) -> "ClusterPipeline":
  1923. raise NotImplementedError("Method is not supported in transactional context.")
  1924. async def execute(
  1925. self, raise_on_error: bool = True, allow_redirections: bool = True
  1926. ) -> List[Any]:
  1927. stack = self._command_queue
  1928. if not stack and (not self._watching or not self._pipeline_slots):
  1929. return []
  1930. return await self._execute_transaction_with_retries(stack, raise_on_error)
  1931. async def _execute_transaction_with_retries(
  1932. self, stack: List["PipelineCommand"], raise_on_error: bool
  1933. ):
  1934. return await self._retry.call_with_retry(
  1935. lambda: self._execute_transaction(stack, raise_on_error),
  1936. self._reinitialize_on_error,
  1937. )
  1938. async def _execute_transaction(
  1939. self, stack: List["PipelineCommand"], raise_on_error: bool
  1940. ):
  1941. if len(self._pipeline_slots) > 1:
  1942. raise CrossSlotTransactionError(
  1943. "All keys involved in a cluster transaction must map to the same slot"
  1944. )
  1945. self._executing = True
  1946. redis_node, connection = self._get_client_and_connection_for_transaction()
  1947. stack = chain(
  1948. [PipelineCommand(0, "MULTI")],
  1949. stack,
  1950. [PipelineCommand(0, "EXEC")],
  1951. )
  1952. commands = [c.args for c in stack if EMPTY_RESPONSE not in c.kwargs]
  1953. packed_commands = connection.pack_commands(commands)
  1954. await connection.send_packed_command(packed_commands)
  1955. errors = []
  1956. # parse off the response for MULTI
  1957. # NOTE: we need to handle ResponseErrors here and continue
  1958. # so that we read all the additional command messages from
  1959. # the socket
  1960. try:
  1961. await redis_node.parse_response(connection, "MULTI")
  1962. except ResponseError as e:
  1963. self._annotate_exception(e, 0, "MULTI")
  1964. errors.append(e)
  1965. except self.CONNECTION_ERRORS as cluster_error:
  1966. self._annotate_exception(cluster_error, 0, "MULTI")
  1967. raise
  1968. # and all the other commands
  1969. for i, command in enumerate(self._command_queue):
  1970. if EMPTY_RESPONSE in command.kwargs:
  1971. errors.append((i, command.kwargs[EMPTY_RESPONSE]))
  1972. else:
  1973. try:
  1974. _ = await redis_node.parse_response(connection, "_")
  1975. except self.SLOT_REDIRECT_ERRORS as slot_error:
  1976. self._annotate_exception(slot_error, i + 1, command.args)
  1977. errors.append(slot_error)
  1978. except self.CONNECTION_ERRORS as cluster_error:
  1979. self._annotate_exception(cluster_error, i + 1, command.args)
  1980. raise
  1981. except ResponseError as e:
  1982. self._annotate_exception(e, i + 1, command.args)
  1983. errors.append(e)
  1984. response = None
  1985. # parse the EXEC.
  1986. try:
  1987. response = await redis_node.parse_response(connection, "EXEC")
  1988. except ExecAbortError:
  1989. if errors:
  1990. raise errors[0]
  1991. raise
  1992. self._executing = False
  1993. # EXEC clears any watched keys
  1994. self._watching = False
  1995. if response is None:
  1996. raise WatchError("Watched variable changed.")
  1997. # put any parse errors into the response
  1998. for i, e in errors:
  1999. response.insert(i, e)
  2000. if len(response) != len(self._command_queue):
  2001. raise InvalidPipelineStack(
  2002. "Unexpected response length for cluster pipeline EXEC."
  2003. " Command stack was {} but response had length {}".format(
  2004. [c.args[0] for c in self._command_queue], len(response)
  2005. )
  2006. )
  2007. # find any errors in the response and raise if necessary
  2008. if raise_on_error or len(errors) > 0:
  2009. self._raise_first_error(
  2010. response,
  2011. self._command_queue,
  2012. )
  2013. # We have to run response callbacks manually
  2014. data = []
  2015. for r, cmd in zip(response, self._command_queue):
  2016. if not isinstance(r, Exception):
  2017. command_name = cmd.args[0]
  2018. if command_name in self._pipe.cluster_client.response_callbacks:
  2019. r = self._pipe.cluster_client.response_callbacks[command_name](
  2020. r, **cmd.kwargs
  2021. )
  2022. data.append(r)
  2023. return data
  2024. async def reset(self):
  2025. self._command_queue = []
  2026. # make sure to reset the connection state in the event that we were
  2027. # watching something
  2028. if self._transaction_connection:
  2029. try:
  2030. # call this manually since our unwatch or
  2031. # immediate_execute_command methods can call reset()
  2032. await self._transaction_connection.send_command("UNWATCH")
  2033. await self._transaction_connection.read_response()
  2034. # we can safely return the connection to the pool here since we're
  2035. # sure we're no longer WATCHing anything
  2036. self._transaction_node.release(self._transaction_connection)
  2037. self._transaction_connection = None
  2038. except self.CONNECTION_ERRORS:
  2039. # disconnect will also remove any previous WATCHes
  2040. if self._transaction_connection:
  2041. await self._transaction_connection.disconnect()
  2042. # clean up the other instance attributes
  2043. self._transaction_node = None
  2044. self._watching = False
  2045. self._explicit_transaction = False
  2046. self._pipeline_slots = set()
  2047. self._executing = False
  2048. def multi(self):
  2049. if self._explicit_transaction:
  2050. raise RedisError("Cannot issue nested calls to MULTI")
  2051. if self._command_queue:
  2052. raise RedisError(
  2053. "Commands without an initial WATCH have already been issued"
  2054. )
  2055. self._explicit_transaction = True
  2056. async def watch(self, *names):
  2057. if self._explicit_transaction:
  2058. raise RedisError("Cannot issue a WATCH after a MULTI")
  2059. return await self.execute_command("WATCH", *names)
  2060. async def unwatch(self):
  2061. if self._watching:
  2062. return await self.execute_command("UNWATCH")
  2063. return True
  2064. async def discard(self):
  2065. await self.reset()
  2066. async def unlink(self, *names):
  2067. return self.execute_command("UNLINK", *names)