From 7d36c4bca5aee3a0334d8fb30a68897d9c3f712d Mon Sep 17 00:00:00 2001
From: Sam Nguyen <samxnguyen@gmail.com>
Date: Thu, 2 Nov 2017 13:03:41 -0700
Subject: [PATCH] Allow default cluster_type

If cluster_type is not specified on the command line, and only 1
cluster_type is configured, then the command will use that cluster_type.

The purpose is to simplify the command-line invocation when you always
have only 1 cluster type.
---
 kafka_utils/kafka_check/main.py            | 3 +--
 kafka_utils/kafka_cluster_manager/main.py  | 4 ++--
 kafka_utils/kafka_consumer_manager/main.py | 4 ++--
 kafka_utils/kafka_corruption_check/main.py | 4 ++--
 kafka_utils/kafka_rolling_restart/main.py  | 4 ++--
 kafka_utils/util/config.py                 | 9 +++++++++
 6 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/kafka_utils/kafka_check/main.py b/kafka_utils/kafka_check/main.py
index 10d8eebd..c2886b15 100644
--- a/kafka_utils/kafka_check/main.py
+++ b/kafka_utils/kafka_check/main.py
@@ -54,8 +54,7 @@ def parse_args():
         "--cluster-type",
         "-t",
         dest='cluster_type',
-        required=True,
-        help='Type of cluster',
+        help='Type of cluster. If not specified and there is only 1 cluster type configured, it will use that one.',
         default=None,
     )
     parser.add_argument(
diff --git a/kafka_utils/kafka_cluster_manager/main.py b/kafka_utils/kafka_cluster_manager/main.py
index bae4472a..6e5733d7 100644
--- a/kafka_utils/kafka_cluster_manager/main.py
+++ b/kafka_utils/kafka_cluster_manager/main.py
@@ -64,9 +64,9 @@ def parse_args():
         '--cluster-type',
         '-t',
         dest='cluster_type',
-        help='Type of the cluster.',
+        help='Type of cluster. If not specified and there is only 1 cluster type configured, it will use that one.',
         type=str,
-        required=True,
+        default=None
     )
     parser.add_argument(
         '--cluster-name',
diff --git a/kafka_utils/kafka_consumer_manager/main.py b/kafka_utils/kafka_consumer_manager/main.py
index 96f80699..bbe2926d 100644
--- a/kafka_utils/kafka_consumer_manager/main.py
+++ b/kafka_utils/kafka_consumer_manager/main.py
@@ -43,8 +43,8 @@ def parse_args():
         "manipulate consumer offsets for a specific consumer group.",
     )
     parser.add_argument(
-        '--cluster-type', '-t', dest='cluster_type', required=True,
-        help='Type of Kafka cluster. This is a mandatory option.',
+        '--cluster-type', '-t', default=None, dest='cluster_type',
+        help='Type of cluster. If not specified and there is only 1 cluster type configured, it will use that one.',
     )
     parser.add_argument(
         '--cluster-name', '-c', dest='cluster_name',
diff --git a/kafka_utils/kafka_corruption_check/main.py b/kafka_utils/kafka_corruption_check/main.py
index a861f212..ba5bf72f 100644
--- a/kafka_utils/kafka_corruption_check/main.py
+++ b/kafka_utils/kafka_corruption_check/main.py
@@ -92,8 +92,8 @@ def parse_args():
     parser.add_argument(
         '--cluster-type',
         '-t',
-        required=True,
-        help='cluster type, e.g. "generic"',
+        default=None,
+        help='Type of cluster. If not specified and there is only 1 cluster type configured, it will use that one.',
     )
     parser.add_argument(
         '--cluster-name',
diff --git a/kafka_utils/kafka_rolling_restart/main.py b/kafka_utils/kafka_rolling_restart/main.py
index a3b77173..cb8d9bd9 100644
--- a/kafka_utils/kafka_rolling_restart/main.py
+++ b/kafka_utils/kafka_rolling_restart/main.py
@@ -62,8 +62,8 @@ def parse_opts():
     parser.add_argument(
         '--cluster-type',
         '-t',
-        required=True,
-        help='cluster type, e.g. "standard"',
+        default=None,
+        help='Type of cluster. If not specified and there is only 1 cluster type configured, it will use that one.',
     )
     parser.add_argument(
         '--cluster-name',
diff --git a/kafka_utils/util/config.py b/kafka_utils/util/config.py
index 2b5fd089..fa274ef6 100644
--- a/kafka_utils/util/config.py
+++ b/kafka_utils/util/config.py
@@ -232,6 +232,15 @@ def get_cluster_config(
     else:
         config_dirs = [kafka_topology_base_path]
 
+    # If cluster_type is not set, and there is only one file in the config dir,
+    # then use that as the default cluster type
+    if cluster_type is None:
+        for config_dir in config_dirs:
+            cluster_types = [x for x in os.listdir(config_dir) if x.endswith(".yaml")]
+            if len(cluster_types) is 1:
+                cluster_type = cluster_types[0][:-5] # cut off the '.yaml' part
+                break
+
     topology = None
     for config_dir in config_dirs:
         try: