aws-bin/aws_kafka.py

#!/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...')