Repository: redis/redis-py
Pull request #2547 (2023-01-09 14:35:21)

add support for custom connection pool class in NodesManager
### Pull Request check-list

_Please make sure to review and check all of these items:_

- [X] Does `$ tox` pass with this change (including linting)?
- [X] Do the CI tests pass with this change (enable it first in your forked repo and wait for the github action build to finish)?
- [X] Is the new or changed code fully tested?
- [ ] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?
- [ ] Is there an example added to the examples folder (if applicable)?
- [ ] Was the change added to CHANGES file?

_NOTE: these things are not required to open a PR and can be done
afterwards / while the PR is open._

### Description of change
currently, RedisCluster doesn't support using any other ConnectionPool implementation
This PR adds support for different ConnectionPool implementations like BlockingConnectionPool.

I've got the idea from Sentinel, which allows custom ConnectionPool implementation
https://github.com/redis/redis-py/blob/e425674d84a63762f16d5b44b19aa70119fcd814/redis/sentinel.py#L273

related issue: https://github.com/redis/redis-py/issues/2350





Patch:
diff --git a/CHANGES b/CHANGES
index fca8d316..02daf5ee 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,3 +1,4 @@
+    * Fix string cleanse in Redis Graph
     * Make PythonParser resumable in case of error (#2510)
     * Add `timeout=None` in `SentinelConnectionManager.read_response`
     * Documentation fix: password protected socket connection (#2374)
diff --git a/redis/cluster.py b/redis/cluster.py
index f1150073..5f730a85 100644
--- a/redis/cluster.py
+++ b/redis/cluster.py
@@ -121,6 +121,7 @@ REDIS_ALLOWED_KEYS = (
     "charset",
     "connection_class",
     "connection_pool",
+    "connection_pool_class",
     "client_name",
     "credential_provider",
     "db",
@@ -1267,6 +1268,7 @@ class NodesManager:
         require_full_coverage=False,
         lock=None,
         dynamic_startup_nodes=True,
+        connection_pool_class=ConnectionPool,
         **kwargs,
     ):
         self.nodes_cache = {}
@@ -1277,6 +1279,7 @@ class NodesManager:
         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 @@ class NodesManager:
             # 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
@@ -1437,6 +1440,8 @@ class NodesManager:
             if target_node is None or target_node.redis_connection is None:
                 # create new cluster node for this cluster
                 target_node = ClusterNode(host, port, role)
+            if target_node.server_type != role:
+                target_node.server_type = role
 
         return target_node
 
diff --git a/redis/commands/helpers.py b/redis/commands/helpers.py
index 6989ab59..b65cd1a9 100644
--- a/redis/commands/helpers.py
+++ b/redis/commands/helpers.py
@@ -115,6 +115,7 @@ def quote_string(v):
     if len(v) == 0:
         return '""'
 
+    v = v.replace("\\", "\\\\")
     v = v.replace('"', '\\"')
 
     return f'"{v}"'


Test patch:
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 27cfee1c..da6a8e4b 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -18,7 +18,7 @@ from redis.cluster import (
     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,
@@ -2262,6 +2262,57 @@ class TestNodesManager:
 
         assert len(n_manager.nodes_cache) == 6
 
+    def test_init_promote_server_type_for_node_in_cache(self):
+        """
+        When replica is promoted to master, nodes_cache must change the server type
+        accordingly
+        """
+        cluster_slots_before_promotion = [
+            [0, 16383, ["127.0.0.1", 7000], ["127.0.0.1", 7003]]
+        ]
+        cluster_slots_after_promotion = [
+            [0, 16383, ["127.0.0.1", 7003], ["127.0.0.1", 7004]]
+        ]
+
+        cluster_slots_results = [
+            cluster_slots_before_promotion,
+            cluster_slots_after_promotion,
+        ]
+
+        with patch.object(Redis, "execute_command") as execute_command_mock:
+
+            def execute_command(*_args, **_kwargs):
+                if _args[0] == "CLUSTER SLOTS":
+                    mock_cluster_slots = cluster_slots_results.pop(0)
+                    return mock_cluster_slots
+                elif _args[0] == "COMMAND":
+                    return {"get": [], "set": []}
+                elif _args[0] == "INFO":
+                    return {"cluster_enabled": True}
+                elif len(_args) > 1 and _args[1] == "cluster-require-full-coverage":
+                    return {"cluster-require-full-coverage": False}
+                else:
+                    return execute_command_mock(*_args, **_kwargs)
+
+            execute_command_mock.side_effect = execute_command
+
+            nm = NodesManager(
+                startup_nodes=[ClusterNode(host=default_host, port=default_port)],
+                from_url=False,
+                require_full_coverage=False,
+                dynamic_startup_nodes=True,
+            )
+
+            assert nm.default_node.host == "127.0.0.1"
+            assert nm.default_node.port == 7000
+            assert nm.default_node.server_type == PRIMARY
+
+            nm.initialize()
+
+            assert nm.default_node.host == "127.0.0.1"
+            assert nm.default_node.port == 7003
+            assert nm.default_node.server_type == PRIMARY
+
     def test_init_slots_cache_cluster_mode_disabled(self):
         """
         Test that creating a RedisCluster failes if one of the startup nodes
@@ -2445,6 +2496,21 @@ class TestNodesManager:
         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:
diff --git a/tests/test_graph.py b/tests/test_graph.py
index d71df486..4721b2f4 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -124,7 +124,7 @@ def test_path(client):
 
 @pytest.mark.redismod
 def test_param(client):
-    params = [1, 2.3, "str", True, False, None, [0, 1, 2]]
+    params = [1, 2.3, "str", True, False, None, [0, 1, 2], r"\" RETURN 1337 //"]
     query = "RETURN $param"
     for param in params:
         result = client.graph().query(query, {"param": param})
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
index 35958290..57a94d2f 100644
--- a/tests/test_helpers.py
+++ b/tests/test_helpers.py
@@ -80,3 +80,9 @@ def test_quote_string():
     assert quote_string("hello world!") == '"hello world!"'
     assert quote_string("") == '""'
     assert quote_string("hello world!") == '"hello world!"'
+    assert quote_string("abc") == '"abc"'
+    assert quote_string("") == '""'
+    assert quote_string('"') == r'"\""'
+    assert quote_string(r"foo \ bar") == r'"foo \\ bar"'
+    assert quote_string(r"foo \" bar") == r'"foo \\\" bar"'
+    assert quote_string('a"a') == r'"a\"a"'