redis 5.2.1__py3-none-any.whl → 5.3.0__py3-none-any.whl

redis/cluster.py

@@ -4,6 +4,7 @@ import sys
 import threading
 import time
 from collections import OrderedDict
+from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from redis._parsers import CommandsParser, Encoder
@@ -15,6 +16,12 @@ from redis.commands import READ_COMMANDS, RedisClusterCommands
 from redis.commands.helpers import list_or_args
 from redis.connection import ConnectionPool, DefaultParser, parse_url
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
+from redis.event import (
+    AfterPooledConnectionsInstantiationEvent,
+    AfterPubSubConnectionInstantiationEvent,
+    ClientType,
+    EventDispatcher,
+)
 from redis.exceptions import (
     AskError,
     AuthenticationError,
@@ -36,6 +43,7 @@ from redis.lock import Lock
 from redis.retry import Retry
 from redis.utils import (
     HIREDIS_AVAILABLE,
+    deprecated_args,
     dict_merge,
     list_keys_to_dict,
     merge_result,
@@ -48,10 +56,13 @@ def get_node_name(host: str, port: Union[str, int]) -> str:
     return f"{host}:{port}"
 
 
+@deprecated_args(
+    allowed_args=["redis_node"],
+    reason="Use get_connection(redis_node) instead",
+    version="5.3.0",
+)
 def get_connection(redis_node, *args, **options):
-    return redis_node.connection or redis_node.connection_pool.get_connection(
-        args[0], **options
-    )
+    return redis_node.connection or redis_node.connection_pool.get_connection()
 
 
 def parse_scan_result(command, res, **options):
@@ -418,7 +429,12 @@ class AbstractRedisCluster:
         list_keys_to_dict(["SCRIPT FLUSH"], lambda command, res: all(res.values())),
     )
 
-    ERRORS_ALLOW_RETRY = (ConnectionError, TimeoutError, ClusterDownError)
+    ERRORS_ALLOW_RETRY = (
+        ConnectionError,
+        TimeoutError,
+        ClusterDownError,
+        SlotNotCoveredError,
+    )
 
     def replace_default_node(self, target_node: "ClusterNode" = None) -> None:
         """Replace the default cluster node.
@@ -490,6 +506,11 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         """
         return cls(url=url, **kwargs)
 
+    @deprecated_args(
+        args_to_warn=["read_from_replicas"],
+        reason="Please configure the 'load_balancing_strategy' instead",
+        version="5.3.0",
+    )
     def __init__(
         self,
         host: Optional[str] = None,
@@ -500,11 +521,13 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         require_full_coverage: bool = False,
         reinitialize_steps: int = 5,
         read_from_replicas: bool = False,
+        load_balancing_strategy: Optional["LoadBalancingStrategy"] = None,
         dynamic_startup_nodes: bool = True,
         url: Optional[str] = None,
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
         **kwargs,
     ):
         """
@@ -527,11 +550,17 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             cluster client. If not all slots are covered, RedisClusterException
             will be thrown.
         :param read_from_replicas:
+             @deprecated - please use load_balancing_strategy instead
              Enable read from replicas in READONLY mode. You can read possibly
              stale data.
              When set to true, read commands will be assigned between the
              primary and its replications in a Round-Robin manner.
-         :param dynamic_startup_nodes:
+        :param load_balancing_strategy:
+             Enable read from replicas in READONLY mode and defines the load balancing
+             strategy that will be used for cluster node selection.
+             The data read from replicas is eventually consistent
+             with the data in primary nodes.
+        :param dynamic_startup_nodes:
              Set the RedisCluster's startup nodes to all of the discovered nodes.
              If true (default value), the cluster's discovered nodes will be used to
              determine the cluster nodes-slots mapping in the next topology refresh.
@@ -636,8 +665,13 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         self.command_flags = self.__class__.COMMAND_FLAGS.copy()
         self.node_flags = self.__class__.NODE_FLAGS.copy()
         self.read_from_replicas = read_from_replicas
+        self.load_balancing_strategy = load_balancing_strategy
         self.reinitialize_counter = 0
         self.reinitialize_steps = reinitialize_steps
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         self.nodes_manager = NodesManager(
             startup_nodes=startup_nodes,
             from_url=from_url,
@@ -646,6 +680,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             address_remap=address_remap,
             cache=cache,
             cache_config=cache_config,
+            event_dispatcher=self._event_dispatcher,
             **kwargs,
         )
 
@@ -683,7 +718,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
         connection.set_parser(ClusterParser)
         connection.on_connect()
 
-        if self.read_from_replicas:
+        if self.read_from_replicas or self.load_balancing_strategy:
             # Sending READONLY command to server to configure connection as
             # readonly. Since each cluster node may change its server type due
             # to a failover, we should establish a READONLY connection
@@ -810,6 +845,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             cluster_response_callbacks=self.cluster_response_callbacks,
             cluster_error_retry_attempts=self.cluster_error_retry_attempts,
             read_from_replicas=self.read_from_replicas,
+            load_balancing_strategy=self.load_balancing_strategy,
             reinitialize_steps=self.reinitialize_steps,
             lock=self._lock,
         )
@@ -927,7 +963,9 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             # get the node that holds the key's slot
             slot = self.determine_slot(*args)
             node = self.nodes_manager.get_node_from_slot(
-                slot, self.read_from_replicas and command in READ_COMMANDS
+                slot,
+                self.read_from_replicas and command in READ_COMMANDS,
+                self.load_balancing_strategy if command in READ_COMMANDS else None,
             )
             return [node]
 
@@ -1151,12 +1189,18 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
                     # refresh the target node
                     slot = self.determine_slot(*args)
                     target_node = self.nodes_manager.get_node_from_slot(
-                        slot, self.read_from_replicas and command in READ_COMMANDS
+                        slot,
+                        self.read_from_replicas and command in READ_COMMANDS,
+                        (
+                            self.load_balancing_strategy
+                            if command in READ_COMMANDS
+                            else None
+                        ),
                     )
                     moved = False
 
                 redis_node = self.get_redis_connection(target_node)
-                connection = get_connection(redis_node, *args, **kwargs)
+                connection = get_connection(redis_node)
                 if asking:
                     connection.send_command("ASKING")
                     redis_node.parse_response(connection, "ASKING", **kwargs)
@@ -1213,13 +1257,19 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
             except AskError as e:
                 redirect_addr = get_node_name(host=e.host, port=e.port)
                 asking = True
-            except ClusterDownError as e:
+            except (ClusterDownError, SlotNotCoveredError):
                 # ClusterDownError can occur during a failover and to get
                 # self-healed, we will try to reinitialize the cluster layout
                 # and retry executing the command
+
+                # SlotNotCoveredError can occur when the cluster is not fully
+                # initialized or can be temporary issue.
+                # We will try to reinitialize the cluster topology
+                # and retry executing the command
+
                 time.sleep(0.25)
                 self.nodes_manager.initialize()
-                raise e
+                raise
             except ResponseError:
                 raise
             except Exception as e:
@@ -1232,7 +1282,7 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands):
 
         raise ClusterError("TTL exhausted.")
 
-    def close(self):
+    def close(self) -> None:
         try:
             with self._lock:
                 if self.nodes_manager:
@@ -1300,6 +1350,12 @@ class ClusterNode:
             self.redis_connection.close()
 
 
+class LoadBalancingStrategy(Enum):
+    ROUND_ROBIN = "round_robin"
+    ROUND_ROBIN_REPLICAS = "round_robin_replicas"
+    RANDOM_REPLICA = "random_replica"
+
+
 class LoadBalancer:
     """
     Round-Robin Load Balancing
@@ -1309,15 +1365,38 @@ class LoadBalancer:
         self.primary_to_idx = {}
         self.start_index = start_index
 
-    def get_server_index(self, primary: str, list_size: int) -> int:
-        server_index = self.primary_to_idx.setdefault(primary, self.start_index)
-        # Update the index
-        self.primary_to_idx[primary] = (server_index + 1) % list_size
-        return server_index
+    def get_server_index(
+        self,
+        primary: str,
+        list_size: int,
+        load_balancing_strategy: LoadBalancingStrategy = LoadBalancingStrategy.ROUND_ROBIN,  # noqa: line too long ignored
+    ) -> int:
+        if load_balancing_strategy == LoadBalancingStrategy.RANDOM_REPLICA:
+            return self._get_random_replica_index(list_size)
+        else:
+            return self._get_round_robin_index(
+                primary,
+                list_size,
+                load_balancing_strategy == LoadBalancingStrategy.ROUND_ROBIN_REPLICAS,
+            )
 
     def reset(self) -> None:
         self.primary_to_idx.clear()
 
+    def _get_random_replica_index(self, list_size: int) -> int:
+        return random.randint(1, list_size - 1)
+
+    def _get_round_robin_index(
+        self, primary: str, list_size: int, replicas_only: bool
+    ) -> int:
+        server_index = self.primary_to_idx.setdefault(primary, self.start_index)
+        if replicas_only and server_index == 0:
+            # skip the primary node index
+            server_index = 1
+        # Update the index for the next round
+        self.primary_to_idx[primary] = (server_index + 1) % list_size
+        return server_index
+
 
 class NodesManager:
     def __init__(
@@ -1332,6 +1411,7 @@ class NodesManager:
         cache: Optional[CacheInterface] = None,
         cache_config: Optional[CacheConfig] = None,
         cache_factory: Optional[CacheFactoryInterface] = None,
+        event_dispatcher: Optional[EventDispatcher] = None,
         **kwargs,
     ):
         self.nodes_cache = {}
@@ -1353,6 +1433,13 @@ class NodesManager:
         if lock is None:
             lock = threading.Lock()
         self._lock = lock
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
+        self._credential_provider = self.connection_kwargs.get(
+            "credential_provider", None
+        )
         self.initialize()
 
     def get_node(self, host=None, port=None, node_name=None):
@@ -1413,7 +1500,21 @@ class NodesManager:
         # Reset moved_exception
         self._moved_exception = None
 
-    def get_node_from_slot(self, slot, read_from_replicas=False, server_type=None):
+    @deprecated_args(
+        args_to_warn=["server_type"],
+        reason=(
+            "In case you need select some load balancing strategy "
+            "that will use replicas, please set it through 'load_balancing_strategy'"
+        ),
+        version="5.3.0",
+    )
+    def get_node_from_slot(
+        self,
+        slot,
+        read_from_replicas=False,
+        load_balancing_strategy=None,
+        server_type=None,
+    ):
         """
         Gets a node that servers this hash slot
         """
@@ -1428,11 +1529,14 @@ class NodesManager:
                 f'"require_full_coverage={self._require_full_coverage}"'
             )
 
