From faf55b65ffb7294b1a639cc9adb0231b987b0e32 Mon Sep 17 00:00:00 2001
From: dvora-h <67596500+dvora-h@users.noreply.github.com>
Date: Sun, 8 May 2022 15:23:00 +0300
Subject: Add support for CLUSTER SHARDS (#2151)

* Add support for CLUSTER SHARDS

* linters

* add docstring

* linters
---
 redis/cluster.py          | 26 +++++++++++++++++++++++++-
 redis/commands/cluster.py |  8 ++++++++
 tests/test_cluster.py     | 23 +++++++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/redis/cluster.py b/redis/cluster.py
index d42d49b..e5baee2 100644
--- a/redis/cluster.py
+++ b/redis/cluster.py
@@ -92,6 +92,26 @@ def parse_cluster_slots(resp, **options):
     return slots
 
 
+def parse_cluster_shards(resp, **options):
+    """
+    Parse CLUSTER SHARDS response.
+    """
+    shards = []
+    for x in resp:
+        shard = {"slots": [], "nodes": []}
+        for i in range(0, len(x[1]), 2):
+            shard["slots"].append((x[1][i], (x[1][i + 1])))
+        nodes = x[3]
+        for node in nodes:
+            dict_node = {}
+            for i in range(0, len(node), 2):
+                dict_node[node[i]] = node[i + 1]
+            shard["nodes"].append(dict_node)
+        shards.append(shard)
+
+    return shards
+
+
 PRIMARY = "primary"
 REPLICA = "replica"
 SLOT_ID = "slot-id"
@@ -274,6 +294,7 @@ class AbstractRedisCluster:
                 "CLUSTER RESET",
                 "CLUSTER SET-CONFIG-EPOCH",
                 "CLUSTER SLOTS",
+                "CLUSTER SHARDS",
                 "CLUSTER COUNT-FAILURE-REPORTS",
                 "CLUSTER KEYSLOT",
                 "COMMAND",
@@ -354,7 +375,10 @@ class AbstractRedisCluster:
         ],
     )
 
-    CLUSTER_COMMANDS_RESPONSE_CALLBACKS = {"CLUSTER SLOTS": parse_cluster_slots}
+    CLUSTER_COMMANDS_RESPONSE_CALLBACKS = {
+        "CLUSTER SLOTS": parse_cluster_slots,
+        "CLUSTER SHARDS": parse_cluster_shards,
+    }
 
     RESULT_CALLBACKS = dict_merge(
         list_keys_to_dict(["PUBSUB NUMSUB"], parse_pubsub_numsub),
diff --git a/redis/commands/cluster.py b/redis/commands/cluster.py
index 06b702f..b91b65f 100644
--- a/redis/commands/cluster.py
+++ b/redis/commands/cluster.py
@@ -581,6 +581,14 @@ class ClusterManagementCommands(ManagementCommands):
         """
         return self.execute_command("CLUSTER SLOTS", target_nodes=target_nodes)
 
+    def cluster_shards(self, target_nodes=None):
+        """
+        Returns details about the shards of the cluster.
+
+        For more information see https://redis.io/commands/cluster-shards
+        """
+        return self.execute_command("CLUSTER SHARDS", target_nodes=target_nodes)
+
     def cluster_links(self, target_node: "TargetNodesT") -> ResponseT:
         """
         Each node in a Redis Cluster maintains a pair of long-lived TCP link with each
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 376e3f8..7e64345 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -856,6 +856,29 @@ class TestClusterRedisCommands:
         assert cluster_slots.get((0, 8191)) is not None
         assert cluster_slots.get((0, 8191)).get("primary") == ("127.0.0.1", 7000)
 
+    @skip_if_server_version_lt("7.0.0")
+    @skip_if_redis_enterprise()
+    def test_cluster_shards(self, r):
+        cluster_shards = r.cluster_shards()
+        assert isinstance(cluster_shards, list)
+        assert isinstance(cluster_shards[0], dict)
+        attributes = [
+            "id",
+            "endpoint",
+            "ip",
+            "hostname",
+            "port",
+            "tls-port",
+            "role",
+            "replication-offset",
+            "health",
+        ]
+        for x in cluster_shards:
+            assert list(x.keys()) == ["slots", "nodes"]
+            for node in x["nodes"]:
+                for attribute in node.keys():
+                    assert attribute in attributes
+
     @skip_if_redis_enterprise()
     def test_cluster_addslots(self, r):
         node = r.get_random_node()
-- 
cgit v1.2.1