uncloud-mravi/scripts/ucloud

#!/usr/bin/env python3

import argparse
import logging
import importlib
import multiprocessing as mp
import sys

import colorama

from logging.handlers import SysLogHandler

from ucloud.configure.main import configure_parser


def exception_hook(exc_type, exc_value, exc_traceback):
    logger.error(
        'Uncaught exception',
        exc_info=(exc_type, exc_value, exc_traceback)
    )
    # print('Error: ', end='')
    # print(exc_type, exc_value, exc_traceback)


class NoTracebackStreamHandler(logging.StreamHandler):
    def handle(self, record):
        info, cache = record.exc_info, record.exc_text
        record.exc_info, record.exc_text = None, None

        if record.levelname == 'WARNING':
            color = colorama.Fore.YELLOW
        elif record.levelname == 'ERROR':
            color = colorama.Fore.LIGHTRED_EX
        else:
            color = colorama.Fore.RED
        try:
            print(color)
            super().handle(record)
        finally:
            record.exc_info = info
            record.exc_text = cache
            print(colorama.Style.RESET_ALL)

if __name__ == '__main__':
    sys.excepthook = exception_hook

    arg_parser = argparse.ArgumentParser()
    subparsers = arg_parser.add_subparsers(dest="command")
    
    api_parser = subparsers.add_parser("api")
    
    host_parser = subparsers.add_parser("host")
    host_parser.add_argument("--hostname", required=True)

    scheduler_parser = subparsers.add_parser("scheduler")

    filescanner_parser = subparsers.add_parser("filescanner")

    imagescanner_parser = subparsers.add_parser("imagescanner")

    metadata_parser = subparsers.add_parser("metadata")

    config_parser = subparsers.add_parser("configure")
    configure_parser(config_parser)
    args = arg_parser.parse_args()

    if not args.command:
        arg_parser.print_help()
    else:
        # Setting up root logger
        logger = logging.getLogger('ucloud')

        syslog_handler = SysLogHandler(address='/dev/log')
        syslog_handler.setLevel(logging.DEBUG)
        syslog_formatter = logging.Formatter('%(pathname)s:%(lineno)d -- %(levelname)-8s %(message)s')
        syslog_handler.setFormatter(syslog_formatter)

        stream_handler = NoTracebackStreamHandler()
        stream_handler.setLevel(logging.WARNING)
        stream_formatter = logging.Formatter('%(message)s')
        stream_handler.setFormatter(stream_formatter)

        logger.addHandler(syslog_handler)
        logger.addHandler(stream_handler)

        # if we start etcd in seperate process with default settings
        # i.e inheriting few things from parent process etcd3 module
        # errors out, so the following command configure multiprocessing
        # module to not inherit anything from parent.
        mp.set_start_method('spawn')

        arguments = vars(args)
        try:
            name = arguments.pop('command')
            mod = importlib.import_module("ucloud.{}.main".format(name))
            main = getattr(mod, "main")
            main(**arguments)
        except Exception as err:
            logger.exception(err)