-        if read_from_replicas is True:
-            # get the server index in a Round-Robin manner
+        if read_from_replicas is True and load_balancing_strategy is None:
+            load_balancing_strategy = LoadBalancingStrategy.ROUND_ROBIN
+
+        if len(self.slots_cache[slot]) > 1 and load_balancing_strategy:
+            # get the server index using the strategy defined in load_balancing_strategy
             primary_name = self.slots_cache[slot][0].name
             node_idx = self.read_load_balancer.get_server_index(
-                primary_name, len(self.slots_cache[slot])
+                primary_name, len(self.slots_cache[slot]), load_balancing_strategy
             )
         elif (
             server_type is None
@@ -1479,11 +1583,19 @@ class NodesManager:
         """
         This function will create a redis connection to all nodes in :nodes:
         """
+        connection_pools = []
         for node in nodes:
             if node.redis_connection is None:
                 node.redis_connection = self.create_redis_node(
                     host=node.host, port=node.port, **self.connection_kwargs
                 )
+                connection_pools.append(node.redis_connection.connection_pool)
+
+        self._event_dispatcher.dispatch(
+            AfterPooledConnectionsInstantiationEvent(
+                connection_pools, ClientType.SYNC, self._credential_provider
+            )
+        )
 
     def create_redis_node(self, host, port, **kwargs):
         if self.from_url:
@@ -1613,7 +1725,7 @@ class NodesManager:
                             if len(disagreements) > 5:
                                 raise RedisClusterException(
                                     f"startup_nodes could not agree on a valid "
-                                    f'slots cache: {", ".join(disagreements)}'
+                                    f"slots cache: {', '.join(disagreements)}"
                                 )
 
             fully_covered = self.check_slots_coverage(tmp_slots)
@@ -1658,7 +1770,7 @@ class NodesManager:
         # If initialize was called after a MovedError, clear it
         self._moved_exception = None
 
-    def close(self):
+    def close(self) -> None:
         self.default_node = None
         for node in self.nodes_cache.values():
             if node.redis_connection:
@@ -1698,6 +1810,7 @@ class ClusterPubSub(PubSub):
         host=None,
         port=None,
         push_handler_func=None,
+        event_dispatcher: Optional["EventDispatcher"] = None,
         **kwargs,
     ):
         """
@@ -1706,7 +1819,7 @@ class ClusterPubSub(PubSub):
         first command execution. The node will be determined by:
          1. Hashing the channel name in the request to find its keyslot
          2. Selecting a node that handles the keyslot: If read_from_replicas is
-            set to true, a replica can be selected.
+            set to true or load_balancing_strategy is set, a replica can be selected.
 
         :type redis_cluster: RedisCluster
         :type node: ClusterNode
@@ -1723,10 +1836,15 @@ class ClusterPubSub(PubSub):
         self.cluster = redis_cluster
         self.node_pubsub_mapping = {}
         self._pubsubs_generator = self._pubsubs_generator()
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
         super().__init__(
             connection_pool=connection_pool,
             encoder=redis_cluster.encoder,
             push_handler_func=push_handler_func,
+            event_dispatcher=self._event_dispatcher,
             **kwargs,
         )
 
@@ -1797,7 +1915,9 @@ class ClusterPubSub(PubSub):
                     channel = args[1]
                     slot = self.cluster.keyslot(channel)
                     node = self.cluster.nodes_manager.get_node_from_slot(
-                        slot, self.cluster.read_from_replicas
+                        slot,
+                        self.cluster.read_from_replicas,
+                        self.cluster.load_balancing_strategy,
                     )
                 else:
                     # Get a random node
@@ -1805,14 +1925,17 @@ class ClusterPubSub(PubSub):
                 self.node = node
                 redis_connection = self.cluster.get_redis_connection(node)
                 self.connection_pool = redis_connection.connection_pool
-            self.connection = self.connection_pool.get_connection(
-                "pubsub", self.shard_hint
-            )
+            self.connection = self.connection_pool.get_connection()
             # register a callback that re-subscribes to any channels we
             # were listening to when we were disconnected
             self.connection.register_connect_callback(self.on_connect)
             if self.push_handler_func is not None and not HIREDIS_AVAILABLE:
                 self.connection._parser.set_pubsub_push_handler(self.push_handler_func)
+            self._event_dispatcher.dispatch(
+                AfterPubSubConnectionInstantiationEvent(
+                    self.connection, self.connection_pool, ClientType.SYNC, self._lock
+                )
+            )
         connection = self.connection
         self._execute(connection, connection.send_command, *args)
 
@@ -1937,6 +2060,7 @@ class ClusterPipeline(RedisCluster):
         cluster_response_callbacks: Optional[Dict[str, Callable]] = None,
         startup_nodes: Optional[List["ClusterNode"]] = None,
         read_from_replicas: bool = False,
+        load_balancing_strategy: Optional[LoadBalancingStrategy] = None,
         cluster_error_retry_attempts: int = 3,
         reinitialize_steps: int = 5,
         lock=None,
@@ -1952,6 +2076,7 @@ class ClusterPipeline(RedisCluster):
         )
         self.startup_nodes = startup_nodes if startup_nodes else []
         self.read_from_replicas = read_from_replicas
+        self.load_balancing_strategy = load_balancing_strategy
         self.command_flags = self.__class__.COMMAND_FLAGS.copy()
         self.cluster_response_callbacks = cluster_response_callbacks
         self.cluster_error_retry_attempts = cluster_error_retry_attempts
@@ -2023,12 +2148,11 @@ class ClusterPipeline(RedisCluster):
         """
         cmd = " ".join(map(safe_str, command))
         msg = (
-            f"Command # {number} ({cmd}) of pipeline "
-            f"caused error: {exception.args[0]}"
+            f"Command # {number} ({cmd}) of pipeline caused error: {exception.args[0]}"
         )
         exception.args = (msg,) + exception.args[1:]
 
-    def execute(self, raise_on_error=True):
+    def execute(self, raise_on_error: bool = True) -> List[Any]:
         """
         Execute all the commands in the current pipeline
         """
@@ -2162,8 +2286,8 @@ class ClusterPipeline(RedisCluster):
                 if node_name not in nodes:
                     redis_node = self.get_redis_connection(node)
                     try:
-                        connection = get_connection(redis_node, c.args)
-                    except ConnectionError:
+                        connection = get_connection(redis_node)
+                    except (ConnectionError, TimeoutError):
                         for n in nodes.values():
                             n.connection_pool.release(n.connection)
                         # Connection retries are being handled in the node's