cluster.py 62 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620
  1. import asyncio
  2. import collections
  3. import random
  4. import socket
  5. import warnings
  6. from typing import (
  7. Any,
  8. Callable,
  9. Deque,
  10. Dict,
  11. Generator,
  12. List,
  13. Mapping,
  14. Optional,
  15. Tuple,
  16. Type,
  17. TypeVar,
  18. Union,
  19. )
  20. from redis._parsers import AsyncCommandsParser, Encoder
  21. from redis._parsers.helpers import (
  22. _RedisCallbacks,
  23. _RedisCallbacksRESP2,
  24. _RedisCallbacksRESP3,
  25. )
  26. from redis.asyncio.client import ResponseCallbackT
  27. from redis.asyncio.connection import Connection, DefaultParser, SSLConnection, parse_url
  28. from redis.asyncio.lock import Lock
  29. from redis.asyncio.retry import Retry
  30. from redis.backoff import default_backoff
  31. from redis.client import EMPTY_RESPONSE, NEVER_DECODE, AbstractRedis
  32. from redis.cluster import (
  33. PIPELINE_BLOCKED_COMMANDS,
  34. PRIMARY,
  35. REPLICA,
  36. SLOT_ID,
  37. AbstractRedisCluster,
  38. LoadBalancer,
  39. block_pipeline_command,
  40. get_node_name,
  41. parse_cluster_slots,
  42. )
  43. from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
  44. from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
  45. from redis.credentials import CredentialProvider
  46. from redis.exceptions import (
  47. AskError,
  48. BusyLoadingError,
  49. ClusterCrossSlotError,
  50. ClusterDownError,
  51. ClusterError,
  52. ConnectionError,
  53. DataError,
  54. MasterDownError,
  55. MaxConnectionsError,
  56. MovedError,
  57. RedisClusterException,
  58. ResponseError,
  59. SlotNotCoveredError,
  60. TimeoutError,
  61. TryAgainError,
  62. )
  63. from redis.typing import AnyKeyT, EncodableT, KeyT
  64. from redis.utils import (
  65. deprecated_function,
  66. dict_merge,
  67. get_lib_version,
  68. safe_str,
  69. str_if_bytes,
  70. )
  71. TargetNodesT = TypeVar(
  72. "TargetNodesT", str, "ClusterNode", List["ClusterNode"], Dict[Any, "ClusterNode"]
  73. )
  74. class ClusterParser(DefaultParser):
  75. EXCEPTION_CLASSES = dict_merge(
  76. DefaultParser.EXCEPTION_CLASSES,
  77. {
  78. "ASK": AskError,
  79. "CLUSTERDOWN": ClusterDownError,
  80. "CROSSSLOT": ClusterCrossSlotError,
  81. "MASTERDOWN": MasterDownError,
  82. "MOVED": MovedError,
  83. "TRYAGAIN": TryAgainError,
  84. },
  85. )
  86. class RedisCluster(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
  87. """
  88. Create a new RedisCluster client.
  89. Pass one of parameters:
  90. - `host` & `port`
  91. - `startup_nodes`
  92. | Use ``await`` :meth:`initialize` to find cluster nodes & create connections.
  93. | Use ``await`` :meth:`close` to disconnect connections & close client.
  94. Many commands support the target_nodes kwarg. It can be one of the
  95. :attr:`NODE_FLAGS`:
  96. - :attr:`PRIMARIES`
  97. - :attr:`REPLICAS`
  98. - :attr:`ALL_NODES`
  99. - :attr:`RANDOM`
  100. - :attr:`DEFAULT_NODE`
  101. Note: This client is not thread/process/fork safe.
  102. :param host:
  103. | Can be used to point to a startup node
  104. :param port:
  105. | Port used if **host** is provided
  106. :param startup_nodes:
  107. | :class:`~.ClusterNode` to used as a startup node
  108. :param require_full_coverage:
  109. | When set to ``False``: the client will not require a full coverage of
  110. the slots. However, if not all slots are covered, and at least one node
  111. has ``cluster-require-full-coverage`` set to ``yes``, the server will throw
  112. a :class:`~.ClusterDownError` for some key-based commands.
  113. | When set to ``True``: all slots must be covered to construct the cluster
  114. client. If not all slots are covered, :class:`~.RedisClusterException` will be
  115. thrown.
  116. | See:
  117. https://redis.io/docs/manual/scaling/#redis-cluster-configuration-parameters
  118. :param read_from_replicas:
  119. | Enable read from replicas in READONLY mode. You can read possibly stale data.
  120. When set to true, read commands will be assigned between the primary and
  121. its replications in a Round-Robin manner.
  122. :param reinitialize_steps:
  123. | Specifies the number of MOVED errors that need to occur before reinitializing
  124. the whole cluster topology. If a MOVED error occurs and the cluster does not
  125. need to be reinitialized on this current error handling, only the MOVED slot
  126. will be patched with the redirected node.
  127. To reinitialize the cluster on every MOVED error, set reinitialize_steps to 1.
  128. To avoid reinitializing the cluster on moved errors, set reinitialize_steps to
  129. 0.
  130. :param cluster_error_retry_attempts:
  131. | Number of times to retry before raising an error when :class:`~.TimeoutError`
  132. or :class:`~.ConnectionError` or :class:`~.ClusterDownError` are encountered
  133. :param connection_error_retry_attempts:
  134. | Number of times to retry before reinitializing when :class:`~.TimeoutError`
  135. or :class:`~.ConnectionError` are encountered.
  136. The default backoff strategy will be set if Retry object is not passed (see
  137. default_backoff in backoff.py). To change it, pass a custom Retry object
  138. using the "retry" keyword.
  139. :param max_connections:
  140. | Maximum number of connections per node. If there are no free connections & the
  141. maximum number of connections are already created, a
  142. :class:`~.MaxConnectionsError` is raised. This error may be retried as defined
  143. by :attr:`connection_error_retry_attempts`
  144. :param address_remap:
  145. | An optional callable which, when provided with an internal network
  146. address of a node, e.g. a `(host, port)` tuple, will return the address
  147. where the node is reachable. This can be used to map the addresses at
  148. which the nodes _think_ they are, to addresses at which a client may
  149. reach them, such as when they sit behind a proxy.
  150. | Rest of the arguments will be passed to the
  151. :class:`~redis.asyncio.connection.Connection` instances when created
  152. :raises RedisClusterException:
  153. if any arguments are invalid or unknown. Eg:
  154. - `db` != 0 or None
  155. - `path` argument for unix socket connection
  156. - none of the `host`/`port` & `startup_nodes` were provided
  157. """
  158. @classmethod
  159. def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
  160. """
  161. Return a Redis client object configured from the given URL.
  162. For example::
  163. redis://[[username]:[password]]@localhost:6379/0
  164. rediss://[[username]:[password]]@localhost:6379/0
  165. Three URL schemes are supported:
  166. - `redis://` creates a TCP socket connection. See more at:
  167. <https://www.iana.org/assignments/uri-schemes/prov/redis>
  168. - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
  169. <https://www.iana.org/assignments/uri-schemes/prov/rediss>
  170. The username, password, hostname, path and all querystring values are passed
  171. through ``urllib.parse.unquote`` in order to replace any percent-encoded values
  172. with their corresponding characters.
  173. All querystring options are cast to their appropriate Python types. Boolean
  174. arguments can be specified with string values "True"/"False" or "Yes"/"No".
  175. Values that cannot be properly cast cause a ``ValueError`` to be raised. Once
  176. parsed, the querystring arguments and keyword arguments are passed to
  177. :class:`~redis.asyncio.connection.Connection` when created.
  178. In the case of conflicting arguments, querystring arguments are used.
  179. """
  180. kwargs.update(parse_url(url))
  181. if kwargs.pop("connection_class", None) is SSLConnection:
  182. kwargs["ssl"] = True
  183. return cls(**kwargs)
  184. __slots__ = (
  185. "_initialize",
  186. "_lock",
  187. "cluster_error_retry_attempts",
  188. "command_flags",
  189. "commands_parser",
  190. "connection_error_retry_attempts",
  191. "connection_kwargs",
  192. "encoder",
  193. "node_flags",
  194. "nodes_manager",
  195. "read_from_replicas",
  196. "reinitialize_counter",
  197. "reinitialize_steps",
  198. "response_callbacks",
  199. "result_callbacks",
  200. )
  201. def __init__(
  202. self,
  203. host: Optional[str] = None,
  204. port: Union[str, int] = 6379,
  205. # Cluster related kwargs
  206. startup_nodes: Optional[List["ClusterNode"]] = None,
  207. require_full_coverage: bool = True,
  208. read_from_replicas: bool = False,
  209. reinitialize_steps: int = 5,
  210. cluster_error_retry_attempts: int = 3,
  211. connection_error_retry_attempts: int = 3,
  212. max_connections: int = 2**31,
  213. # Client related kwargs
  214. db: Union[str, int] = 0,
  215. path: Optional[str] = None,
  216. credential_provider: Optional[CredentialProvider] = None,
  217. username: Optional[str] = None,
  218. password: Optional[str] = None,
  219. client_name: Optional[str] = None,
  220. lib_name: Optional[str] = "redis-py",
  221. lib_version: Optional[str] = get_lib_version(),
  222. # Encoding related kwargs
  223. encoding: str = "utf-8",
  224. encoding_errors: str = "strict",
  225. decode_responses: bool = False,
  226. # Connection related kwargs
  227. health_check_interval: float = 0,
  228. socket_connect_timeout: Optional[float] = None,
  229. socket_keepalive: bool = False,
  230. socket_keepalive_options: Optional[Mapping[int, Union[int, bytes]]] = None,
  231. socket_timeout: Optional[float] = None,
  232. retry: Optional["Retry"] = None,
  233. retry_on_error: Optional[List[Type[Exception]]] = None,
  234. # SSL related kwargs
  235. ssl: bool = False,
  236. ssl_ca_certs: Optional[str] = None,
  237. ssl_ca_data: Optional[str] = None,
  238. ssl_cert_reqs: str = "required",
  239. ssl_certfile: Optional[str] = None,
  240. ssl_check_hostname: bool = False,
  241. ssl_keyfile: Optional[str] = None,
  242. protocol: Optional[int] = 2,
  243. address_remap: Optional[Callable[[str, int], Tuple[str, int]]] = None,
  244. ) -> None:
  245. if db:
  246. raise RedisClusterException(
  247. "Argument 'db' must be 0 or None in cluster mode"
  248. )
  249. if path:
  250. raise RedisClusterException(
  251. "Unix domain socket is not supported in cluster mode"
  252. )
  253. if (not host or not port) and not startup_nodes:
  254. raise RedisClusterException(
  255. "RedisCluster requires at least one node to discover the cluster.\n"
  256. "Please provide one of the following or use RedisCluster.from_url:\n"
  257. ' - host and port: RedisCluster(host="localhost", port=6379)\n'
  258. " - startup_nodes: RedisCluster(startup_nodes=["
  259. 'ClusterNode("localhost", 6379), ClusterNode("localhost", 6380)])'
  260. )
  261. kwargs: Dict[str, Any] = {
  262. "max_connections": max_connections,
  263. "connection_class": Connection,
  264. "parser_class": ClusterParser,
  265. # Client related kwargs
  266. "credential_provider": credential_provider,
  267. "username": username,
  268. "password": password,
  269. "client_name": client_name,
  270. "lib_name": lib_name,
  271. "lib_version": lib_version,
  272. # Encoding related kwargs
  273. "encoding": encoding,
  274. "encoding_errors": encoding_errors,
  275. "decode_responses": decode_responses,
  276. # Connection related kwargs
  277. "health_check_interval": health_check_interval,
  278. "socket_connect_timeout": socket_connect_timeout,
  279. "socket_keepalive": socket_keepalive,
  280. "socket_keepalive_options": socket_keepalive_options,
  281. "socket_timeout": socket_timeout,
  282. "retry": retry,
  283. "protocol": protocol,
  284. }
  285. if ssl:
  286. # SSL related kwargs
  287. kwargs.update(
  288. {
  289. "connection_class": SSLConnection,
  290. "ssl_ca_certs": ssl_ca_certs,
  291. "ssl_ca_data": ssl_ca_data,
  292. "ssl_cert_reqs": ssl_cert_reqs,
  293. "ssl_certfile": ssl_certfile,
  294. "ssl_check_hostname": ssl_check_hostname,
  295. "ssl_keyfile": ssl_keyfile,
  296. }
  297. )
  298. if read_from_replicas:
  299. # Call our on_connect function to configure READONLY mode
  300. kwargs["redis_connect_func"] = self.on_connect
  301. self.retry = retry
  302. if retry or retry_on_error or connection_error_retry_attempts > 0:
  303. # Set a retry object for all cluster nodes
  304. self.retry = retry or Retry(
  305. default_backoff(), connection_error_retry_attempts
  306. )
  307. if not retry_on_error:
  308. # Default errors for retrying
  309. retry_on_error = [ConnectionError, TimeoutError]
  310. self.retry.update_supported_errors(retry_on_error)
  311. kwargs.update({"retry": self.retry})
  312. kwargs["response_callbacks"] = _RedisCallbacks.copy()
  313. if kwargs.get("protocol") in ["3", 3]:
  314. kwargs["response_callbacks"].update(_RedisCallbacksRESP3)
  315. else:
  316. kwargs["response_callbacks"].update(_RedisCallbacksRESP2)
  317. self.connection_kwargs = kwargs
  318. if startup_nodes:
  319. passed_nodes = []
  320. for node in startup_nodes:
  321. passed_nodes.append(
  322. ClusterNode(node.host, node.port, **self.connection_kwargs)
  323. )
  324. startup_nodes = passed_nodes
  325. else:
  326. startup_nodes = []
  327. if host and port:
  328. startup_nodes.append(ClusterNode(host, port, **self.connection_kwargs))
  329. self.nodes_manager = NodesManager(
  330. startup_nodes,
  331. require_full_coverage,
  332. kwargs,
  333. address_remap=address_remap,
  334. )
  335. self.encoder = Encoder(encoding, encoding_errors, decode_responses)
  336. self.read_from_replicas = read_from_replicas
  337. self.reinitialize_steps = reinitialize_steps
  338. self.cluster_error_retry_attempts = cluster_error_retry_attempts
  339. self.connection_error_retry_attempts = connection_error_retry_attempts
  340. self.reinitialize_counter = 0
  341. self.commands_parser = AsyncCommandsParser()
  342. self.node_flags = self.__class__.NODE_FLAGS.copy()
  343. self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  344. self.response_callbacks = kwargs["response_callbacks"]
  345. self.result_callbacks = self.__class__.RESULT_CALLBACKS.copy()
  346. self.result_callbacks[
  347. "CLUSTER SLOTS"
  348. ] = lambda cmd, res, **kwargs: parse_cluster_slots(
  349. list(res.values())[0], **kwargs
  350. )
  351. self._initialize = True
  352. self._lock: Optional[asyncio.Lock] = None
  353. async def initialize(self) -> "RedisCluster":
  354. """Get all nodes from startup nodes & creates connections if not initialized."""
  355. if self._initialize:
  356. if not self._lock:
  357. self._lock = asyncio.Lock()
  358. async with self._lock:
  359. if self._initialize:
  360. try:
  361. await self.nodes_manager.initialize()
  362. await self.commands_parser.initialize(
  363. self.nodes_manager.default_node
  364. )
  365. self._initialize = False
  366. except BaseException:
  367. await self.nodes_manager.aclose()
  368. await self.nodes_manager.aclose("startup_nodes")
  369. raise
  370. return self
  371. async def aclose(self) -> None:
  372. """Close all connections & client if initialized."""
  373. if not self._initialize:
  374. if not self._lock:
  375. self._lock = asyncio.Lock()
  376. async with self._lock:
  377. if not self._initialize:
  378. self._initialize = True
  379. await self.nodes_manager.aclose()
  380. await self.nodes_manager.aclose("startup_nodes")
  381. @deprecated_function(version="5.0.0", reason="Use aclose() instead", name="close")
  382. async def close(self) -> None:
  383. """alias for aclose() for backwards compatibility"""
  384. await self.aclose()
  385. async def __aenter__(self) -> "RedisCluster":
  386. return await self.initialize()
  387. async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  388. await self.aclose()
  389. def __await__(self) -> Generator[Any, None, "RedisCluster"]:
  390. return self.initialize().__await__()
  391. _DEL_MESSAGE = "Unclosed RedisCluster client"
  392. def __del__(
  393. self,
  394. _warn: Any = warnings.warn,
  395. _grl: Any = asyncio.get_running_loop,
  396. ) -> None:
  397. if hasattr(self, "_initialize") and not self._initialize:
  398. _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
  399. try:
  400. context = {"client": self, "message": self._DEL_MESSAGE}
  401. _grl().call_exception_handler(context)
  402. except RuntimeError:
  403. pass
  404. async def on_connect(self, connection: Connection) -> None:
  405. await connection.on_connect()
  406. # Sending READONLY command to server to configure connection as
  407. # readonly. Since each cluster node may change its server type due
  408. # to a failover, we should establish a READONLY connection
  409. # regardless of the server type. If this is a primary connection,
  410. # READONLY would not affect executing write commands.
  411. await connection.send_command("READONLY")
  412. if str_if_bytes(await connection.read_response()) != "OK":
  413. raise ConnectionError("READONLY command failed")
  414. def get_nodes(self) -> List["ClusterNode"]:
  415. """Get all nodes of the cluster."""
  416. return list(self.nodes_manager.nodes_cache.values())
  417. def get_primaries(self) -> List["ClusterNode"]:
  418. """Get the primary nodes of the cluster."""
  419. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  420. def get_replicas(self) -> List["ClusterNode"]:
  421. """Get the replica nodes of the cluster."""
  422. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  423. def get_random_node(self) -> "ClusterNode":
  424. """Get a random node of the cluster."""
  425. return random.choice(list(self.nodes_manager.nodes_cache.values()))
  426. def get_default_node(self) -> "ClusterNode":
  427. """Get the default node of the client."""
  428. return self.nodes_manager.default_node
  429. def set_default_node(self, node: "ClusterNode") -> None:
  430. """
  431. Set the default node of the client.
  432. :raises DataError: if None is passed or node does not exist in cluster.
  433. """
  434. if not node or not self.get_node(node_name=node.name):
  435. raise DataError("The requested node does not exist in the cluster.")
  436. self.nodes_manager.default_node = node
  437. def get_node(
  438. self,
  439. host: Optional[str] = None,
  440. port: Optional[int] = None,
  441. node_name: Optional[str] = None,
  442. ) -> Optional["ClusterNode"]:
  443. """Get node by (host, port) or node_name."""
  444. return self.nodes_manager.get_node(host, port, node_name)
  445. def get_node_from_key(
  446. self, key: str, replica: bool = False
  447. ) -> Optional["ClusterNode"]:
  448. """
  449. Get the cluster node corresponding to the provided key.
  450. :param key:
  451. :param replica:
  452. | Indicates if a replica should be returned
  453. |
  454. None will returned if no replica holds this key
  455. :raises SlotNotCoveredError: if the key is not covered by any slot.
  456. """
  457. slot = self.keyslot(key)
  458. slot_cache = self.nodes_manager.slots_cache.get(slot)
  459. if not slot_cache:
  460. raise SlotNotCoveredError(f'Slot "{slot}" is not covered by the cluster.')
  461. if replica:
  462. if len(self.nodes_manager.slots_cache[slot]) < 2:
  463. return None
  464. node_idx = 1
  465. else:
  466. node_idx = 0
  467. return slot_cache[node_idx]
  468. def keyslot(self, key: EncodableT) -> int:
  469. """
  470. Find the keyslot for a given key.
  471. See: https://redis.io/docs/manual/scaling/#redis-cluster-data-sharding
  472. """
  473. return key_slot(self.encoder.encode(key))
  474. def get_encoder(self) -> Encoder:
  475. """Get the encoder object of the client."""
  476. return self.encoder
  477. def get_connection_kwargs(self) -> Dict[str, Optional[Any]]:
  478. """Get the kwargs passed to :class:`~redis.asyncio.connection.Connection`."""
  479. return self.connection_kwargs
  480. def get_retry(self) -> Optional["Retry"]:
  481. return self.retry
  482. def set_retry(self, retry: "Retry") -> None:
  483. self.retry = retry
  484. for node in self.get_nodes():
  485. node.connection_kwargs.update({"retry": retry})
  486. for conn in node._connections:
  487. conn.retry = retry
  488. def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
  489. """Set a custom response callback."""
  490. self.response_callbacks[command] = callback
  491. async def _determine_nodes(
  492. self, command: str, *args: Any, node_flag: Optional[str] = None
  493. ) -> List["ClusterNode"]:
  494. # Determine which nodes should be executed the command on.
  495. # Returns a list of target nodes.
  496. if not node_flag:
  497. # get the nodes group for this command if it was predefined
  498. node_flag = self.command_flags.get(command)
  499. if node_flag in self.node_flags:
  500. if node_flag == self.__class__.DEFAULT_NODE:
  501. # return the cluster's default node
  502. return [self.nodes_manager.default_node]
  503. if node_flag == self.__class__.PRIMARIES:
  504. # return all primaries
  505. return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
  506. if node_flag == self.__class__.REPLICAS:
  507. # return all replicas
  508. return self.nodes_manager.get_nodes_by_server_type(REPLICA)
  509. if node_flag == self.__class__.ALL_NODES:
  510. # return all nodes
  511. return list(self.nodes_manager.nodes_cache.values())
  512. if node_flag == self.__class__.RANDOM:
  513. # return a random node
  514. return [random.choice(list(self.nodes_manager.nodes_cache.values()))]
  515. # get the node that holds the key's slot
  516. return [
  517. self.nodes_manager.get_node_from_slot(
  518. await self._determine_slot(command, *args),
  519. self.read_from_replicas and command in READ_COMMANDS,
  520. )
  521. ]
  522. async def _determine_slot(self, command: str, *args: Any) -> int:
  523. if self.command_flags.get(command) == SLOT_ID:
  524. # The command contains the slot ID
  525. return int(args[0])
  526. # Get the keys in the command
  527. # EVAL and EVALSHA are common enough that it's wasteful to go to the
  528. # redis server to parse the keys. Besides, there is a bug in redis<7.0
  529. # where `self._get_command_keys()` fails anyway. So, we special case
  530. # EVAL/EVALSHA.
  531. # - issue: https://github.com/redis/redis/issues/9493
  532. # - fix: https://github.com/redis/redis/pull/9733
  533. if command.upper() in ("EVAL", "EVALSHA"):
  534. # command syntax: EVAL "script body" num_keys ...
  535. if len(args) < 2:
  536. raise RedisClusterException(
  537. f"Invalid args in command: {command, *args}"
  538. )
  539. keys = args[2 : 2 + int(args[1])]
  540. # if there are 0 keys, that means the script can be run on any node
  541. # so we can just return a random slot
  542. if not keys:
  543. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  544. else:
  545. keys = await self.commands_parser.get_keys(command, *args)
  546. if not keys:
  547. # FCALL can call a function with 0 keys, that means the function
  548. # can be run on any node so we can just return a random slot
  549. if command.upper() in ("FCALL", "FCALL_RO"):
  550. return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
  551. raise RedisClusterException(
  552. "No way to dispatch this command to Redis Cluster. "
  553. "Missing key.\nYou can execute the command by specifying "
  554. f"target nodes.\nCommand: {args}"
  555. )
  556. # single key command
  557. if len(keys) == 1:
  558. return self.keyslot(keys[0])
  559. # multi-key command; we need to make sure all keys are mapped to
  560. # the same slot
  561. slots = {self.keyslot(key) for key in keys}
  562. if len(slots) != 1:
  563. raise RedisClusterException(
  564. f"{command} - all keys must map to the same key slot"
  565. )
  566. return slots.pop()
  567. def _is_node_flag(self, target_nodes: Any) -> bool:
  568. return isinstance(target_nodes, str) and target_nodes in self.node_flags
  569. def _parse_target_nodes(self, target_nodes: Any) -> List["ClusterNode"]:
  570. if isinstance(target_nodes, list):
  571. nodes = target_nodes
  572. elif isinstance(target_nodes, ClusterNode):
  573. # Supports passing a single ClusterNode as a variable
  574. nodes = [target_nodes]
  575. elif isinstance(target_nodes, dict):
  576. # Supports dictionaries of the format {node_name: node}.
  577. # It enables to execute commands with multi nodes as follows:
  578. # rc.cluster_save_config(rc.get_primaries())
  579. nodes = list(target_nodes.values())
  580. else:
  581. raise TypeError(
  582. "target_nodes type can be one of the following: "
  583. "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
  584. "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
  585. f"The passed type is {type(target_nodes)}"
  586. )
  587. return nodes
  588. async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  589. """
  590. Execute a raw command on the appropriate cluster node or target_nodes.
  591. It will retry the command as specified by :attr:`cluster_error_retry_attempts` &
  592. then raise an exception.
  593. :param args:
  594. | Raw command args
  595. :param kwargs:
  596. - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
  597. or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
  598. - Rest of the kwargs are passed to the Redis connection
  599. :raises RedisClusterException: if target_nodes is not provided & the command
  600. can't be mapped to a slot
  601. """
  602. command = args[0]
  603. target_nodes = []
  604. target_nodes_specified = False
  605. retry_attempts = self.cluster_error_retry_attempts
  606. passed_targets = kwargs.pop("target_nodes", None)
  607. if passed_targets and not self._is_node_flag(passed_targets):
  608. target_nodes = self._parse_target_nodes(passed_targets)
  609. target_nodes_specified = True
  610. retry_attempts = 0
  611. # Add one for the first execution
  612. execute_attempts = 1 + retry_attempts
  613. for _ in range(execute_attempts):
  614. if self._initialize:
  615. await self.initialize()
  616. if (
  617. len(target_nodes) == 1
  618. and target_nodes[0] == self.get_default_node()
  619. ):
  620. # Replace the default cluster node
  621. self.replace_default_node()
  622. try:
  623. if not target_nodes_specified:
  624. # Determine the nodes to execute the command on
  625. target_nodes = await self._determine_nodes(
  626. *args, node_flag=passed_targets
  627. )
  628. if not target_nodes:
  629. raise RedisClusterException(
  630. f"No targets were found to execute {args} command on"
  631. )
  632. if len(target_nodes) == 1:
  633. # Return the processed result
  634. ret = await self._execute_command(target_nodes[0], *args, **kwargs)
  635. if command in self.result_callbacks:
  636. return self.result_callbacks[command](
  637. command, {target_nodes[0].name: ret}, **kwargs
  638. )
  639. return ret
  640. else:
  641. keys = [node.name for node in target_nodes]
  642. values = await asyncio.gather(
  643. *(
  644. asyncio.create_task(
  645. self._execute_command(node, *args, **kwargs)
  646. )
  647. for node in target_nodes
  648. )
  649. )
  650. if command in self.result_callbacks:
  651. return self.result_callbacks[command](
  652. command, dict(zip(keys, values)), **kwargs
  653. )
  654. return dict(zip(keys, values))
  655. except Exception as e:
  656. if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  657. # The nodes and slots cache were should be reinitialized.
  658. # Try again with the new cluster setup.
  659. retry_attempts -= 1
  660. continue
  661. else:
  662. # raise the exception
  663. raise e
  664. async def _execute_command(
  665. self, target_node: "ClusterNode", *args: Union[KeyT, EncodableT], **kwargs: Any
  666. ) -> Any:
  667. asking = moved = False
  668. redirect_addr = None
  669. ttl = self.RedisClusterRequestTTL
  670. while ttl > 0:
  671. ttl -= 1
  672. try:
  673. if asking:
  674. target_node = self.get_node(node_name=redirect_addr)
  675. await target_node.execute_command("ASKING")
  676. asking = False
  677. elif moved:
  678. # MOVED occurred and the slots cache was updated,
  679. # refresh the target node
  680. slot = await self._determine_slot(*args)
  681. target_node = self.nodes_manager.get_node_from_slot(
  682. slot, self.read_from_replicas and args[0] in READ_COMMANDS
  683. )
  684. moved = False
  685. return await target_node.execute_command(*args, **kwargs)
  686. except (BusyLoadingError, MaxConnectionsError):
  687. raise
  688. except (ConnectionError, TimeoutError):
  689. # Connection retries are being handled in the node's
  690. # Retry object.
  691. # Remove the failed node from the startup nodes before we try
  692. # to reinitialize the cluster
  693. self.nodes_manager.startup_nodes.pop(target_node.name, None)
  694. # Hard force of reinitialize of the node/slots setup
  695. # and try again with the new setup
  696. await self.aclose()
  697. raise
  698. except ClusterDownError:
  699. # ClusterDownError can occur during a failover and to get
  700. # self-healed, we will try to reinitialize the cluster layout
  701. # and retry executing the command
  702. await self.aclose()
  703. await asyncio.sleep(0.25)
  704. raise
  705. except MovedError as e:
  706. # First, we will try to patch the slots/nodes cache with the
  707. # redirected node output and try again. If MovedError exceeds
  708. # 'reinitialize_steps' number of times, we will force
  709. # reinitializing the tables, and then try again.
  710. # 'reinitialize_steps' counter will increase faster when
  711. # the same client object is shared between multiple threads. To
  712. # reduce the frequency you can set this variable in the
  713. # RedisCluster constructor.
  714. self.reinitialize_counter += 1
  715. if (
  716. self.reinitialize_steps
  717. and self.reinitialize_counter % self.reinitialize_steps == 0
  718. ):
  719. await self.aclose()
  720. # Reset the counter
  721. self.reinitialize_counter = 0
  722. else:
  723. self.nodes_manager._moved_exception = e
  724. moved = True
  725. except AskError as e:
  726. redirect_addr = get_node_name(host=e.host, port=e.port)
  727. asking = True
  728. except TryAgainError:
  729. if ttl < self.RedisClusterRequestTTL / 2:
  730. await asyncio.sleep(0.05)
  731. raise ClusterError("TTL exhausted.")
  732. def pipeline(
  733. self, transaction: Optional[Any] = None, shard_hint: Optional[Any] = None
  734. ) -> "ClusterPipeline":
  735. """
  736. Create & return a new :class:`~.ClusterPipeline` object.
  737. Cluster implementation of pipeline does not support transaction or shard_hint.
  738. :raises RedisClusterException: if transaction or shard_hint are truthy values
  739. """
  740. if shard_hint:
  741. raise RedisClusterException("shard_hint is deprecated in cluster mode")
  742. if transaction:
  743. raise RedisClusterException("transaction is deprecated in cluster mode")
  744. return ClusterPipeline(self)
  745. def lock(
  746. self,
  747. name: KeyT,
  748. timeout: Optional[float] = None,
  749. sleep: float = 0.1,
  750. blocking: bool = True,
  751. blocking_timeout: Optional[float] = None,
  752. lock_class: Optional[Type[Lock]] = None,
  753. thread_local: bool = True,
  754. ) -> Lock:
  755. """
  756. Return a new Lock object using key ``name`` that mimics
  757. the behavior of threading.Lock.
  758. If specified, ``timeout`` indicates a maximum life for the lock.
  759. By default, it will remain locked until release() is called.
  760. ``sleep`` indicates the amount of time to sleep per loop iteration
  761. when the lock is in blocking mode and another client is currently
  762. holding the lock.
  763. ``blocking`` indicates whether calling ``acquire`` should block until
  764. the lock has been acquired or to fail immediately, causing ``acquire``
  765. to return False and the lock not being acquired. Defaults to True.
  766. Note this value can be overridden by passing a ``blocking``
  767. argument to ``acquire``.
  768. ``blocking_timeout`` indicates the maximum amount of time in seconds to
  769. spend trying to acquire the lock. A value of ``None`` indicates
  770. continue trying forever. ``blocking_timeout`` can be specified as a
  771. float or integer, both representing the number of seconds to wait.
  772. ``lock_class`` forces the specified lock implementation. Note that as
  773. of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
  774. a Lua-based lock). So, it's unlikely you'll need this parameter, unless
  775. you have created your own custom lock class.
  776. ``thread_local`` indicates whether the lock token is placed in
  777. thread-local storage. By default, the token is placed in thread local
  778. storage so that a thread only sees its token, not a token set by
  779. another thread. Consider the following timeline:
  780. time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
  781. thread-1 sets the token to "abc"
  782. time: 1, thread-2 blocks trying to acquire `my-lock` using the
  783. Lock instance.
  784. time: 5, thread-1 has not yet completed. redis expires the lock
  785. key.
  786. time: 5, thread-2 acquired `my-lock` now that it's available.
  787. thread-2 sets the token to "xyz"
  788. time: 6, thread-1 finishes its work and calls release(). if the
  789. token is *not* stored in thread local storage, then
  790. thread-1 would see the token value as "xyz" and would be
  791. able to successfully release the thread-2's lock.
  792. In some use cases it's necessary to disable thread local storage. For
  793. example, if you have code where one thread acquires a lock and passes
  794. that lock instance to a worker thread to release later. If thread
  795. local storage isn't disabled in this case, the worker thread won't see
  796. the token set by the thread that acquired the lock. Our assumption
  797. is that these cases aren't common and as such default to using
  798. thread local storage."""
  799. if lock_class is None:
  800. lock_class = Lock
  801. return lock_class(
  802. self,
  803. name,
  804. timeout=timeout,
  805. sleep=sleep,
  806. blocking=blocking,
  807. blocking_timeout=blocking_timeout,
  808. thread_local=thread_local,
  809. )
  810. class ClusterNode:
  811. """
  812. Create a new ClusterNode.
  813. Each ClusterNode manages multiple :class:`~redis.asyncio.connection.Connection`
  814. objects for the (host, port).
  815. """
  816. __slots__ = (
  817. "_connections",
  818. "_free",
  819. "connection_class",
  820. "connection_kwargs",
  821. "host",
  822. "max_connections",
  823. "name",
  824. "port",
  825. "response_callbacks",
  826. "server_type",
  827. )
  828. def __init__(
  829. self,
  830. host: str,
  831. port: Union[str, int],
  832. server_type: Optional[str] = None,
  833. *,
  834. max_connections: int = 2**31,
  835. connection_class: Type[Connection] = Connection,
  836. **connection_kwargs: Any,
  837. ) -> None:
  838. if host == "localhost":
  839. host = socket.gethostbyname(host)
  840. connection_kwargs["host"] = host
  841. connection_kwargs["port"] = port
  842. self.host = host
  843. self.port = port
  844. self.name = get_node_name(host, port)
  845. self.server_type = server_type
  846. self.max_connections = max_connections
  847. self.connection_class = connection_class
  848. self.connection_kwargs = connection_kwargs
  849. self.response_callbacks = connection_kwargs.pop("response_callbacks", {})
  850. self._connections: List[Connection] = []
  851. self._free: Deque[Connection] = collections.deque(maxlen=self.max_connections)
  852. def __repr__(self) -> str:
  853. return (
  854. f"[host={self.host}, port={self.port}, "
  855. f"name={self.name}, server_type={self.server_type}]"
  856. )
  857. def __eq__(self, obj: Any) -> bool:
  858. return isinstance(obj, ClusterNode) and obj.name == self.name
  859. _DEL_MESSAGE = "Unclosed ClusterNode object"
  860. def __del__(
  861. self,
  862. _warn: Any = warnings.warn,
  863. _grl: Any = asyncio.get_running_loop,
  864. ) -> None:
  865. for connection in self._connections:
  866. if connection.is_connected:
  867. _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
  868. try:
  869. context = {"client": self, "message": self._DEL_MESSAGE}
  870. _grl().call_exception_handler(context)
  871. except RuntimeError:
  872. pass
  873. break
  874. async def disconnect(self) -> None:
  875. ret = await asyncio.gather(
  876. *(
  877. asyncio.create_task(connection.disconnect())
  878. for connection in self._connections
  879. ),
  880. return_exceptions=True,
  881. )
  882. exc = next((res for res in ret if isinstance(res, Exception)), None)
  883. if exc:
  884. raise exc
  885. def acquire_connection(self) -> Connection:
  886. try:
  887. return self._free.popleft()
  888. except IndexError:
  889. if len(self._connections) < self.max_connections:
  890. connection = self.connection_class(**self.connection_kwargs)
  891. self._connections.append(connection)
  892. return connection
  893. raise MaxConnectionsError()
  894. async def parse_response(
  895. self, connection: Connection, command: str, **kwargs: Any
  896. ) -> Any:
  897. try:
  898. if NEVER_DECODE in kwargs:
  899. response = await connection.read_response(disable_decoding=True)
  900. kwargs.pop(NEVER_DECODE)
  901. else:
  902. response = await connection.read_response()
  903. except ResponseError:
  904. if EMPTY_RESPONSE in kwargs:
  905. return kwargs[EMPTY_RESPONSE]
  906. raise
  907. if EMPTY_RESPONSE in kwargs:
  908. kwargs.pop(EMPTY_RESPONSE)
  909. # Return response
  910. if command in self.response_callbacks:
  911. return self.response_callbacks[command](response, **kwargs)
  912. return response
  913. async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
  914. # Acquire connection
  915. connection = self.acquire_connection()
  916. # Execute command
  917. await connection.send_packed_command(connection.pack_command(*args), False)
  918. # Read response
  919. try:
  920. return await self.parse_response(connection, args[0], **kwargs)
  921. finally:
  922. # Release connection
  923. self._free.append(connection)
  924. async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
  925. # Acquire connection
  926. connection = self.acquire_connection()
  927. # Execute command
  928. await connection.send_packed_command(
  929. connection.pack_commands(cmd.args for cmd in commands), False
  930. )
  931. # Read responses
  932. ret = False
  933. for cmd in commands:
  934. try:
  935. cmd.result = await self.parse_response(
  936. connection, cmd.args[0], **cmd.kwargs
  937. )
  938. except Exception as e:
  939. cmd.result = e
  940. ret = True
  941. # Release connection
  942. self._free.append(connection)
  943. return ret
  944. class NodesManager:
  945. __slots__ = (
  946. "_moved_exception",
  947. "connection_kwargs",
  948. "default_node",
  949. "nodes_cache",
  950. "read_load_balancer",
  951. "require_full_coverage",
  952. "slots_cache",
  953. "startup_nodes",
  954. "address_remap",
  955. )
  956. def __init__(
  957. self,
  958. startup_nodes: List["ClusterNode"],
  959. require_full_coverage: bool,
  960. connection_kwargs: Dict[str, Any],
  961. address_remap: Optional[Callable[[str, int], Tuple[str, int]]] = None,
  962. ) -> None:
  963. self.startup_nodes = {node.name: node for node in startup_nodes}
  964. self.require_full_coverage = require_full_coverage
  965. self.connection_kwargs = connection_kwargs
  966. self.address_remap = address_remap
  967. self.default_node: "ClusterNode" = None
  968. self.nodes_cache: Dict[str, "ClusterNode"] = {}
  969. self.slots_cache: Dict[int, List["ClusterNode"]] = {}
  970. self.read_load_balancer = LoadBalancer()
  971. self._moved_exception: MovedError = None
  972. def get_node(
  973. self,
  974. host: Optional[str] = None,
  975. port: Optional[int] = None,
  976. node_name: Optional[str] = None,
  977. ) -> Optional["ClusterNode"]:
  978. if host and port:
  979. # the user passed host and port
  980. if host == "localhost":
  981. host = socket.gethostbyname(host)
  982. return self.nodes_cache.get(get_node_name(host=host, port=port))
  983. elif node_name:
  984. return self.nodes_cache.get(node_name)
  985. else:
  986. raise DataError(
  987. "get_node requires one of the following: "
  988. "1. node name "
  989. "2. host and port"
  990. )
  991. def set_nodes(
  992. self,
  993. old: Dict[str, "ClusterNode"],
  994. new: Dict[str, "ClusterNode"],
  995. remove_old: bool = False,
  996. ) -> None:
  997. if remove_old:
  998. for name in list(old.keys()):
  999. if name not in new:
  1000. task = asyncio.create_task(old.pop(name).disconnect()) # noqa
  1001. for name, node in new.items():
  1002. if name in old:
  1003. if old[name] is node:
  1004. continue
  1005. task = asyncio.create_task(old[name].disconnect()) # noqa
  1006. old[name] = node
  1007. def _update_moved_slots(self) -> None:
  1008. e = self._moved_exception
  1009. redirected_node = self.get_node(host=e.host, port=e.port)
  1010. if redirected_node:
  1011. # The node already exists
  1012. if redirected_node.server_type != PRIMARY:
  1013. # Update the node's server type
  1014. redirected_node.server_type = PRIMARY
  1015. else:
  1016. # This is a new node, we will add it to the nodes cache
  1017. redirected_node = ClusterNode(
  1018. e.host, e.port, PRIMARY, **self.connection_kwargs
  1019. )
  1020. self.set_nodes(self.nodes_cache, {redirected_node.name: redirected_node})
  1021. if redirected_node in self.slots_cache[e.slot_id]:
  1022. # The MOVED error resulted from a failover, and the new slot owner
  1023. # had previously been a replica.
  1024. old_primary = self.slots_cache[e.slot_id][0]
  1025. # Update the old primary to be a replica and add it to the end of
  1026. # the slot's node list
  1027. old_primary.server_type = REPLICA
  1028. self.slots_cache[e.slot_id].append(old_primary)
  1029. # Remove the old replica, which is now a primary, from the slot's
  1030. # node list
  1031. self.slots_cache[e.slot_id].remove(redirected_node)
  1032. # Override the old primary with the new one
  1033. self.slots_cache[e.slot_id][0] = redirected_node
  1034. if self.default_node == old_primary:
  1035. # Update the default node with the new primary
  1036. self.default_node = redirected_node
  1037. else:
  1038. # The new slot owner is a new server, or a server from a different
  1039. # shard. We need to remove all current nodes from the slot's list
  1040. # (including replications) and add just the new node.
  1041. self.slots_cache[e.slot_id] = [redirected_node]
  1042. # Reset moved_exception
  1043. self._moved_exception = None
  1044. def get_node_from_slot(
  1045. self, slot: int, read_from_replicas: bool = False
  1046. ) -> "ClusterNode":
  1047. if self._moved_exception:
  1048. self._update_moved_slots()
  1049. try:
  1050. if read_from_replicas:
  1051. # get the server index in a Round-Robin manner
  1052. primary_name = self.slots_cache[slot][0].name
  1053. node_idx = self.read_load_balancer.get_server_index(
  1054. primary_name, len(self.slots_cache[slot])
  1055. )
  1056. return self.slots_cache[slot][node_idx]
  1057. return self.slots_cache[slot][0]
  1058. except (IndexError, TypeError):
  1059. raise SlotNotCoveredError(
  1060. f'Slot "{slot}" not covered by the cluster. '
  1061. f'"require_full_coverage={self.require_full_coverage}"'
  1062. )
  1063. def get_nodes_by_server_type(self, server_type: str) -> List["ClusterNode"]:
  1064. return [
  1065. node
  1066. for node in self.nodes_cache.values()
  1067. if node.server_type == server_type
  1068. ]
  1069. async def initialize(self) -> None:
  1070. self.read_load_balancer.reset()
  1071. tmp_nodes_cache: Dict[str, "ClusterNode"] = {}
  1072. tmp_slots: Dict[int, List["ClusterNode"]] = {}
  1073. disagreements = []
  1074. startup_nodes_reachable = False
  1075. fully_covered = False
  1076. exception = None
  1077. for startup_node in self.startup_nodes.values():
  1078. try:
  1079. # Make sure cluster mode is enabled on this node
  1080. if not (await startup_node.execute_command("INFO")).get(
  1081. "cluster_enabled"
  1082. ):
  1083. raise RedisClusterException(
  1084. "Cluster mode is not enabled on this node"
  1085. )
  1086. cluster_slots = await startup_node.execute_command("CLUSTER SLOTS")
  1087. startup_nodes_reachable = True
  1088. except Exception as e:
  1089. # Try the next startup node.
  1090. # The exception is saved and raised only if we have no more nodes.
  1091. exception = e
  1092. continue
  1093. # CLUSTER SLOTS command results in the following output:
  1094. # [[slot_section[from_slot,to_slot,master,replica1,...,replicaN]]]
  1095. # where each node contains the following list: [IP, port, node_id]
  1096. # Therefore, cluster_slots[0][2][0] will be the IP address of the
  1097. # primary node of the first slot section.
  1098. # If there's only one server in the cluster, its ``host`` is ''
  1099. # Fix it to the host in startup_nodes
  1100. if (
  1101. len(cluster_slots) == 1
  1102. and not cluster_slots[0][2][0]
  1103. and len(self.startup_nodes) == 1
  1104. ):
  1105. cluster_slots[0][2][0] = startup_node.host
  1106. for slot in cluster_slots:
  1107. for i in range(2, len(slot)):
  1108. slot[i] = [str_if_bytes(val) for val in slot[i]]
  1109. primary_node = slot[2]
  1110. host = primary_node[0]
  1111. if host == "":
  1112. host = startup_node.host
  1113. port = int(primary_node[1])
  1114. host, port = self.remap_host_port(host, port)
  1115. target_node = tmp_nodes_cache.get(get_node_name(host, port))
  1116. if not target_node:
  1117. target_node = ClusterNode(
  1118. host, port, PRIMARY, **self.connection_kwargs
  1119. )
  1120. # add this node to the nodes cache
  1121. tmp_nodes_cache[target_node.name] = target_node
  1122. for i in range(int(slot[0]), int(slot[1]) + 1):
  1123. if i not in tmp_slots:
  1124. tmp_slots[i] = []
  1125. tmp_slots[i].append(target_node)
  1126. replica_nodes = [slot[j] for j in range(3, len(slot))]
  1127. for replica_node in replica_nodes:
  1128. host = replica_node[0]
  1129. port = replica_node[1]
  1130. host, port = self.remap_host_port(host, port)
  1131. target_replica_node = tmp_nodes_cache.get(
  1132. get_node_name(host, port)
  1133. )
  1134. if not target_replica_node:
  1135. target_replica_node = ClusterNode(
  1136. host, port, REPLICA, **self.connection_kwargs
  1137. )
  1138. tmp_slots[i].append(target_replica_node)
  1139. # add this node to the nodes cache
  1140. tmp_nodes_cache[
  1141. target_replica_node.name
  1142. ] = target_replica_node
  1143. else:
  1144. # Validate that 2 nodes want to use the same slot cache
  1145. # setup
  1146. tmp_slot = tmp_slots[i][0]
  1147. if tmp_slot.name != target_node.name:
  1148. disagreements.append(
  1149. f"{tmp_slot.name} vs {target_node.name} on slot: {i}"
  1150. )
  1151. if len(disagreements) > 5:
  1152. raise RedisClusterException(
  1153. f"startup_nodes could not agree on a valid "
  1154. f'slots cache: {", ".join(disagreements)}'
  1155. )
  1156. # Validate if all slots are covered or if we should try next startup node
  1157. fully_covered = True
  1158. for i in range(REDIS_CLUSTER_HASH_SLOTS):
  1159. if i not in tmp_slots:
  1160. fully_covered = False
  1161. break
  1162. if fully_covered:
  1163. break
  1164. if not startup_nodes_reachable:
  1165. raise RedisClusterException(
  1166. f"Redis Cluster cannot be connected. Please provide at least "
  1167. f"one reachable node: {str(exception)}"
  1168. ) from exception
  1169. # Check if the slots are not fully covered
  1170. if not fully_covered and self.require_full_coverage:
  1171. # Despite the requirement that the slots be covered, there
  1172. # isn't a full coverage
  1173. raise RedisClusterException(
  1174. f"All slots are not covered after query all startup_nodes. "
  1175. f"{len(tmp_slots)} of {REDIS_CLUSTER_HASH_SLOTS} "
  1176. f"covered..."
  1177. )
  1178. # Set the tmp variables to the real variables
  1179. self.slots_cache = tmp_slots
  1180. self.set_nodes(self.nodes_cache, tmp_nodes_cache, remove_old=True)
  1181. # Populate the startup nodes with all discovered nodes
  1182. self.set_nodes(self.startup_nodes, self.nodes_cache, remove_old=True)
  1183. # Set the default node
  1184. self.default_node = self.get_nodes_by_server_type(PRIMARY)[0]
  1185. # If initialize was called after a MovedError, clear it
  1186. self._moved_exception = None
  1187. async def aclose(self, attr: str = "nodes_cache") -> None:
  1188. self.default_node = None
  1189. await asyncio.gather(
  1190. *(
  1191. asyncio.create_task(node.disconnect())
  1192. for node in getattr(self, attr).values()
  1193. )
  1194. )
  1195. def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
  1196. """
  1197. Remap the host and port returned from the cluster to a different
  1198. internal value. Useful if the client is not connecting directly
  1199. to the cluster.
  1200. """
  1201. if self.address_remap:
  1202. return self.address_remap((host, port))
  1203. return host, port
  1204. class ClusterPipeline(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
  1205. """
  1206. Create a new ClusterPipeline object.
  1207. Usage::
  1208. result = await (
  1209. rc.pipeline()
  1210. .set("A", 1)
  1211. .get("A")
  1212. .hset("K", "F", "V")
  1213. .hgetall("K")
  1214. .mset_nonatomic({"A": 2, "B": 3})
  1215. .get("A")
  1216. .get("B")
  1217. .delete("A", "B", "K")
  1218. .execute()
  1219. )
  1220. # result = [True, "1", 1, {"F": "V"}, True, True, "2", "3", 1, 1, 1]
  1221. Note: For commands `DELETE`, `EXISTS`, `TOUCH`, `UNLINK`, `mset_nonatomic`, which
  1222. are split across multiple nodes, you'll get multiple results for them in the array.
  1223. Retryable errors:
  1224. - :class:`~.ClusterDownError`
  1225. - :class:`~.ConnectionError`
  1226. - :class:`~.TimeoutError`
  1227. Redirection errors:
  1228. - :class:`~.TryAgainError`
  1229. - :class:`~.MovedError`
  1230. - :class:`~.AskError`
  1231. :param client:
  1232. | Existing :class:`~.RedisCluster` client
  1233. """
  1234. __slots__ = ("_command_stack", "_client")
  1235. def __init__(self, client: RedisCluster) -> None:
  1236. self._client = client
  1237. self._command_stack: List["PipelineCommand"] = []
  1238. async def initialize(self) -> "ClusterPipeline":
  1239. if self._client._initialize:
  1240. await self._client.initialize()
  1241. self._command_stack = []
  1242. return self
  1243. async def __aenter__(self) -> "ClusterPipeline":
  1244. return await self.initialize()
  1245. async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  1246. self._command_stack = []
  1247. def __await__(self) -> Generator[Any, None, "ClusterPipeline"]:
  1248. return self.initialize().__await__()
  1249. def __enter__(self) -> "ClusterPipeline":
  1250. self._command_stack = []
  1251. return self
  1252. def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
  1253. self._command_stack = []
  1254. def __bool__(self) -> bool:
  1255. return bool(self._command_stack)
  1256. def __len__(self) -> int:
  1257. return len(self._command_stack)
  1258. def execute_command(
  1259. self, *args: Union[KeyT, EncodableT], **kwargs: Any
  1260. ) -> "ClusterPipeline":
  1261. """
  1262. Append a raw command to the pipeline.
  1263. :param args:
  1264. | Raw command args
  1265. :param kwargs:
  1266. - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
  1267. or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
  1268. - Rest of the kwargs are passed to the Redis connection
  1269. """
  1270. self._command_stack.append(
  1271. PipelineCommand(len(self._command_stack), *args, **kwargs)
  1272. )
  1273. return self
  1274. async def execute(
  1275. self, raise_on_error: bool = True, allow_redirections: bool = True
  1276. ) -> List[Any]:
  1277. """
  1278. Execute the pipeline.
  1279. It will retry the commands as specified by :attr:`cluster_error_retry_attempts`
  1280. & then raise an exception.
  1281. :param raise_on_error:
  1282. | Raise the first error if there are any errors
  1283. :param allow_redirections:
  1284. | Whether to retry each failed command individually in case of redirection
  1285. errors
  1286. :raises RedisClusterException: if target_nodes is not provided & the command
  1287. can't be mapped to a slot
  1288. """
  1289. if not self._command_stack:
  1290. return []
  1291. try:
  1292. for _ in range(self._client.cluster_error_retry_attempts):
  1293. if self._client._initialize:
  1294. await self._client.initialize()
  1295. try:
  1296. return await self._execute(
  1297. self._client,
  1298. self._command_stack,
  1299. raise_on_error=raise_on_error,
  1300. allow_redirections=allow_redirections,
  1301. )
  1302. except BaseException as e:
  1303. if type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  1304. # Try again with the new cluster setup.
  1305. exception = e
  1306. await self._client.aclose()
  1307. await asyncio.sleep(0.25)
  1308. else:
  1309. # All other errors should be raised.
  1310. raise
  1311. # If it fails the configured number of times then raise an exception
  1312. raise exception
  1313. finally:
  1314. self._command_stack = []
  1315. async def _execute(
  1316. self,
  1317. client: "RedisCluster",
  1318. stack: List["PipelineCommand"],
  1319. raise_on_error: bool = True,
  1320. allow_redirections: bool = True,
  1321. ) -> List[Any]:
  1322. todo = [
  1323. cmd for cmd in stack if not cmd.result or isinstance(cmd.result, Exception)
  1324. ]
  1325. nodes = {}
  1326. for cmd in todo:
  1327. passed_targets = cmd.kwargs.pop("target_nodes", None)
  1328. if passed_targets and not client._is_node_flag(passed_targets):
  1329. target_nodes = client._parse_target_nodes(passed_targets)
  1330. else:
  1331. target_nodes = await client._determine_nodes(
  1332. *cmd.args, node_flag=passed_targets
  1333. )
  1334. if not target_nodes:
  1335. raise RedisClusterException(
  1336. f"No targets were found to execute {cmd.args} command on"
  1337. )
  1338. if len(target_nodes) > 1:
  1339. raise RedisClusterException(f"Too many targets for command {cmd.args}")
  1340. node = target_nodes[0]
  1341. if node.name not in nodes:
  1342. nodes[node.name] = (node, [])
  1343. nodes[node.name][1].append(cmd)
  1344. errors = await asyncio.gather(
  1345. *(
  1346. asyncio.create_task(node[0].execute_pipeline(node[1]))
  1347. for node in nodes.values()
  1348. )
  1349. )
  1350. if any(errors):
  1351. if allow_redirections:
  1352. # send each errored command individually
  1353. for cmd in todo:
  1354. if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
  1355. try:
  1356. cmd.result = await client.execute_command(
  1357. *cmd.args, **cmd.kwargs
  1358. )
  1359. except Exception as e:
  1360. cmd.result = e
  1361. if raise_on_error:
  1362. for cmd in todo:
  1363. result = cmd.result
  1364. if isinstance(result, Exception):
  1365. command = " ".join(map(safe_str, cmd.args))
  1366. msg = (
  1367. f"Command # {cmd.position + 1} ({command}) of pipeline "
  1368. f"caused error: {result.args}"
  1369. )
  1370. result.args = (msg,) + result.args[1:]
  1371. raise result
  1372. default_node = nodes.get(client.get_default_node().name)
  1373. if default_node is not None:
  1374. # This pipeline execution used the default node, check if we need
  1375. # to replace it.
  1376. # Note: when the error is raised we'll reset the default node in the
  1377. # caller function.
  1378. for cmd in default_node[1]:
  1379. # Check if it has a command that failed with a relevant
  1380. # exception
  1381. if type(cmd.result) in self.__class__.ERRORS_ALLOW_RETRY:
  1382. client.replace_default_node()
  1383. break
  1384. return [cmd.result for cmd in stack]
  1385. def _split_command_across_slots(
  1386. self, command: str, *keys: KeyT
  1387. ) -> "ClusterPipeline":
  1388. for slot_keys in self._client._partition_keys_by_slot(keys).values():
  1389. self.execute_command(command, *slot_keys)
  1390. return self
  1391. def mset_nonatomic(
  1392. self, mapping: Mapping[AnyKeyT, EncodableT]
  1393. ) -> "ClusterPipeline":
  1394. encoder = self._client.encoder
  1395. slots_pairs = {}
  1396. for pair in mapping.items():
  1397. slot = key_slot(encoder.encode(pair[0]))
  1398. slots_pairs.setdefault(slot, []).extend(pair)
  1399. for pairs in slots_pairs.values():
  1400. self.execute_command("MSET", *pairs)
  1401. return self
  1402. for command in PIPELINE_BLOCKED_COMMANDS:
  1403. command = command.replace(" ", "_").lower()
  1404. if command == "mset_nonatomic":
  1405. continue
  1406. setattr(ClusterPipeline, command, block_pipeline_command(command))
  1407. class PipelineCommand:
  1408. def __init__(self, position: int, *args: Any, **kwargs: Any) -> None:
  1409. self.args = args
  1410. self.kwargs = kwargs
  1411. self.position = position
  1412. self.result: Union[Any, Exception] = None
  1413. def __repr__(self) -> str:
  1414. return f"[{self.position}] {self.args} ({self.kwargs})"