Skip to content
This repository has been archived by the owner on Mar 3, 2023. It is now read-only.

Commit

Permalink
Use click instead of argparse in tracker
Browse files Browse the repository at this point in the history
  • Loading branch information
Code0x58 committed Aug 5, 2020
1 parent 9d4541b commit 23345dd
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 162 deletions.
5 changes: 3 additions & 2 deletions heron/tools/tracker/src/python/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
256 changes: 97 additions & 159 deletions heron/tools/tracker/src/python/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
# under the License.

''' main.py '''
import argparse
import logging
import os
import signal
import sys
Expand All @@ -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):
Expand Down Expand Up @@ -93,173 +95,109 @@ 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.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)
Expand All @@ -278,15 +216,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
1 change: 0 additions & 1 deletion heron/tools/tracker/src/python/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 23345dd

Please sign in to comment.