From ea766f8f71339034b453ce13b2b4a537048ec144 Mon Sep 17 00:00:00 2001 From: Oliver Bristow Date: Wed, 5 Aug 2020 00:41:20 +0100 Subject: [PATCH] Use click instead of argparse in tracker --- heron/tools/tracker/src/python/BUILD | 5 +- heron/tools/tracker/src/python/main.py | 257 +++++++++--------------- heron/tools/tracker/src/python/utils.py | 1 - 3 files changed, 101 insertions(+), 162 deletions(-) diff --git a/heron/tools/tracker/src/python/BUILD b/heron/tools/tracker/src/python/BUILD index 7b06a787fd4..14ed3bde775 100644 --- a/heron/tools/tracker/src/python/BUILD +++ b/heron/tools/tracker/src/python/BUILD @@ -7,10 +7,11 @@ pex_library( exclude = ["main.py"], ), reqs = [ - "protobuf==3.8.0", - "tornado==4.0.2", + "click==7.1.2", "javaobj-py3==0.4.1", "networkx==2.4", + "protobuf==3.8.0", + "tornado==4.0.2", ], deps = [ "//heron/common/src/python:common-py", diff --git a/heron/tools/tracker/src/python/main.py b/heron/tools/tracker/src/python/main.py index 70fa9da1948..345ee55437a 100644 --- a/heron/tools/tracker/src/python/main.py +++ b/heron/tools/tracker/src/python/main.py @@ -19,7 +19,7 @@ # under the License. ''' main.py ''' -import argparse +import logging import os import signal import sys @@ -29,14 +29,16 @@ from tornado.options import define from tornado.httpclient import AsyncHTTPClient -import heron.tools.common.src.python.utils.config as common_config -import heron.common.src.python.utils.log as log +from heron.tools.common.src.python.utils import config as common_config +from heron.common.src.python.utils import log from heron.tools.tracker.src.python import constants from heron.tools.tracker.src.python import handlers from heron.tools.tracker.src.python import utils from heron.tools.tracker.src.python.config import Config, STATEMGRS_KEY from heron.tools.tracker.src.python.tracker import Tracker +import click + Log = log.Log class Application(tornado.web.Application): @@ -93,173 +95,110 @@ def __init__(self, config): def stop(self): self.tracker.stop_sync() -# pylint: disable=protected-access -class _HelpAction(argparse._HelpAction): - """ HelpAction """ - def __call__(self, parser, namespace, values, option_string=None): - parser.print_help() - - # retrieve subparsers from parser - subparsers_actions = [ - action for action in parser._actions - if isinstance(action, argparse._SubParsersAction)] - - # there will probably only be one subparser_action, - # but better save than sorry - for subparsers_action in subparsers_actions: - # get all subparsers and print help - for choice, subparser in list(subparsers_action.choices.items()): - print("Subparser '{}'".format(choice)) - print(subparser.format_help()) - - parser.exit() - -# pylint: disable=bad-super-call -class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter): - """ Subcommand help formatter """ - def _format_action(self, action): - parts = super(argparse.RawDescriptionHelpFormatter, self)._format_action(action) - if action.nargs == argparse.PARSER: - parts = "\n".join(parts.split("\n")[1:]) - return parts - - -def add_titles(parser): - """ add titles """ - parser._positionals.title = "Required arguments" - parser._optionals.title = "Optional arguments" - return parser - - -def add_arguments(parser): - """ add arguments """ - default_config_file = os.path.join( - utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE) - - parser.add_argument( - '--config-file', - metavar='(a string; path to config file; default: "' + default_config_file + '")', - default=default_config_file) - - parser.add_argument( - '--type', - metavar='(an string; type of state manager (zookeeper or file, etc.); example: ' \ - + str(constants.DEFAULT_STATE_MANAGER_TYPE) + ')', - choices=["file", "zookeeper"]) - - parser.add_argument( - '--name', - metavar='(an string; name to be used for the state manager; example: ' \ - + str(constants.DEFAULT_STATE_MANAGER_NAME) + ')') - - parser.add_argument( - '--rootpath', - metavar='(an string; where all the states are stored; example: ' \ - + str(constants.DEFAULT_STATE_MANAGER_ROOTPATH) + ')') - - parser.add_argument( - '--tunnelhost', - metavar='(an string; if ssh tunneling needs to be established to connect to it; example: ' \ - + str(constants.DEFAULT_STATE_MANAGER_TUNNELHOST) + ')') - - parser.add_argument( - '--hostport', - metavar='(an string; only used to connect to zk, must be of the form \'host:port\';'\ - ' example: ' + str(constants.DEFAULT_STATE_MANAGER_HOSTPORT) + ')') - - parser.add_argument( - '--port', - metavar='(an integer; port to listen; default: ' + str(constants.DEFAULT_PORT) + ')', - type=int, - default=constants.DEFAULT_PORT) - parser.add_argument( - '--verbose', - action='store_true') - - return parser - -def create_parsers(): - """ create argument parser """ - parser = argparse.ArgumentParser( - epilog='For detailed documentation, go to http://github.com/apache/incubator-heron', - usage="%(prog)s [options] [help]", - add_help=False) - - parser = add_titles(parser) - parser = add_arguments(parser) - - ya_parser = argparse.ArgumentParser( - parents=[parser], - formatter_class=SubcommandHelpFormatter, - add_help=False) - - subparsers = ya_parser.add_subparsers( - title="Available commands") - - help_parser = subparsers.add_parser( - 'help', - help='Prints help', - add_help=False) - - help_parser.set_defaults(help=True) - - subparsers.add_parser( - 'version', - help='Prints version', - add_help=True) - - return parser, ya_parser - -def define_options(port, config_file): +def define_options(port: int, config_file: str) -> None: """ define Tornado global variables """ define("port", default=port) define("config_file", default=config_file) -def create_tracker_config(namespace): + +def create_tracker_config(config_file: str, stmgr_override: dict) -> dict: # try to parse the config file if we find one - config_file = namespace["config_file"] config = utils.parse_config_file(config_file) if config is None: - Log.debug("Config file does not exists: %s" % config_file) + Log.debug(f"Config file does not exists: {config_file}") config = {STATEMGRS_KEY:[{}]} - # update the config if we have any flags - config_flags = ["type", "name", "rootpath", "tunnelhost", "hostport"] - config_to_update = config[STATEMGRS_KEY][0] - for flag in config_flags: - value = namespace.get(flag, None) - if value is not None: - config_to_update[flag] = value - + # update non-null options + config[STATEMGRS_KEY][0].update( + (k, v) + for k, v in stmgr_override.items() + if v is not None + ) return config -def main(): - """ main """ - # create the parser and parse the arguments - (parser, _) = create_parsers() - (args, remaining) = parser.parse_known_args() - if remaining == ['help']: - parser.print_help() - parser.exit() - - elif remaining == ['version']: +def show_version(_, __, value): + if value: common_config.print_build_info() - parser.exit() - - elif remaining != []: - Log.error('Invalid subcommand') - sys.exit(1) - - namespace = vars(args) - - log.set_logging_level(namespace) + sys.exit(0) + + +@click.command() +@click.option( + "--version", + is_flag=True, + is_eager=True, + expose_value=False, + callback=show_version, +) +@click.option('--verbose', is_flag=True) +@click.option( + '--config-file', + help="path to a tracker config file", + default=os.path.join(utils.get_heron_tracker_conf_dir(), constants.DEFAULT_CONFIG_FILE), + show_default=True, +) +@click.option( + '--port', + type=int, + default=constants.DEFAULT_PORT, + show_default=True, + help="local port to serve on", +) +@click.option( + '--type', + "stmgr_type", + help=f"statemanager type e.g. {constants.DEFAULT_STATE_MANAGER_TYPE}", + type=click.Choice(choices=["file", "zookeeper"]), +) +@click.option( + '--name', + help=f"statemanager name e.g. {constants.DEFAULT_STATE_MANAGER_NAME}", +) +@click.option( + '--rootpath', + help=f"statemanager rootpath e.g. {constants.DEFAULT_STATE_MANAGER_ROOTPATH}", +) +@click.option( + '--tunnelhost', + help=f"statemanager tunnelhost e.g. {constants.DEFAULT_STATE_MANAGER_TUNNELHOST}", +) +@click.option( + '--hostport', + help=f"statemanager hostport e.g. {constants.DEFAULT_STATE_MANAGER_HOSTPORT}", +) +def cli( + config_file: str, + stmgr_type: str, + name: str, + rootpath: str, + tunnelhost: str, + hostport: str, + port: int, + verbose: bool, +) -> None: + """ + A HTTP service for serving data about clusters. + + The statemanager's config from the given config file can be overrided using + options on this executable. + + """ + + log.configure(logging.DEBUG if verbose else logging.INFO) # set Tornado global option - define_options(namespace['port'], namespace['config_file']) + define_options(port, config_file) - config = Config(create_tracker_config(namespace)) + stmgr_override = { + "type": stmgr_type, + "name": name, + "rootpath": rootpath, + "tunnelhost": tunnelhost, + "hostport": hostport, + } + config = Config(create_tracker_config(config_file, stmgr_override)) # create Tornado application application = Application(config) @@ -278,15 +217,15 @@ def signal_handler(signum, frame): signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) - Log.info("Running on port: %d", namespace['port']) - if namespace["config_file"]: - Log.info("Using config file: %s", namespace['config_file']) - Log.info("Using state manager:\n" + str(config)) + Log.info("Running on port: %d", port) + if config_file: + Log.info("Using config file: %s", config_file) + Log.info(f"Using state manager:\n{config}") http_server = tornado.httpserver.HTTPServer(application) - http_server.listen(namespace['port']) + http_server.listen(port) tornado.ioloop.IOLoop.instance().start() if __name__ == "__main__": - main() + cli() # pylint: disable=no-value-for-parameter diff --git a/heron/tools/tracker/src/python/utils.py b/heron/tools/tracker/src/python/utils.py index a6d317410e7..af43e2a91a9 100644 --- a/heron/tools/tracker/src/python/utils.py +++ b/heron/tools/tracker/src/python/utils.py @@ -162,7 +162,6 @@ def parse_config_file(config_file): if not os.path.lexists(expanded_config_file_path): return None - configs = {} # Read the configuration file with open(expanded_config_file_path, 'r') as f: configs = yaml.load(f)