osc_cli

#!/usr/bin/env python3
"""
Copyright (c) 2017-present, Facebook, Inc.
All rights reserved.

This source code is licensed under the BSD-style license found in the
LICENSE file in the root directory of this source tree.
"""

import argparse
import logging
import logging.handlers
import sys
import textwrap

from collections import defaultdict
from core.commands import commands
from core.lib.error import OSCError
from core.lib import hook


def main():
    # Provide general help information for this command and a list of
    # commands provided.
    parser = argparse.ArgumentParser(description=textwrap.dedent("""
    CLI for running an online schema change (OSC) that can perform schema
    changes with minimal impact to active MySQL instances.
    """))

    choices = list(logging._nameToLevel.keys())
    parser.add_argument("--log-level", dest='level', default="INFO",
                        help="Severity levels for log output",
                        choices=choices)
    parser.add_argument("--log-file", default="/tmp/osc.log",
                        help="File path for storing all the logs")
    parser.add_argument("--quiet", action='store_true',
                        help="Avoid printing log output to screen")

    command_parser = parser.add_subparsers(title='supported commands',
                                           dest='command',
                                           metavar='mode')

    # set up a subparser for each registered command
    commands_by_name = {}
    for cmd_class in commands:
        cmd = cmd_class()
        # abort if there is no NAME property
        try:
            name = cmd.name()
        except Exception:
            print("Skipping command %s because it has no NAME" % cmd)
            continue
        help = cmd.description()
        cmd.parser = parser
        cmd.subparser = command_parser.add_parser(name, help=help)
        commands_by_name[name] = cmd
        cmd.setup_parser(
            cmd.subparser, require_user=True, require_password=True)

    args = parser.parse_args()

    # Setting up log handlers
    if isinstance(args.level, int):
        level = args.level
    else:
        level = getattr(logging, args.level)
    log = logging.getLogger()
    log.setLevel(level)
    formatter = logging.Formatter(
        fmt='%(levelname)-9s %(asctime)s.%(msecs)03d %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')
    if not args.quiet:
        scr = logging.StreamHandler()
        scr.setFormatter(formatter)
        log.addHandler(scr)
    max_bytes = 1024 * 1024 * 10
    dh = logging.handlers.RotatingFileHandler(
        args.log_file,
        maxBytes=max_bytes,
        backupCount=5
    )
    dh.setFormatter(formatter)
    log.addHandler(dh)

    try:
        if not args.command:
            parser.print_help()
            sys.exit(0)
        # Arguments sanity check
        cmd = commands_by_name[args.command]
        cmd.args = args
        cmd.validate_args()
        if args.command == 'copy':
            hook_map = defaultdict(lambda: hook.NoopHook())
            # uncomment following line if you want to add a hook config
            # hook_map['after_init_connection'] = hook.NoopHook()
            cmd.args.hook_map = hook_map

        # Execute main logic
        cmd.op()
    except OSCError as e:
        log.error(e.desc)
        sys.exit(e.code)


if __name__ == '__main__':
    main()