Merge pull request #10 from devshawn/kafka-preferred-replica-election

feat: add kafka-preferred-replica-election command (closes #3)

Author: devshawn

README.md

@@ -79,6 +79,7 @@ Currently, the following commands are supported:
 * `kafka-avro-console-producer`
 * `kafka-verifiable-consumer`
 * `kafka-verifiable-producer`
+* `kafka-preferred-replica-election`
 * `kafka-replica-verification`
 * `kafka-broker-api-versions`
 * `kafka-consumer-groups`

kafkashell/constants.py

@@ -24,6 +24,7 @@
 COMMAND_KAFKA_VERIFIABLE_CONSUMER = "kafka-verifiable-consumer"
 COMMAND_KAFKA_VERIFIABLE_PRODUCER = "kafka-verifiable-producer"
 COMMAND_KAFKA_CONSUMER_GROUPS = "kafka-consumer-groups"
+COMMAND_KAFKA_PREFERRED_REPLICA_ELECTION = "kafka-preferred-replica-election"
 COMMAND_KAFKA_REPLICA_VERIFICATION = "kafka-replica-verification"
 COMMAND_KAFKA_BROKER_API_VERSIONS = "kafka-broker-api-versions"
 COMMAND_KAFKA_DELETE_RECORDS = "kafka-delete-records"

kafkashell/data/completer.json

@@ -792,6 +792,32 @@
                 }
             }
         },
+        "kafka-preferred-replica-election": {
+            "name": "kafka-preferred-replica-election",
+            "description": "Tool to cause leadership for each partition to be transferred back to the 'preferred replica'.",
+            "options": {
+                "--admin.config": {
+                    "name": "--admin.config",
+                    "description": "Property file containing configs to be passed to the admin client."
+                },
+                "--bootstrap-server": {
+                    "name": "--bootstrap-server",
+                    "description": "The Kafka broker(s) to connect to. NOTE: This'll overwrite the selected cluster value."
+                },
+                "--help": {
+                    "name": "--help",
+                    "description": "Print usage information."
+                },
+                "--path-to-json-file": {
+                    "name": "--path-to-json-file",
+                    "description": "The JSON file with the list of partitions to perform preferred replica leader election on."
+                },
+                "--zookeeper": {
+                    "name": "--zookeeper",
+                    "description": "DEPRECATED. Zookeeper connection string."
+                }
+            }
+        },
         "kafka-replica-verification": {
             "name": "kafka-replica-verification",
             "description": "Validate that all replicas for a set of topics have the same data.",

kafkashell/executor.py

@@ -87,6 +87,9 @@ def execute_valid_command(self, command):
         elif command.startswith(constants.COMMAND_KAFKA_CONSUMER_GROUPS):
             final_command = self.handle_kafka_consumer_groups_command(command)
 
+        elif command.startswith(constants.COMMAND_KAFKA_PREFERRED_REPLICA_ELECTION):
+            final_command = self.handle_kafka_preferred_replica_election(command)
+
         elif command.startswith(constants.COMMAND_KAFKA_REPLICA_VERIFICATION):
             final_command = self.handle_kafka_replica_verification_command(command)
 
@@ -193,6 +196,11 @@ def handle_kafka_consumer_groups_command(self, command):
         command += self.handle_admin_client_settings(command)
         return command
 
+    def handle_kafka_preferred_replica_election(self, command):
+        command += self.handle_bootstrap_server_flag(command)
+        command += self.handle_admin_client_settings(command)
+        return command
+
     def handle_kafka_replica_verification_command(self, command):
         command += self.handle_broker_list_flag(command)
         return command
@@ -263,7 +271,9 @@ def handle_zookeeper_shell_input(self, command):
 
     def handle_admin_client_settings(self, command):
         if "admin_client_settings" in self.settings.get_cluster_details().keys():
-            return self.handle_config(command, "command-config", "admin_client")
+            admin_client_option = "command-config" if constants.COMMAND_KAFKA_PREFERRED_REPLICA_ELECTION \
+                                                      not in command else "admin.config"
+            return self.handle_config(command, admin_client_option, "admin_client")
         else:
             return ""
 

tests/test_completer_data.py

@@ -4,7 +4,7 @@
     (
         "",
         ["version", "cluster-select", "cluster-describe", "exit", "clear", "kafka-acls", "kafka-avro-console-consumer",
-         "kafka-avro-console-producer", "kafka-replica-verification",
+         "kafka-avro-console-producer", "kafka-replica-verification", "kafka-preferred-replica-election",
          "kafka-broker-api-versions", "kafka-configs", "kafka-console-consumer", "kafka-console-producer",
          "kafka-consumer-groups", "kafka-delete-records", "kafka-dump-log", "kafka-log-dirs", "kafka-topics",
          "kafka-verifiable-consumer", "kafka-verifiable-producer", "ksql", "zookeeper-shell"]
@@ -14,12 +14,12 @@
         ["kafka-acls", "kafka-avro-console-consumer", "kafka-avro-console-producer", "kafka-broker-api-versions",
          "kafka-configs", "kafka-console-consumer", "kafka-console-producer", "kafka-consumer-groups",
          "kafka-delete-records", "kafka-dump-log", "kafka-log-dirs", "kafka-topics", "kafka-verifiable-consumer",
-         "kafka-verifiable-producer", "kafka-replica-verification", ]
+         "kafka-verifiable-producer", "kafka-replica-verification", "kafka-preferred-replica-election"]
     ),
     (
         "k",
         ["ksql", "kafka-acls", "kafka-avro-console-consumer", "kafka-avro-console-producer",
-         "kafka-broker-api-versions", "kafka-replica-verification",
+         "kafka-broker-api-versions", "kafka-replica-verification", "kafka-preferred-replica-election",
          "kafka-configs", "kafka-console-consumer", "kafka-console-producer", "kafka-consumer-groups",
          "kafka-delete-records", "kafka-dump-log", "kafka-log-dirs", "kafka-topics", "kafka-verifiable-consumer",
          "kafka-verifiable-producer", "zookeeper-shell"]

tests/test_executor.py

@@ -113,6 +113,22 @@
         "kafka-consumer-groups --list --bootstrap-server test:9092",
         "kafka-consumer-groups --list --bootstrap-server test:9092"
     ),
