barman-montagu

#!/usr/bin/env python3
"""Set up and use barman (Postgres streaming backup) for montagu

Usage:
  barman-montagu setup [options] --slot=<slot> <host>
  barman-montagu status
  barman-montagu barman [--] [<args>...]
  barman-montagu recover [--wipe-first] [<backup-id>]
  barman-montagu update-nightly
  barman-montagu wipe-recover
  barman-montagu destroy

Options:
  --password-group=<group>  Password group [default: production]
  --image-tag=<tag>         Barman image tag [default: master]
  --image-source=<repo>     The Docker registry to pull from
                            [default: vimc]
  --pull-image              Pull image before running?
  --no-clean-on-error       Don't clean up container/volumes on error
  --slot=<slot>             Replication slot to use on the db
  --volume_data=<volume>    Volume name or path for data [default: barman_data]
  --volume_logs=<volume>    Volume name or path for logs [default: barman_logs]
  --volume_recover=<volume> Volume name or path for recovery
                            [default: barman_recover]
  --volume_nightly<volume>  Volume name or path for nightly recover
                            [default: montagu_db_volume]
"""

import os
import socket
import subprocess

import docker
import docopt
import hvac

# These defaults are all fairly montagu specific
DEFAULT_VAULT_ADDR = "https://vault.dide.ic.ac.uk:8200"
BARMAN_CONTAINER = 'barman-montagu'
# The 'volume' element is the name of the volume (so can be adjusted)
# but the 'path' element corresponds to a path within the container
# and are assumed by either barman (for data and logs) or us (for
# recover).
BARMAN_VOLUMES = {
    'data': {'volume': 'barman_data', 'path': '/var/lib/barman'},
    'logs': {'volume': 'barman_logs', 'path': '/var/log/barman'},
    'metrics': {'volume': 'barman_metrics', 'path': '/metrics'},
    'recover': {'volume': 'barman_recover', 'path': '/recover'},
    'nightly': {'volume': 'montagu_db_volume', 'path': '/nightly'}
}
DATABASE_NAME = "montagu"
DOCKER_IMAGE = "montagu-barman"


def status():
    barman = get_barman_container(error=False)
    if barman:
        print("{} is running".format(BARMAN_CONTAINER))
        exec_safely(barman, ["barman", "status", DATABASE_NAME],
                    check=True, stream=True)
    else:
        print("{} is not running".format(BARMAN_CONTAINER))


def barman(args):
    barman = get_barman_container()
    exec_safely(barman, ["barman"] + args, check=True, stream=True)


def recover(id, wipe_first):
    barman = get_barman_container()
    if not id:
        id = "latest"
    if wipe_first:
        wipe_recover(barman)
    exec_safely(barman, ["barman", "recover", DATABASE_NAME, id,
                         BARMAN_VOLUMES['recover']['path']],
                stream=True, check=True)
    print("\nThe volume '{}' now contains the recovered PostgreSQL data".format(
        BARMAN_VOLUMES["recover"]['volume']))


def update_nightly():
    barman = get_barman_container()
    wipe_recover(barman)
    print("Dumping backup into nightly volume")
    exec_safely(barman, ["barman", "recover", DATABASE_NAME, "latest",
                         "/nightly/"],
                stream=True, check=True)
    replay_wal_nightly()
    exec_safely(barman, ["write-nightly-timestamp"], stream=True, check=True)


def replay_wal_nightly():
    image = "vimc/montagu-db:master"
    d = docker.client.from_env()
    print("Running postgres on restored files")
    img = d.images.pull(image)
    try:
        config = "/etc/montagu/postgresql.production.conf"
        volumes = {BARMAN_VOLUMES["nightly"]["volume"]: {"bind": "/pgdata"}}
        container = d.containers.run(img, config, detach=True, environment=env,
                                     volumes=volumes)
        print("Waiting for db to complete recovery")
        exec_safely(container, ["montagu-wait.sh", "3600"], check=True)
    finally:
        container.stop()
        container.remove()


def wipe_recover(barman=None):
    barman = barman or get_barman_container()
    print("Wiping recovery volume '{}'".format(
        BARMAN_VOLUMES["recover"]['volume']))
    exec_safely(barman, ["wipe-recover"], check=True)


def destroy():
    print("Removing barman container and volumes")
    d = docker.client.from_env()
    container_remove(BARMAN_CONTAINER, d)
    for v in BARMAN_VOLUMES.values():
        volume_remove(v['volume'], d)


