Import MongoDB databases from tarballs

initdb.py

#!/usr/bin/env python3

import argparse
import os
import subprocess
import sys
from collections import ChainMap
from pathlib import Path
from tempfile import TemporaryDirectory

CLINICS = (
    'cdkb',
    'csm',
    'fbmse',
    # 'fmbc',
    'fnkc',
    'kb122-spb',
    'lpu165',
    'lpu169',
    'lpu85',
    'lpu86',
    'lpuccd',
    'lpuckbvl',
    'lpugncii',
    'lpuniifhm',
    'lpusernovodsk',
    'skc'
)


def callcmd(*args):
    cmd = [str(arg) for arg in args]
    subprocess.check_call(cmd, stdout=subprocess.DEVNULL)


def untar(tarball, output):
    callcmd('tar', 'zxf', tarball, '-C', output)


def dropdb(db):
    host = 'localhost/{}'.format(db)
    callcmd('mongo', host, '--eval', 'db.dropDatabase()')


def restore(dbdir, target_db):
    callcmd('mongorestore', '--db', target_db, dbdir)


def parse_args():
    defaults = {'clinic_db_dir': '.'}
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--clinic-db-dir')
    namespace = parser.parse_args()
    command_line_args = {
        key: val for key, val in vars(namespace).items() if val}
    return ChainMap(command_line_args, os.environ, defaults)


def import_dbs(clinic_db_dir):
    with TemporaryDirectory() as tempdir:
        for clinic in CLINICS:
            arcpath = Path(clinic_db_dir, '{}.{}'.format(clinic, 'tar.gz'))
            untar(arcpath, tempdir)
            dump_path = Path(tempdir).joinpath(clinic)
            for dbdir in dump_path.iterdir():
                target_db = '{}-{}'.format(clinic, dbdir.name)
                dropdb(target_db)
                restore(dbdir, target_db)


def main():
    clinic_db_dir = parse_args()['clinic_db_dir']
    try:
        import_dbs(clinic_db_dir)
    except subprocess.CalledProcessError:
        sys.exit(1)

if __name__ == '__main__':
    main()