图片解析应用
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1625 lines
62 KiB

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