def setup(host, slot, password_group, tag, registry, pull, clean_on_error):
    d = docker.client.from_env()
    barman = get_barman_container(error=False)
    if barman:
        print("{} already running ({})".format(
            BARMAN_CONTAINER, barman.short_id))
        return barman

    pw = get_barman_passwords(password_group)

    image_name = '{}:{}'.format(DOCKER_IMAGE, tag)
    if registry:
        image_name = '{}/{}'.format(registry, image_name)
    volume_mappings = {x['volume']: {'bind': x['path']}
                       for x in BARMAN_VOLUMES.values()}
    print("Volume mappings: ")
    print(volume_mappings)

    db = DATABASE_NAME.upper()
    environment = {
        # These values are substituted into montagu.json in barman.d.in, and
        # affect Barman
        'BARMAN_{}_PASS_BACKUP'.format(db): pw['barman'],
        'BARMAN_{}_PASS_STREAM'.format(db): pw['streaming_barman'],
        'BARMAN_{}_SLOT_NAME'.format(db): slot,
        # This value is used by the Flask metrics app (in ./metrics)
        'BARMAN_DATABASE_NAME': DATABASE_NAME
    }
    extra_hosts = {'db': resolve_host_ip(host)}
    network = "host" if host == "localhost" else None
    if pull:
        print("Pulling image...")
        d.images.pull(image_name)
    volumes = [volume_create(v, d) for v in named_volumes()]
    barman = d.containers.run(image_name, name=BARMAN_CONTAINER, detach=True,
                              volumes=volume_mappings, environment=environment,
                              extra_hosts=extra_hosts, network=network,
                              restart_policy={"Name": "always"})
    return barman


def named_volumes():
    return (v['volume'] for v in BARMAN_VOLUMES.values()
            if not v['volume'].startswith('/'))


def vault_client():
    vault_url = os.environ.get("VAULT_ADDR", DEFAULT_VAULT_ADDR)
    vault_token = os.environ.get("VAULT_AUTH_GITHUB_TOKEN")
    if not vault_token:
        print("Please paste your vault GitHub access token: ", end="")
        vault_token = input().strip()
    vault = hvac.Client(url=vault_url)
    print("Authenticating vault with GitHub")
    vault.auth_github(vault_token)
    return vault


def get_barman_passwords(password_group):
    users = ['barman', 'streaming_barman']

    if password_group == "fake":
        def get_password(user):
            return user
    else:
        def get_password(user):
            password = os.environ.get("MONTAGU_DB_PASSWORD_" + user)
            if password:
                return password
            else:
                vault = vault_client()
                path = "secret/vimc/database/{}/users/{}".format(password_group, user)
                return vault.read(path)['data']['password']

    return {x: get_password(x) for x in users}


def get_barman_container(docker_client=docker.client.from_env(), error=True):
    try:
        return docker_client.containers.get(BARMAN_CONTAINER)
    except docker.errors.NotFound:
        if error:
            raise Exception("{} is not running; run setup first".format(
                BARMAN_CONTAINER))
        else:
            return None


def volume_create(name, docker_client=docker.client.from_env()):
    """Create a docker volume.  Returns the a tuple of the volume and a
    boolean indicating if the volume was freshly created"""
    try:
        v = docker_client.volumes.get(name)
        return (v, False)
    except docker.errors.NotFound:
        print("Creating volume '{}'".format(name))
        v = docker_client.volumes.create(name)
        return (v, True)


def volume_remove(name, docker_client=docker.client.from_env()):
    try:
        v = docker_client.volumes.get(name)
        print("Removing volume '{}'".format(v.name))
        v.remove()
        return True
    except docker.errors.NotFound:
        return False


def container_remove(name, docker_client=docker.client.from_env()):
    try:
        x = docker_client.containers.get(name)
        print("Removing container '{}'".format(name))
        x.remove(force=True)
        return True
    except docker.errors.NotFound:
        return False


# See comments in montagu:
# https://github.com/vimc/montagu/pull/87/files#diff-6927d77cf5038b5393885eb816eb9d7a
def exec_safely(container, cmd, check=False, stream=True):
    api = container.client.api
    exec_id = api.exec_create(container.id, cmd)['Id']
    out = api.exec_start(exec_id, stream=stream)
    if stream:
        ret = []
        for line in out:
            print(line.decode("UTF-8"), end="")
            ret.append(line)
        out = b''.join(ret)
    dat = api.exec_inspect(exec_id)
    dat['Output'] = out
    if check and dat['ExitCode'] != 0:
        msg = out.decode("UTF-8").strip()
        raise Exception("Exec failed with message '{}'".format(msg))
    return dat


def resolve_host_ip(host):
    if host == "localhost" or host == "127.0.0.1":
        return "127.0.0.1"
    try:
        parsed = socket.inet_aton(host)
        if len(parsed) != 4:
            raise Exception("Expected 4 octets in host if given as IP")
        return host
    except socket.error:
        return socket.gethostbyname(host)


def handle_volume_args(args):
    for volume in ['data', 'logs', 'nightly', 'recover']:
        arg = "--volume_" + volume
        if arg in args:
            BARMAN_VOLUMES[volume]['volume'] = args[arg]


def main():
    args = docopt.docopt(__doc__)
    handle_volume_args(args)

    if args['setup']:
        setup(host=args['<host>'],
              slot=args['--slot'],
              password_group=args['--password-group'],
              tag=args['--image-tag'],
              registry=args['--image-source'],
              pull=args['--pull-image'],
              clean_on_error=not args['--no-clean-on-error'])
    elif args['status']:
        status()
    elif args['barman']:
        barman(args['<args>'])
    elif args['recover']:
        recover(args['<backup-id>'], args['--wipe-first'])
    elif args['destroy']:
        destroy()
    elif args['wipe-recover']:
        wipe_recover()
    elif args['update-nightly']:
        update_nightly()
    else:
        raise Exception("(this should never happen)")


if __name__ == "__main__":
    main()