+    (
+        "kafka-preferred-replica-election",
+        "kafka-preferred-replica-election --bootstrap-server localhost:9092"
+    ),
+    (
+        "kafka-preferred-replica-election",
+        "kafka-preferred-replica-election --bootstrap-server localhost:9092"
+    ),
+    (
+        "kafka-preferred-replica-election --bootstrap-server test:9092",
+        "kafka-preferred-replica-election --bootstrap-server test:9092"
+    ),
+    (
+        "kafka-preferred-replica-election --bootstrap-server test:9092 --admin.config test.properties",
+        "kafka-preferred-replica-election --bootstrap-server test:9092 --admin.config test.properties"
+    ),
     (
         "kafka-replica-verification",
         "kafka-replica-verification --broker-list localhost:9092"
@@ -512,6 +528,21 @@
         "kafka-log-dirs --broker-list 0,1 --command-config test.properties --bootstrap-server test:9092",
         "test-admin-client-settings"
     ),
+    (
+        "kafka-preferred-replica-election",
+        "kafka-preferred-replica-election --bootstrap-server localhost:9092 --admin.config admin.properties",
+        "test-admin-client-settings"
+    ),
+    (
+        "kafka-preferred-replica-election --admin.config testing.properties",
+        "kafka-preferred-replica-election --admin.config testing.properties --bootstrap-server localhost:9092",
+        "test-admin-client-settings"
+    ),
+    (
+        "kafka-preferred-replica-election --bootstrap-server test:9092",
+        "kafka-preferred-replica-election --bootstrap-server test:9092 --admin.config admin.properties",
+        "test-admin-client-settings"
+    ),
 ]
 
 prefix_test_data = [