viewsize_v1.py

#!/usr/bin/env python

import argparse
import sys
import requests
import couchdb
import random
import string
import uuid
import time
import copy
import itertools

URL = 'http://adm:pass@127.0.0.1:15984'
DBNAME = 'thedb'
TEMPLATE =  "%-10s %-10s %-10s"

def add_view(db, args):
    if args.reduce_fun:
        view = {'reduce': args.reduce_fun}
    else:
        view = {}
    if args.map_fun:
        view['map'] = args.map_fun
    else:
        view['map'] = 'function(d){emit("K"+d._id, d.v);}'
    db['_design/des1'] = {"views": {"v1": view}}

def main(args):
    if not args.param:
        return run(args)
    param_names = []
    param_values = []
    for p in args.param.split('&'):
        n, vs = _parse_param(p)
        param_names.append(n)
        param_values.append(vs)
    print TEMPLATE % ("file", "active", "external")
    for vtup in itertools.product(*param_values):
        zipped = zip(param_names, vtup)
        paramstr = ",".join(["%s=%s" % (n, v) for (n, v) in zipped])
        run_args = copy.copy(args)
        for (n, v) in zipped:
            setattr(run_args, n, v)
        run(run_args, paramstr)
    print

def run(args, paramstr):
    random.seed(args.random_seed)
    s = couchdb.Server(args.url)
    version = s.version()
    if args.dbname in s:
        s.delete(args.dbname)
    res = requests.put(args.url + '/%s?q=1&n=1' % args.dbname)
    res.raise_for_status()
    db = s[args.dbname]
    add_view(db, args)
    n = args.num
    b = args.batch_size
    for i in xrange(n / b):
       db.update([_doc(i * b + j, args) for j in xrange(b)])
       if args.query:
           len(db.view('des1/v1'))
    stopped_at = n - n % b
    db.update([_doc(i, args) for i in xrange(stopped_at, n)])
    len(db.view('des1/v1'))
    time.sleep(args.view_query_timeout)
    sizes = db.info('des1')['view_index']['sizes']
    fsize, asize, esize = sizes['file'], sizes['active'], sizes['external']

    print TEMPLATE % (fsize, asize, esize), paramstr

def _parse_param(pstr):
    pstr = pstr.strip()
    pname,vstr = pstr.split('=', 1)
    pname = pname.replace('-','_')
    vlist = vstr.split('|')
    return pname, [_parse_value(v) for v in vlist]

def _parse_value(vstr):
    vstr = vstr.strip()
    vstr = vstr.strip("'")
    vstr = vstr.strip('"')
    if vstr.lower() == "none":
        return None
    if vstr.lower() == "null":
        return None
    if vstr.lower() == "true":
        return True
    if vstr.lower() == "false":
        return False
    try:
        return int(vstr)
    except ValueError:
        pass
    try:
        return float(vstr)
    except ValueError:
        pass
    return vstr

def _doc(i, args):
   return {'_id': _id(i, args), 'v': _data(args)}

def _data(args):
    if args.alphabet:
        alphabet = args.alphabet
    else:
        alphabet = string.ascii_letters + string.digits
    return ''.join(random.choice(alphabet) for _ in xrange(args.size))

def _id(i, args):
    if args.random_keys:
        key = uuid.uuid4().hex
    else:
        key = '%06d' % i
    extend = args.min_key_size - len(key)
    if extend > 0:
        key = key + 'x' * extend
    return key

def _args():
    description = "Make a view, add docs and measure size"
    p = argparse.ArgumentParser(description = description)
    p.add_argument('-u', '--url', default=URL,
                   help = "Server URL")
    p.add_argument('-d', '--dbname', default=DBNAME,
                   help = "DB name")
    p.add_argument('-n', '--num', type=int, default=1000,
                   help = "Number of documents")
    p.add_argument('-b', '--batch-size', type=int, default=1,
                   help = "Batch size")
    p.add_argument('-s', '--size', type=int, default=1,
                   help = "Emit value size")
    p.add_argument('-q', '--query', action="store_true", default=False,
                   help = "Query view after every batch insert")
    p.add_argument('-k', '--min-key-size', type=int, default=1,
                   help = "Min key size. If key is smaller it will be expanded")
    p.add_argument('-x', '--random-keys', action="store_true", default=False,
                   help = "Use random (UUID4) ids")
    p.add_argument('-a', '--alphabet', default=None,
                   help = "Alphabet used to generate random data from")
    p.add_argument('-t', '--view-query-timeout', type=int, default=6,
                   help = "Hold-off used by db to commit view changes")
    p.add_argument('-m', '--map-fun', default=None,
                   help = "Use a custom map function body")
    p.add_argument('-r', '--reduce-fun', default=None,
                   help = "Provide a custom reduce function body")
    p.add_argument('-p', '--param', default=None,
                   help = "A list of parameters. Separated by ';'. Each looks "
                   "like param-name=v1|v2|v3|... For example, "
                   "-p 'size=10|1000; query=True|False'")
    return p.parse_args()

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