[add] Enabled enterprise cluster setup steps

Author: filipecosta90

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redisbench-admin"
-version = "0.1.15"
+version = "0.1.16"
 description = "Redis benchmark run helper. A wrapper around Redis and Redis Modules benchmark tools ( ftsb_redisearch, memtier_benchmark, redis-benchmark, aibench, etc... )."
 authors = ["filipecosta90 <[email protected]>"]
 readme = "README.md"
@@ -21,6 +21,7 @@ toml = "^0.10.1"
 seaborn = "^0.10.1"
 matplotlib = "^3.2.1"
 redistimeseries = "^0.8.0"
+redis-py-cluster = "^2.1.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^4.6"

pyproject.toml

@@ -9,6 +9,7 @@
 import redis
 import requests
 from cpuinfo import cpuinfo
+from rediscluster import RedisCluster
 from tqdm import tqdm
 
 from redisbench_admin.compare.compare import generate_comparison_dataframe_configs
@@ -119,7 +120,7 @@ def run_command_logic(args):
     progress = tqdm(unit="bench steps", total=total_steps)
     for repetition in range(1, args.repetitions + 1):
         if benchmark_repetitions_require_teardown is True or repetition == 1:
-            aux_client = run_setup_commands(args, aux_client, benchmark_config)
+            aux_client = run_setup_commands(args, aux_client, benchmark_config, oss_cluster_mode)
             if "setup" in run_stages_inputs:
                 setup_run_key = "setup-run-{}.json".format(repetition)
                 setup_run_json_output_fullpath = "{}/{}".format(local_path, setup_run_key)
@@ -197,16 +198,27 @@ def run_command_logic(args):
         upload_artifacts_to_s3(artifacts, s3_bucket_name, s3_bucket_path)
 
 
-def run_setup_commands(args, aux_client, benchmark_config):
+def run_setup_commands(args, aux_client, benchmark_config, cluster_enabled):
     print("Running setup steps...")
-    for command in benchmark_config["setup"]["commands"]:
-        try:
+    try:
+        if cluster_enabled:
+            host_port_arr = args.redis_url.split(":")
+            host = host_port_arr[0]
+            port = host_port_arr[1]
+            startup_nodes = [{"host": host, "port": port}]
+            aux_client = RedisCluster(startup_nodes=startup_nodes, decode_responses=True)
+            cluster_nodes = aux_client.cluster_nodes()
+            for command in benchmark_config["setup"]["commands"]:
+                for master_node in aux_client.connection_pool.nodes.all_masters():
+                    redis.from_url("redis://"+master_node["name"]).execute_command(" ".join(command))
+        else:
             aux_client = redis.from_url(args.redis_url)
-            aux_client.execute_command(" ".join(command))
-        except redis.connection.ConnectionError as e:
-            print('Error while issuing setup command to Redis.Command {}! Error message: {} Exiting..'.format(command,
-                                                                                                              e.__str__()))
-            sys.exit(1)
+            for command in benchmark_config["setup"]["commands"]:
+                aux_client.execute_command(" ".join(command))
+    except redis.connection.ConnectionError as e:
+        print('Error while issuing setup command to Redis.Command {}! Error message: {} Exiting..'.format(command,
+                                                                                                          e.__str__()))
+        sys.exit(1)
     return aux_client