Apache CouchDB DB Stampede

dbcreate.py

#!/usr/bin/env python

# Uses CouchDB (pip install CouchDB)

import sys, couchdb, time, os, uuid, argparse, random, traceback
from multiprocessing.dummy import Pool as ThreadPool

DB_URLS = [
    'http://adm:pass@localhost:15984',
    'http://adm:pass@localhost:25984',
    'http://adm:pass@localhost:35984'
]

def log(*args):
    msg = " ".join(["%s" % a for a in args])
    sys.stdout.write(msg + '\n')
    sys.stdout.flush()

def pick_server(URLS):
    if isinstance(URLS, list):
        return random.choice(URLS)
    return URLS

def create_db(srv, dbname):
    return srv.create(dbname)

def worker((i, args)):
    url = pick_server(args.urls)
    srv = couchdb.Server(url)
    dbname = get_dbname(i)
    if args.create:
        create_db(srv, dbname)
    else:
        db = srv[dbname]
        db.info()

def cleaner((i, args)):
    url = pick_server(args.urls)
    srv = couchdb.Server(url)
    dbname = get_dbname(i)
    if dbname in srv:
        srv.delete(dbname)

def get_dbname(i):
    return "db_%06d" % i

def _args():
    description = "Create lots of DBs. DBs are deleted if they already exist"
    p = argparse.ArgumentParser(description = description)
    p.add_argument('-u', '--urls', action="append", default=DB_URLS, help = "Server URL(s)")
    p.add_argument('-d', '--dbs', type=int, default=1, help = "DB count")
    p.add_argument('-c', '--create', action="store_true", default=False, help = "Create the DB")
    p.add_argument('-l', '--loop', action="store_true", default=False, help = "Repeat")
    p.add_argument('-w', '--worker-count', type=int, default=1)
    return p.parse_args()

def clean(args):
    wcount = args.worker_count
    pool = ThreadPool(wcount)
    worker_args = ((i, args) for i in xrange(args.dbs))
    pool.map(cleaner, worker_args)
    pool.close()
    pool.join()

def run(args):
    wcount = args.worker_count
    pool = ThreadPool(wcount)
    worker_args = ((i, args) for i in xrange(args.dbs))
    pool.map(worker, worker_args)
    pool.close()
    pool.join()

def main(args):
    if args.create:
        clean(args)
    if args.loop:
        while True:
            run(args)
    else:
        run(args)

if __name__=='__main__':
    args = _args()
    main(_args())