| #!/usr/bin/env python3 |
| |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import pprint |
| import traceback |
| from cmd import Cmd |
| import boto3 |
| |
| |
| class AwsKafkaCmd(Cmd): |
| """ |
| Refer to AWS document on Kafka client interface: |
| https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/kafka.html |
| """ |
| |
| def __init__(self): |
| super().__init__() |
| self.client = boto3.client('kafka') |
| self.cluster_info_list = None |
| self.current_cluster_info = None |
| |
| def onecmd(self, line): |
| """Override the command processor so it does not exit on exception in the command""" |
| try: |
| return super().onecmd(line) |
| except Exception as ex: |
| traceback.print_exc() |
| return False # don't stop |
| |
| def _print_pretty_dict(self, d, indent=4): |
| """Print a dictionary in a pretty format""" |
| if not d: |
| return "" |
| if isinstance(d, list): |
| return "\n" + "\t" * indent + "\n ".join([f'({i}) {self._print_pretty_dict(item)}' for i, item in enumerate(d)]) |
| elif isinstance(d, dict): |
| return pprint.pformat(d, indent=indent) |
| else: |
| return str(d) |
| |
| def _confirm(self, prompt="Confirm "): |
| confirm = input(prompt + "(yes/no)" ) |
| return confirm == "yes" |
| |
| def get_cluster_names(self): |
| return [x['ClusterName'] for x in self.get_cluster_info_list()] |
| |
| def get_cluster_info_list(self): |
| if not self.cluster_info_list: |
| print("Finding current list of clusters") |
| self.do_list(None) |
| return self.cluster_info_list |
| |
| def set_current_cluster(self, cluster_name): |
| cluster_names = self.get_cluster_names() |
| for cluster_info in self.get_cluster_info_list(): |
| if cluster_name == cluster_info['ClusterName']: |
| self.current_cluster_info = cluster_info |
| return True |
| else: # try partial match |
| print(f'No cluster-name matched "{cluster_name}", trying partial match') |
| for cluster_info in self.get_cluster_info_list(): |
| if cluster_name in cluster_info['ClusterName']: |
| print(f'Selecting cluster "{cluster_info["ClusterName"]}" as partial match to {cluster_name}') |
| self.current_cluster_info = cluster_info |
| return True |
| return False |
| |
| def do_info(self, args): |
| """List session information like clusters, selected cluster""" |
| print(f'Clusters: {self.get_cluster_names()}') |
| if self.current_cluster_info: |
| print(f'Current cluster: {self.current_cluster_info["ClusterName"]}') |
| else: |
| print("Current cluster: None") |
| |
| def do_list(self, args): |
| """Lists the clusters""" |
| clusters = self.client.list_clusters() |
| if clusters: |
| self.cluster_info_list = clusters['ClusterInfoList'] |
| if self.cluster_info_list: |
| for i, cluster_info in enumerate(self.cluster_info_list): |
| cluster_name = cluster_info["ClusterName"] |
| cluster_arn = cluster_info["ClusterArn"] |
| res = self._print_pretty_dict(cluster_info, indent=4) |
| print(f'({i})\t{cluster_name}\t{cluster_arn}\n\n\t{res}\n') |
| else: |
| print('No Kafka clusters found') |
| else: |
| print('No response from AWS') |
| |
| def run_client_method_and_print_tag(self, cluster_name, client_method_name, result_tag, result_label=None, keyword_args=None): |
| """ |
| Run a client method for the supplied cluster name, or current cluster name or all clusters. |
| If cluster_name is supplied (in the args) then use it. Otherwise if the current cluster is set, |
| return information on it. Otherwise, return the same information on all the clusters. |
| |
| :param cluster_name: arguments to the command, typically the name of the cluster |
| :param client_method_name: method on the client to call |
| :param result_tag: tag in the response that we are interested in |
| :param result_label: Descriptive label for the result |
| :param keyword_args: |
| :return: |
| """ |
| client_func = getattr(self.client, client_method_name) |
| keyword_args = dict(keyword_args) if keyword_args else {} |
| use_args = False if keyword_args else True # do not use args if keyword_args is supplied |
| if not result_label: |
| result_label = result_tag |
| if cluster_name: |
| for cluster_info in self.get_cluster_info_list(): |
| if cluster_name in cluster_info['ClusterName']: |
| keyword_args['ClusterArn'] = cluster_info['ClusterArn'] |
| result = client_func(**keyword_args) |
| info = self._print_pretty_dict(result[result_tag] if (result and result_tag) else result) |
| print(f'Cluster {cluster_info["ClusterName"]} {result_label}={info}') |
| else: |
| if self.current_cluster_info: |
| keyword_args['ClusterArn'] = self.current_cluster_info['ClusterArn'] |
| result = client_func(**keyword_args) |
| info = self._print_pretty_dict(result[result_tag] if (result and result_tag) else result) |
| print(f'Cluster {self.current_cluster_info["ClusterName"]} {result_label}={info}') |
| else: |
| for cluster_info in self.get_cluster_info_list(): |
| keyword_args['ClusterArn'] = cluster_info['ClusterArn'] |
| result = client_func(**keyword_args) |
| info = self._print_pretty_dict(result[result_tag] if (result and result_tag) else result) |
| print(f'Cluster {cluster_info["ClusterName"]} {result_label}={info}') |
| |
| def do_bootstrap_brokers(self, args): |
| """List Bootstrap brokers for the supplied cluster_name or or current cluster name or all clusters""" |
| self.run_client_method_and_print_tag(args, 'get_bootstrap_brokers', 'BootstrapBrokerStringSaslIam', |
| result_label='brokers') |
| |
| def do_describe(self, args): |
| """Describe the supplied cluster_name or or current cluster name or all clusters""" |
| self.run_client_method_and_print_tag(args, 'describe_cluster', 'ClusterInfo', |
| result_label='brokers') |
| |
| def do_describev2(self, args): |
| """Describe the supplied cluster_name or or current cluster name or all clusters""" |
| self.run_client_method_and_print_tag(args, 'describe_cluster_v2', 'ClusterInfo', |
| result_label='brokers') |
| |
| def do_compatible_kafka_versions(self, args): |
| """Gets the Apache Kafka versions to which you can update the MSK cluster""" |
| self.run_client_method_and_print_tag(args, 'get_compatible_kafka_versions', 'CompatibleKafkaVersions') |
| |
| def do_list_cluster_operations(self, args): |
| """Returns a list of all the operations that have been performed on the specified MSK cluster.""" |
| self.run_client_method_and_print_tag(args, 'list_cluster_operations', 'ClusterOperationInfoList') |
| |
| def do_list_configurations(self, args): |
| """Returns a list of all the MSK configurations in this Region""" |
| self.run_client_method_and_print_tag(args, 'list_configurations', 'Configurations') |
| |
| def do_list_kafka_versions(self, args): |
| """Returns a list of Apache Kafka versions""" |
| self.run_client_method_and_print_tag(args, 'list_kafka_versions', 'KafkaVersions') |
| |
| def do_list_nodes(self, args): |
| """Returns a list of the broker nodes in the cluster""" |
| self.run_client_method_and_print_tag(args, 'list_nodes', 'NodeInfoList') |
| |
| def do_list_scram_secrets(self, args): |
| """Returns a list of the Scram Secrets associated with an Amazon MSK cluster""" |
| self.run_client_method_and_print_tag(args, 'list_scram_secrets', 'SecretArnList') |
| |
| def do_reboot_broker(self, args): |
| """Reboots brokers""" |
| if not self.current_cluster_info: |
| print('Use select_cluster command first') |
| return |
| if self._confirm(prompt=f"Confirm reboot of broker {args}"): |
| keyword_args = { |
| 'BrokerIds': [args], |
| 'ClusterArn': self.current_cluster_info['ClusterArn']} |
| self.run_client_method_and_print_tag(None, 'reboot_broker', result_tag=None, keyword_args=keyword_args) |
| |
| def do_update_broker_count(self, args): |
| """Updates the number of broker nodes in the cluster""" |
| if not self.current_cluster_info: |
| print('Use select_cluster command first') |
| return |
| if self._confirm(prompt=f"Confirm updating broker count to {args}"): |
| keyword_args = { |
| 'ClusterArn': self.current_cluster_info['ClusterArn'], |
| 'CurrentVersion': self.current_cluster_info['CurrentVersion'], |
| 'TargetNumberOfBrokerNodes': int(args) |
| } |
| self.run_client_method_and_print_tag(None, 'update_broker_count', result_tag=None, keyword_args=keyword_args) |
| |
| def do_topics(self, args): |
| """Lists the topics in the current cluster""" |
| print("Not implemented yet - need to go thru MSK Client EC2 instance") |
| |
| def do_select_cluster(self, args): |
| """Set the current cluster to the specified name if valid""" |
| cluster_names = self.get_cluster_names() |
| if len(self.get_cluster_info_list()) == 0: |
| print("No clusters exist") |
| return |
| if args: |
| if self.set_current_cluster(args): |
| return |
| print(f'Cluster "{args}" not not exist') |
| print(f'syntax: select_cluster {cluster_names}') |
| |
| def do_quit(self, args): |
| """Quits the program.""" |
| print("Quitting.") |
| raise SystemExit |
| |
| |
| if __name__ == '__main__': |
| prompt = AwsKafkaCmd() |
| prompt.prompt = '> ' |
| prompt.cmdloop('Starting prompt...') |
