Extended "CLUSTER NODES" parser to support special slot entries (importing, migrating). (#2080)

See https://redis.io/commands/cluster-nodes/#special-slot-entries

Co-authored-by: dvora-h <67596500+dvora-h@users.noreply.github.com>

2 files changed, 65 insertions, 5 deletions

diff --git a/redis/client.py b/redis/client.py
index 8c17665..d8d7a75 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -466,7 +466,6 @@ def _parse_node_line(line):
     line_items = line.split(" ")
     node_id, addr, flags, master_id, ping, pong, epoch, connected = line.split(" ")[:8]
     addr = addr.split("@")[0]
-    slots = [sl.split("-") for sl in line_items[8:]]
     node_dict = {
         "node_id": node_id,
         "flags": flags,
@@ -474,18 +473,42 @@ def _parse_node_line(line):
         "last_ping_sent": ping,
         "last_pong_rcvd": pong,
         "epoch": epoch,
-        "slots": slots,
+        "slots": [],
+        "migrations": [],
         "connected": True if connected == "connected" else False,
     }
+    if len(line_items) >= 9:
+        slots, migrations = _parse_slots(line_items[8:])
+        node_dict["slots"], node_dict["migrations"] = slots, migrations
     return addr, node_dict
 
 
+def _parse_slots(slot_ranges):
+    slots, migrations = [], []
+    for s_range in slot_ranges:
+        if "->-" in s_range:
+            slot_id, dst_node_id = s_range[1:-1].split("->-", 1)
+            migrations.append(
+                {"slot": slot_id, "node_id": dst_node_id, "state": "migrating"}
+            )
+        elif "-<-" in s_range:
+            slot_id, src_node_id = s_range[1:-1].split("-<-", 1)
+            migrations.append(
+                {"slot": slot_id, "node_id": src_node_id, "state": "importing"}
+            )
+        else:
+            s_range = [sl for sl in s_range.split("-")]
+            slots.append(s_range)
+
+    return slots, migrations
+
+
 def parse_cluster_nodes(response, **options):
     """
-    @see: https://redis.io/commands/cluster-nodes  # string
-    @see: https://redis.io/commands/cluster-replicas # list of string
+    @see: https://redis.io/commands/cluster-nodes  # string / bytes
+    @see: https://redis.io/commands/cluster-replicas # list of string / bytes
     """
-    if isinstance(response, str):
+    if isinstance(response, (str, bytes)):
         response = response.splitlines()
     return dict(_parse_node_line(str_if_bytes(node)) for node in response)
 
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 1ee5cdf..3794c31 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -968,6 +968,43 @@ class TestClusterRedisCommands:
         )
 
     @skip_if_redis_enterprise()
+    def test_cluster_nodes_importing_migrating(self, r):
+        response = (
+            "488ead2fcce24d8c0f158f9172cb1f4a9e040fe5 127.0.0.1:16381@26381 "
+            "master - 0 1648975557664 3 connected 10923-16383\n"
+            "8ae2e70812db80776f739a72374e57fc4ae6f89d 127.0.0.1:16380@26380 "
+            "master - 0 1648975555000 2 connected 1 5461-10922 ["
+            "2-<-ed8007ccfa2d91a7b76f8e6fba7ba7e257034a16]\n"
+            "ed8007ccfa2d91a7b76f8e6fba7ba7e257034a16 127.0.0.1:16379@26379 "
+            "myself,master - 0 1648975556000 1 connected 0 2-5460 ["
+            "2->-8ae2e70812db80776f739a72374e57fc4ae6f89d]\n"
+        )
+        mock_all_nodes_resp(r, response)
+        nodes = r.cluster_nodes()
+        assert len(nodes) == 3
+        node_16379 = nodes.get("127.0.0.1:16379")
+        node_16380 = nodes.get("127.0.0.1:16380")
+        node_16381 = nodes.get("127.0.0.1:16381")
+        assert node_16379.get("migrations") == [
+            {
+                "slot": "2",
+                "node_id": "8ae2e70812db80776f739a72374e57fc4ae6f89d",
+                "state": "migrating",
+            }
+        ]
+        assert node_16379.get("slots") == [["0"], ["2", "5460"]]
+        assert node_16380.get("migrations") == [
+            {
+                "slot": "2",
+                "node_id": "ed8007ccfa2d91a7b76f8e6fba7ba7e257034a16",
+                "state": "importing",
+            }
+        ]
+        assert node_16380.get("slots") == [["1"], ["5461", "10922"]]
+        assert node_16381.get("slots") == [["10923", "16383"]]
+        assert node_16381.get("migrations") == []
+
+    @skip_if_redis_enterprise()
     def test_cluster_replicate(self, r):
         node = r.get_random_node()
         all_replicas = r.get_replicas()