Skip to content

Commit 7dd73a3

Browse files
zakafzach-iee
andauthored
add support for custom connection pool class in NodesManager (#2547)
Co-authored-by: zach.lee <[email protected]>
1 parent bae6385 commit 7dd73a3

File tree

2 files changed

+20
-2
lines changed

2 files changed

+20
-2
lines changed

redis/cluster.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ def parse_cluster_shards(resp, **options):
121121
"charset",
122122
"connection_class",
123123
"connection_pool",
124+
"connection_pool_class",
124125
"client_name",
125126
"credential_provider",
126127
"db",
@@ -1267,6 +1268,7 @@ def __init__(
12671268
require_full_coverage=False,
12681269
lock=None,
12691270
dynamic_startup_nodes=True,
1271+
connection_pool_class=ConnectionPool,
12701272
**kwargs,
12711273
):
12721274
self.nodes_cache = {}
@@ -1277,6 +1279,7 @@ def __init__(
12771279
self.from_url = from_url
12781280
self._require_full_coverage = require_full_coverage
12791281
self._dynamic_startup_nodes = dynamic_startup_nodes
1282+
self.connection_pool_class = connection_pool_class
12801283
self._moved_exception = None
12811284
self.connection_kwargs = kwargs
12821285
self.read_load_balancer = LoadBalancer()
@@ -1420,7 +1423,7 @@ def create_redis_node(self, host, port, **kwargs):
14201423
# Create a redis node with a costumed connection pool
14211424
kwargs.update({"host": host})
14221425
kwargs.update({"port": port})
1423-
r = Redis(connection_pool=ConnectionPool(**kwargs))
1426+
r = Redis(connection_pool=self.connection_pool_class(**kwargs))
14241427
else:
14251428
r = Redis(host=host, port=port, **kwargs)
14261429
return r

tests/test_cluster.py

+16-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
get_node_name,
1919
)
2020
from redis.commands import CommandsParser
21-
from redis.connection import Connection
21+
from redis.connection import BlockingConnectionPool, Connection, ConnectionPool
2222
from redis.crc import key_slot
2323
from redis.exceptions import (
2424
AskError,
@@ -2496,6 +2496,21 @@ def test_init_slots_dynamic_startup_nodes(self, dynamic_startup_nodes):
24962496
else:
24972497
assert startup_nodes == ["[email protected]:7000"]
24982498

2499+
@pytest.mark.parametrize(
2500+
"connection_pool_class", [ConnectionPool, BlockingConnectionPool]
2501+
)
2502+
def test_connection_pool_class(self, connection_pool_class):
2503+
rc = get_mocked_redis_client(
2504+
url="redis://[email protected]:7000",
2505+
cluster_slots=default_cluster_slots,
2506+
connection_pool_class=connection_pool_class,
2507+
)
2508+
2509+
for node in rc.nodes_manager.nodes_cache.values():
2510+
assert isinstance(
2511+
node.redis_connection.connection_pool, connection_pool_class
2512+
)
2513+
24992514

25002515
@pytest.mark.onlycluster
25012516
class TestClusterPubSubObject:

0 commit comments

Comments
 (0)