add support for custom connection pool class in NodesManager (#2547)

zakaf · zach-iee · web-flow · commit 7dd73a306add · 2023-01-11T11:55:45.000+02:00
Co-authored-by: zach.lee &lt;zach.lee@sendbird.com&gt;
diff --git a/redis/cluster.py b/redis/cluster.py
@@ -121,6 +121,7 @@ def parse_cluster_shards(resp, **options):
     "charset",
     "connection_class",
     "connection_pool",
+    "connection_pool_class",
     "client_name",
     "credential_provider",
     "db",
@@ -1267,6 +1268,7 @@ def __init__(
         require_full_coverage=False,
         lock=None,
         dynamic_startup_nodes=True,
+        connection_pool_class=ConnectionPool,
         **kwargs,
     ):
         self.nodes_cache = {}
@@ -1277,6 +1279,7 @@ def __init__(
         self.from_url = from_url
         self._require_full_coverage = require_full_coverage
         self._dynamic_startup_nodes = dynamic_startup_nodes
+        self.connection_pool_class = connection_pool_class
         self._moved_exception = None
         self.connection_kwargs = kwargs
         self.read_load_balancer = LoadBalancer()
@@ -1420,7 +1423,7 @@ def create_redis_node(self, host, port, **kwargs):
             # Create a redis node with a costumed connection pool
             kwargs.update({"host": host})
             kwargs.update({"port": port})
-            r = Redis(connection_pool=ConnectionPool(**kwargs))
+            r = Redis(connection_pool=self.connection_pool_class(**kwargs))
         else:
             r = Redis(host=host, port=port, **kwargs)
         return r
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
@@ -18,7 +18,7 @@
     get_node_name,
 )
 from redis.commands import CommandsParser
-from redis.connection import Connection
+from redis.connection import BlockingConnectionPool, Connection, ConnectionPool
 from redis.crc import key_slot
 from redis.exceptions import (
     AskError,
@@ -2496,6 +2496,21 @@ def test_init_slots_dynamic_startup_nodes(self, dynamic_startup_nodes):
         else:
             assert startup_nodes == ["my@DNS.com:7000"]
 
+    @pytest.mark.parametrize(
+        "connection_pool_class", [ConnectionPool, BlockingConnectionPool]
+    )
+    def test_connection_pool_class(self, connection_pool_class):
+        rc = get_mocked_redis_client(
+            url="redis://my@DNS.com:7000",
+            cluster_slots=default_cluster_slots,
+            connection_pool_class=connection_pool_class,
+        )
+
+        for node in rc.nodes_manager.nodes_cache.values():
+            assert isinstance(
+                node.redis_connection.connection_pool, connection_pool_class
+            )
+
 
 @pytest.mark.onlycluster
 class TestClusterPubSubObject:
