Download ZIP archive with a MongoDB database, unpack it and import it to the local MongoDB instance

mongounzip.py

# -*- coding: utf-8 -*-
'''Download ZIP archive with a MongoDB database, unpack it
and import it to the local MongoDB instance
Usage: python mongounzip.py
'''
import os
import os.path
import shutil
import subprocess
import sys
import tempfile
import zipfile

import requests
from requests_ntlm import HttpNtlmAuth


SOURCE_URL = 'http://files.example.com/data/db/{dbname}.latest.zip'
DATABASES = ['database1', 'database2']
SITE_LOGIN = 'MYSERV\\siteuser'
SITE_PASSWORD = 'KKjdhh8y'
MONGORESTORE = 'C:\\mongodb\\bin\\mongorestore.exe' # Windows


def download(url, to_file):
    '''Download ZIP file from URL and store it in the file-like object.'''
    res = requests.get(url, auth=HttpNtlmAuth(SITE_LOGIN, SITE_PASSWORD))
    if res.status_code == 200:
        to_file.write(res.content)


def unzip(ziparc):
    '''Unpack ZIP archive to the archive folder.'''
    rootdir, _ = os.path.split(ziparc.name)
    zfile = zipfile.ZipFile(ziparc)
    for arpath in zfile.namelist():
        dirname, filename = os.path.split(arpath)
        absdirpath = os.path.join(rootdir, dirname)
        if not os.path.exists(absdirpath):
            os.mkdir(absdirpath)
        abspath = os.path.join(absdirpath, filename)
        with open(abspath, 'wb') as outfile:
            outfile.write(zfile.read(arpath))


def restore(dumpdir):
    '''Restore MongoDB databases from dubp folder with --drop option.'''
    args = [MONGORESTORE, '--drop', dumpdir]
    subprocess.call(args)


def main():
    '''The entry point of the script.'''
    temp_dir = tempfile.mkdtemp()
    try:
        for database in DATABASES:
            url = SOURCE_URL.format(dbname=database)
            arname = url.split('?')[0].split('/')[-1]
            arname = os.path.join(temp_dir, arname)
            with open(arname, 'w+b') as tempzip:
                download(url, tempzip)
                unzip(tempzip)
            os.remove(arname)
        restore(temp_dir)
    except StandardError as err:
        print >> sys.stderr, str(err)
        return 1
    finally:
        shutil.rmtree(temp_dir)


if __name__ == '__main__':
    sys.exit(main())