FabledSteward/steward/cli.py

import asyncio
import logging
import signal
from pathlib import Path

import click

from .core.logging_setup import configure_logging

# How long a stop waits for requests in flight (and SSE streams) to finish.
GRACEFUL_TIMEOUT_SECONDS = 10.0


def hypercorn_config(host: str, port: int):
    """Hypercorn as the production server (#5159): one worker, because Steward
    keeps run, stream and scheduler state in-process; its access and error
    lines go through Steward's logging, so the log.level setting governs them."""
    from hypercorn.config import Config

    config = Config()
    config.bind = [f"{host}:{port}"]
    config.workers = 1
    config.graceful_timeout = GRACEFUL_TIMEOUT_SECONDS
    config.accesslog = logging.getLogger("hypercorn.access")
    config.errorlog = logging.getLogger("hypercorn.error")
    config.access_log_format = "%(h)s %(r)s %(s)s %(b)s %(D)s"
    return config


async def _serve(app, config) -> None:
    from hypercorn.asyncio import serve

    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, stop.set)
    await serve(app, config, shutdown_trigger=stop.wait)


@click.group(invoke_without_command=True)
@click.option("--config", default="config.yaml", help="Path to config.yaml")
@click.option("--host", default="0.0.0.0")
@click.option("--port", default=5000, type=int)
@click.option("--debug", is_flag=True, default=False,
              help="Quart's development server with the reloader, not Hypercorn.")
@click.pass_context
def main(ctx: click.Context, config: str, host: str, port: int, debug: bool) -> None:
    """Start the Steward server, or run one of the commands below."""
    # Before anything logs; create_app applies the log.level setting once
    # settings load.
    configure_logging()
    ctx.obj = {"config": config}
    if ctx.invoked_subcommand is not None:
        return
    from .app import create_app
    app = create_app(config_path=Path(config))
    if debug:
        app.run(host=host, port=port, debug=True)
        return
    asyncio.run(_serve(app, hypercorn_config(host, port)))


@main.command("backup")
@click.option("--to", "out_dir", default=".", type=click.Path(file_okay=False, path_type=Path),
              help="The folder to write the backup file into.")
@click.pass_context
def backup_command(ctx: click.Context, out_dir: Path) -> None:
    """Write a backup file of this Steward (settings and the secret key: what
    recovery needs, no logs or history) into a folder, without sending it anywhere."""
    from .backup.dump import BackupError, make_backup
    from .backup.job import migration_heads
    from .config import load_bootstrap

    boot = load_bootstrap(Path(ctx.obj["config"]))
    out_dir.mkdir(parents=True, exist_ok=True)
    try:
        heads = asyncio.run(migration_heads(boot["database_url"]))
        made = asyncio.run(make_backup(boot["database_url"], boot["secret_key"], heads, out_dir))
    except BackupError as exc:
        raise click.ClickException(str(exc)) from None
    click.echo(f"Wrote {made.path} ({made.size} bytes). It holds the secret key: keep it safe.")


@main.command("restore")
@click.argument("backup_file", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--wipe", is_flag=True, default=False,
              help="Drop everything in the database first. Without it the database must be empty.")
@click.pass_context
def restore_command(ctx: click.Context, backup_file: Path, wipe: bool) -> None:
    """Bring Steward back from a backup file into its (empty) database, and
    put the backup's secret key in place. Start Steward afterwards: it runs
    any migrations the backup is behind on."""
    from .backup.dump import BackupError
    from .backup.restore import restore
    from .config import _SECRET_KEY_FILE, _env, load_bootstrap

    boot = load_bootstrap(Path(ctx.obj["config"]))
    try:
        done = asyncio.run(restore(backup_file, boot["database_url"], wipe=wipe,
                                   plugin_dirs=[Path(d) for d in boot["plugin_dirs"]]))
    except BackupError as exc:
        raise click.ClickException(str(exc)) from None
    made = done.manifest.get("created_at", "?")
    click.echo(f"Restored the backup made {made} (Steward {done.manifest.get('steward_version', '?')}).")
    if done.fixed:
        click.echo(f"Cleared {done.fixed} references to rows a backup leaves out (e.g. Ansible runs).")
    env_key = _env("SECRET_KEY")
    if env_key:
        if env_key != done.secret_key:
            click.echo("STEWARD_SECRET_KEY is set and differs from the backup's key: set it to the "
                       "backup's key, or the restored passwords and keys can't be read.", err=True)
            raise SystemExit(2)
    else:
        _SECRET_KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
        _SECRET_KEY_FILE.write_text(done.secret_key)
        _SECRET_KEY_FILE.chmod(0o600)
        click.echo(f"Wrote the backup's secret key to {_SECRET_KEY_FILE}.")
    click.echo("Start Steward now; it brings the database up to date on start.")