Wilddog backup and restore

backup_restore.py

# coding: utf-8
from argparse import ArgumentParser
import collections
import itertools
import json
from multiprocessing.pool import ThreadPool
import os

import requests
from wilddog import WilddogApplication


def _pool(concurrency):
    if concurrency:
        return ThreadPool(concurrency)
    else:
        return itertools


def traverse(wa, wilddog_space, wilddog_path="", concurrency=None):
    def f(path):
        try:
            data = wa.get(wilddog_space, path)
        except requests.exceptions.HTTPError:
            data = wa.get(
                wilddog_space, path, params={'shallow': 'true'})
            data = [os.path.join(path, n) for n in data]
            return path, None, data
        else:
            return path, data, []

    pool = _pool(concurrency)
    next_paths = collections.deque([wilddog_path])
    while next_paths:
        for path, data, _next_paths in pool.imap(f, next_paths):
            next_paths.extend(_next_paths)
            if data:
                yield path, data


def backup(output, wilddog_space, wilddog_path, concurrency=None):
    wa = WilddogApplication(wilddog_space)
    with open(output, 'wb') as f:
        t = traverse(
            wa, wilddog_space, wilddog_path=wilddog_path,
            concurrency=concurrency)
        for path, data in t:
            data["__key__"] = path
            line = json.dumps(data)
            f.write(line)
            f.write(u"\n")
    print "DONE"


def restore(input, wilddog_space, concurrency=None):
    def func(line):
        data = json.loads(line)
        path = data.pop("__key__")
        wa.put(wilddog_space, path, data)
        return path

    pool = _pool(concurrency)
    wa = WilddogApplication(wilddog_space)
    with open(input, 'wb') as f:
        for path in pool.imap(func, f):
            print path
    print "DONE"


if __name__ == "__main__":
    parser = ArgumentParser()
    subparsers = parser.add_subparsers()

    backup_parser = subparsers.add_parser("backup")
    backup_parser.set_defaults(command="backup")
    backup_parser.add_argument("output")
    backup_parser.add_argument(
        "-ws", "--wilddog-space", dest="wilddog_space",
        default='https://site-eumt.wilddogio.com')
    backup_parser.add_argument(
        "-wp", "--wilddog-path", dest="wilddog_path", default='')
    backup_parser.add_argument(
        "--concurrency", dest="concurrency", type=int, default=10)

    restore_parser = subparsers.add_parser("restore")
    restore_parser.set_defaults(command="restore")
    restore_parser.add_argument("input")
    restore_parser.add_argument(
        "-ws", "--wilddog-space", dest="wilddog_space",
        default='https://site-eumt.wilddogio.com')
    restore_parser.add_argument(
        "--concurrency", dest="concurrency", type=int, default=10)

    args = parser.parse_args()
    print args

    if args.command == "backup":
        backup(args.output, args.wilddog_space, args.wilddog_path)
    elif args.command == "restore":
        restore(args.input, args.wilddog_space)