Source code for preprocessor.cli

from os.path import join, dirname
import logging.config
import structlog

import click
import yaml
import jsonschema

from .preprocess import preprocess_file
from .daemon import run_daemon
from .config import load_config


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,
    )


def validate_config(config):
    with open(join(dirname(__file__), "config-schema.yaml")) as f:
        schema = yaml.load(f)

    jsonschema.validate(config, schema)


@click.group()
def cli():
    pass


@cli.command(help="Run the preprocess daemon, attaching to a Redis queue")
@click.option("--config-file", type=click.File("r"))
@click.option("--validate/--no-validate", default=False)
@click.option("--host", type=str)
@click.option("--port", type=int)
@click.option("--listen-queue", type=str)
@click.option("--write-queue", type=str)
@click.option("--debug/--no-debug", default=False)
def daemon(
    config_file=None,
    validate=False,
    host=None,
    port=None,
    listen_queue=None,
    write_queue=None,
    debug=False,
):
    setup_logging(debug)
    config = load_config(config_file)
    if validate:
        validate_config(config)
    run_daemon(
        config,
        host,
        port,
        listen_queue,
        write_queue,
    )


@cli.command(help="Run a single, one-off preprocessing")
@click.argument("file_path", type=str)
@click.option("--config-file", type=click.File("r"))
@click.option("--use-dir", type=str)  # TODO: check dir
@click.option("--validate/--no-validate", default=False)
@click.option("--debug/--no-debug", default=False)
def preprocess(
    file_path,
    config_file=None,
    use_dir=None,
    validate=False,
    debug=False,
):
    with structlog.contextvars.bound_contextvars(file_path=file_path):
        setup_logging(debug)
        config = load_config(config_file)
        if validate:
            validate_config(config)

        preprocess_file(config, file_path, use_dir)


if __name__ == "__main__":
    cli()