Source code for harvester.cli

"""
cli.py
======

Contains command line interface
"""

import logging.config
import sys
from os.path import join, dirname
from typing import Tuple

import click
from redis import Redis
import structlog
import structlog.contextvars
from vs_common.config import load_config, validate_config

from .app import App
from .harvester import main
from .model import HarvesterAppConfig

SCHEMA_FILE = join(dirname(__file__), "config-schema.json")


def setup_logging(debug=False):
    logging.config.dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {
                "json": {
                    "()": structlog.stdlib.ProcessorFormatter,
                    "processor": structlog.dev.ConsoleRenderer()
                    if debug
                    else structlog.processors.JSONRenderer(),
                },
            },
            "handlers": {
                "console": {
                    "class": "logging.StreamHandler",
                    "level": "DEBUG" if debug else "INFO",
                    "formatter": "json",
                },
            },
            "root": {
                "handlers": ["console"],
                "level": "DEBUG" if debug else "INFO",
            },
        }
    )

    structlog.configure(
        processors=[
            structlog.contextvars.merge_contextvars,
            structlog.stdlib.filter_by_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.stdlib.add_logger_name,
            structlog.stdlib.add_log_level,
            structlog.stdlib.PositionalArgumentsFormatter(),
            structlog.processors.StackInfoRenderer(),
            structlog.processors.format_exc_info,
            structlog.processors.UnicodeDecoder(),
            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
        ],
        context_class=structlog.threadlocal.wrap_dict(dict),
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
        cache_logger_on_first_use=True,
    )


@click.group()
@click.option("--debug/--no-debug", default=False)
def cli(debug: bool = False):
    setup_logging(debug)


@cli.command()
@click.argument("listen_queue", default="harvester_queue", type=str)
@click.option(
    "--config-file",
    type=click.File("r"),
    default="/config.yaml",
    help="Path to config file. Defaults to /config.yaml",
)
@click.option(
    "--validate/--no-validate",
    type=bool,
    default=False,
    help="Perform configuration validation. Defaults to False",
)
@click.option(
    "-co",
    "--config-override",
    multiple=True,
    help="Dot notaion override. EXAMPLE redis.port=6378",
)
def daemon(
    listen_queue: str,
    config_file: str,
    validate: bool,
    config_override: Tuple[str],
):
    """\b
    Run the harvester daemon, attaching to a Redis queue
    ARGS: LISTEN_QUEUE - Name of queue to listen to. Defaults to harvester_queue
    """
    if validate:
        validate_config(config_file, SCHEMA_FILE)
    config = load_config(config_file, HarvesterAppConfig, config_override)

    client = Redis(
        config.redis.host, config.redis.port, charset="utf-8", decode_responses=True
    )

    app = App(config, client, listen_queue)

    while not app.shutdown:
        app.run()

    sys.exit(0)


@cli.command()
@click.argument("harvester_name", type=str)
@click.option(
    "--config-file",
    type=click.File("r"),
    default="/config.yaml",
    help="Path to config file. Defaults to /config.yaml",
)
@click.option(
    "--validate/--no-validate",
    type=bool,
    default=False,
    help="Perform configuration validation. Defaults to False",
)
@click.option(
    "-co",
    "--config-override",
    multiple=True,
    help="Dot notaion override. EXAMPLE redis.port=6378",
)
def harvest(
    harvester_name: str,
    config_file: str,
    validate: bool,
    config_override: Tuple[str],
):
    """\b
    Run a single, one-off harvest
    ARGS: HARVESTER_NAME - Name of harvester"
    """

    if validate:
        validate_config(config_file, SCHEMA_FILE)
    config = load_config(config_file, HarvesterAppConfig, config_override)

    main(config, harvester_name)


if __name__ == "__main__":
    cli()