socorro-cmd

#!/usr/bin/env python

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""
Script that helps discover and run the various runnable services and scripts in
Socorro.

This command should be run in the Docker-based local development environment in
one of the Docker containers.

Run "socorro-cmd --help" for help.

"""

import argparse
import importlib
import inspect
import sys
import textwrap
import types

import click


sys.path.insert(0, "bin")
sys.path.insert(0, "webapp")


def wrapped(text, width=80):
    """Wraps a block of text"""
    return "\n".join(textwrap.wrap(text, width=width))


def showcommands_cmd(argv):
    """Sub command to show location of all command runners"""
    parser = argparse.ArgumentParser(
        description="Shows socorro-cmd subcommands and runner paths."
    )
    parser.parse_args(argv)

    print("Available commands and runners for socorro-cmd:")
    print("")

    for group in COMMANDS:
        print(f"{group.name}:")
        for cmd, runner in group:
            if not isinstance(runner, str):
                runner = f"{runner.__name__} in {inspect.getfile(runner)}"
            print(f"  {cmd:<24} {runner}")
        print("")


def django_cmd(cmd):
    """Run a Django command."""

    def _django_cmd(argv):
        argv = ["socorro-cmd", cmd] + argv
        import manage

        manage.main(argv)

    return _django_cmd


def import_path(python_path):
    """Import a name from a module."""
    module_path, name = python_path.rsplit(".", 1)
    module = importlib.import_module(module_path)
    return getattr(module, name)


class Group:
    """Defines a grouping of commands"""

    def __init__(self, name, cmd_to_run_map):
        self.name = name
        self.cmd_to_run_map = cmd_to_run_map

    def __iter__(self):
        return iter(sorted(self.cmd_to_run_map.items()))

    def lookup(self, cmd):
        return self.cmd_to_run_map.get(cmd)


COMMANDS = [
    Group(
        "Crash processing utilities",
        {
            "fetch_crashids": import_path("socorro.scripts.fetch_crashids.main"),
            "fetch_crash_data": import_path("socorro.scripts.fetch_crash_data.main"),
            "reprocess": import_path("socorro.scripts.reprocess.main"),
            "fetch_missing": import_path("socorro.scripts.fetch_missing.main"),
        },
    ),
    Group(
        "Local development environment maintenance",
        {
            "upload_telemetry_schema": import_path("upload_telemetry_schema.main"),
            "db": import_path("socorro.scripts.db.db_group"),
            "es": import_path("es_cli.main"),
            "legacy_es": import_path("legacy_es_cli.main"),
            "pubsub": import_path("pubsub_cli.main"),
            "gcs": import_path("gcs_cli.main"),
        },
    ),
    Group(
        "Services",
        {
            "cronrun": django_cmd("cronrun"),
        },
    ),
    Group(
        "Miscellaneous",
        {
            "showcommands": showcommands_cmd,
            "signature": import_path("socorro.signature.cmd_signature.main"),
            "signature-doc": import_path("socorro.signature.cmd_doc.main"),
        },
    ),
]


def import_and_run(app):
    """Takes a runner and runs it."""
    # If the app is a main function, we run it as is
    if isinstance(app, types.FunctionType):
        sys.exit(app(sys.argv[1:]))

    # If the app is a click app, then run it; click apps exit on their
    # own
    if isinstance(app, click.core.Group):
        app()

    print("ProgrammerError: Unknown runner type")
    sys.exit(1)


def build_epilog():
    """Builds the epilog containing the groups and commands"""
    output = []
    output.append("Available commands:")
    output.append("\n")
    for group in COMMANDS:
        output.append(f"{group.name}:")
        for cmd, _ in group:
            output.append(f"  {cmd}")
        output.append("\n")

    output.append("")
    output.append('Type "socorro-cmd CMD --help" for help on any command.')
    output.append("")
    output.append(
        wrapped(
            "All commands need to be run in a Docker container. Service commands are run in "
            'the related container. For example, the "processor" command should run in the '
            '"processor" container. All other commands should be run in the processor '
            "container."
        )
    )
    output.append("")
    output.append(
        wrapped("For more documentation, see <https://socorro.readthedocs.io/>.")
    )
    return "\n".join(output)


def get_runner(cmd):
    """Given a cmd, returns the runner"""
    for group in COMMANDS:
        runner = group.lookup(cmd)
        if runner is not None:
            return runner


def cmd_main():
    # Build a basic parser so we can take advantage of --help
    parser = argparse.ArgumentParser(
        prog="socorro-cmd",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=build_epilog(),
    )
    parser.add_argument("cmd", nargs="?", help="command to run")
    parser.add_argument("arg", nargs="*", help="command arguments")

    if not sys.argv[1:] or sys.argv[1] in ("-h", "--help"):
        parser.print_help()
        return 0

    cmd = sys.argv[1]
    runner = get_runner(cmd)
    if runner is None:
        parser.error(f"{cmd!r} is not a valid command; see --help for command list")

    # Rewrite sys.argv so anything that parses sys.argv has the right one
    sys.argv = [sys.argv[0] + " " + cmd] + sys.argv[2:]
    import_and_run(runner)


if __name__ == "__main__":
    cmd_main()