restore_k8s_jobs

import click

from python.import_cleanstart import KubeClient


def backup_rover(kube_client, pvc_name, backup_dir, host, password):
    postgres_image = 'postgres:12.2'
    sql_command = f'pg_dump -v -d mediaos -j 8 -F d --exclude-table-data=content_versions -f {backup_dir}'
    args = ["-c", sql_command]
    volumes = {
        'volumes': [
            {
                'name': 'rover-backup',
                'persistentVolumeClaim': {
                    'claimName': pvc_name
                }
            }
        ]
    }
    volume_mount = {
        'volumeMounts': [
            {
                'name': 'rover-backup',
                'mountPath': '/backups'
            }
        ]
    }
    config = {
        'postgres': {
            'PGPASSWORD': password,
            'PGHOST': host,
            'PGDATABASE': 'mediaos',
            'PGPORT': '5432',
            'PGUSER': 'mediaos'
        }
    }
    kube_client.create_kube_database_job(
        'rover',
        config,
        'rover-db-backup',
        'mediaos',
        args,
        namespace='rover',
        volumes=volumes,
        volume_mount=volume_mount,
        postgres_image=postgres_image
    )


def restore_rover(kube_client, pvc_name, backup_dir, host, password):
    postgres_image = 'postgres:12.2'
    sql_command = ('pg_restore --clean --if-exists --verbose --jobs=8 '
                   '--format=directory '
                   f'--dbname=mediaos --schema=public {backup_dir}')
    args = ["-c", sql_command]
    volumes = {
        'volumes': [
            {
                'name': 'rover-backup',
                'persistentVolumeClaim': {
                    'claimName': pvc_name
                }
            }
        ]
    }
    volume_mount = {
        'volumeMounts': [
            {
                'name': 'rover-backup',
                'mountPath': '/backups'
            }
        ]
    }
    config = {
        'postgres': {
            'PGPASSWORD': password,
            'PGHOST': host,
            'PGDATABASE': 'mediaos',
            'PGPORT': '5432',
            'PGUSER': 'mediaos'
        }
    }
    kube_client.create_kube_database_job(
        'rover',
        config,
        'rover-db-restore',
        'mediaos',
        args,
        namespace='rover',
        volumes=volumes,
        volume_mount=volume_mount,
        postgres_image=postgres_image
    )


def create_pv_claim(kube_client, pvc_name):
    kube_client.create_pv_claim(
        'rover',
        pvc_name,
        storage_size='500Gi'
    )


@click.command()
@click.option(
    '-kc', '--kube-context', required=True, default='kubestage-s4mnj3p3',
    help='Context to use for kube client'
)
@click.option('--db-backup-host', required=False)
@click.option('--db-backup-host-password', required=False)
@click.option('--db-restore-host', required=False)
@click.option('--db-restore-host-password', required=False)
@click.option('--backup-dir', required=False)
def main(kube_context, db_backup_host, db_backup_host_password,
         db_restore_host, db_restore_host_password, backup_dir):
    kube_client = KubeClient(kube_context)
    pvc_name = 'rover-db-backup'
    backup_dir = '/backups/rover-backup' if not backup_dir else backup_dir
    create_pv_claim(kube_client, pvc_name)
    if all([db_backup_host, db_backup_host_password]):
        backup_rover(kube_client, pvc_name, backup_dir, db_backup_host, db_backup_host_password)
    if all([db_restore_host, db_restore_host_password]):
        restore_rover(kube_client, pvc_name, backup_dir, db_restore_host, db_restore_host_password)


if __name__ == '__main__':
    main()