dbtumble.py

#!/usr/bin/env python

"""
Examples:
 1. Using 50 concurrent workers create 100k dbs
    $ dbtumble.py -u https://user:pass@host -w 50 -d 100000 -c

 2. Open 100k dbs continuously by getting their info
    $ dbtumble.py -u https://user:pass@host -w 50 -d 100000 -l
"""

import sys
import time
import argparse
import threading
import couchdb
from multiprocessing.dummy import Pool as ThreadPool

DB_URL = 'http://adm:pass@localhost:15984'


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

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


_conns = {}
def get_conn(url):
    tid = threading.current_thread().ident
    if tid not in _conns:
        _conns[tid] =  couchdb.Server(url)
    return _conns[tid]


def creator((i, args)):
    srv = get_conn(args.url)
    dbname = get_dbname(i)
    srv.create(dbname)
    logp("CREATED", get_dbname(i))


def opener((i, args)):
    srv = get_conn(args.url)
    dbname = get_dbname(i)
    db = couchdb.Database(srv.resource(dbname), dbname)
    db.info()
    logp("DBINFO ", dbname)


def cleaner((i, args)):
    srv = get_conn(args.url)
    dbname = get_dbname(i)
    if dbname in srv:
        srv.delete(dbname)
        logp("DELETED ", dbname)

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


def _args():
    description = "Create lots of DBs. DBs are deleted if they already exist"
    p = argparse.ArgumentParser(description = description)
    p.add_argument('-u', '--url', default=DB_URL, help = "Server URL")
    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 run(worker, args):
    wcount = args.worker_count
    pool = ThreadPool(wcount)
    worker_args = ((i, args) for i in xrange(args.dbs))
    t0 = time.time()
    resg = pool.imap_unordered(worker, worker_args)
    resl = len(list(resg))
    dt = time.time() - t0
    log()
    log("Got ", resl, "results in %.2f seconds" % dt)
    pool.close()
    pool.join()


def main(args):
    if args.create:
        run(cleaner, args)
        run(creator, args)
        return
    if args.loop:
        while True:
            run(opener, args)
            log()
    else:
        run(opener, args)


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