ycsb

#!/usr/bin/env python

import argparse
import io
import os
import shlex
import sys
import subprocess

BASE_URL = "https://github.com/brianfrankcooper/YCSB/tree/master/"
COMMANDS = {
    "shell" : {
        "command"     : "",
        "description" : "Interactive mode",
        "main"        : "com.yahoo.ycsb.CommandLine",
    },
    "load" : {
        "command"     : "-load",
        "description" : "Execute the load phase",
        "main"        : "com.yahoo.ycsb.Client",
    },
    "run" : {
        "command"     : "-t",
        "description" : "Execute the transaction phase",
        "main"        : "com.yahoo.ycsb.Client",
    },
}

DATABASES = {
    "accumulo"     : "com.yahoo.ycsb.db.AccumuloClient",
    "basic"        : "com.yahoo.ycsb.BasicDB",
    "cassandra-7"  : "com.yahoo.ycsb.db.CassandraClient7",
    "cassandra-8"  : "com.yahoo.ycsb.db.CassandraClient8",
    "cassandra-10" : "com.yahoo.ycsb.db.CassandraClient10",
    "cassandra-cql": "com.yahoo.ycsb.db.CassandraCQLClient",
    "couchbase"    : "com.yahoo.ycsb.db.CouchbaseClient",
    "dynamodb"     : "com.yahoo.ycsb.db.DynamoDBClient",
    "elasticsearch": "com.yahoo.ycsb.db.ElasticSearchClient",
    "gemfire"      : "com.yahoo.ycsb.db.GemFireClient",
    "hbase"        : "com.yahoo.ycsb.db.HBaseClient",
    "hbase-10"     : "com.yahoo.ycsb.db.HBaseClient10",
    "hypertable"   : "com.yahoo.ycsb.db.HypertableClient",
    "infinispan"   : "com.yahoo.ycsb.db.InfinispanClient",
    "jdbc"         : "com.yahoo.ycsb.db.JdbcDBClient",
    "mapkeeper"    : "com.yahoo.ycsb.db.MapKeeperClient",
    "mongodb"      : "com.yahoo.ycsb.db.MongoDbClient",
    "mongodb-async": "com.yahoo.ycsb.db.AsyncMongoDbClient",
    "nosqldb"      : "com.yahoo.ycsb.db.NoSqlDbClient",
    "orientdb"     : "com.yahoo.ycsb.db.OrientDBClient",
    "redis"        : "com.yahoo.ycsb.db.RedisClient",
    "tarantool"    : "com.yahoo.ycsb.db.TarantoolClient",
    "voldemort"    : "com.yahoo.ycsb.db.VoldemortClient"
}

OPTIONS = {
    "-P file"        : "Specify workload file",
    "-p key=value"   : "Override workload property",
    "-s"             : "Print status to stderr",
    "-target n"      : "Target ops/sec (default: unthrottled)",
    "-threads n"     : "Number of client threads (default: 1)",
    "-cp path"       : "Additional Java classpath entries",
    "-jvm-args args" : "Additional arguments to the JVM",
}

def usage():
    output = io.BytesIO()
    print >> output, "%s command database [options]" % sys.argv[0]

    print >> output, "\nCommands:"
    for command in sorted(COMMANDS.keys()):
        print >> output, "    %s %s" % (command.ljust(14),
                                        COMMANDS[command]["description"])

    print >> output, "\nDatabases:"
    for db in sorted(DATABASES.keys()):
        print >> output, "    %s %s" % (db.ljust(14), BASE_URL +
                                        db.split("-")[0])

    print >> output, "\nOptions:"
    for option in sorted(OPTIONS.keys()):
        print >> output, "    %s %s" % (option.ljust(14), OPTIONS[option])

    print >> output, """\nWorkload Files:
    There are various predefined workloads under workloads/ directory.
    See https://github.com/brianfrankcooper/YCSB/wiki/Core-Properties
    for the list of workload properties."""

    return output.getvalue()


def find_jars(dir, database):
    jars = []
    for (dirpath, dirnames, filenames) in os.walk(dir):
        for filename in filenames:
            if filename.endswith(".jar"):
                jars.append(os.path.join(dirpath, filename))
    return jars


def get_ycsb_home():
    dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    while "CHANGELOG" not in os.listdir(dir):
        dir = os.path.join(dir, os.path.pardir)
    return os.path.abspath(dir)


def main():
    p = argparse.ArgumentParser(
            usage=usage(),
            formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('-cp', dest='classpath', help="""Additional classpath
                   entries, e.g.  '-cp /tmp/hbase-1.0.1.1/conf'. Will be
                   prepended to the YCSB classpath.""")
    p.add_argument("-jvm-args", default=[], type=shlex.split,
                   help="""Additional arguments to pass to 'java', e.g.
                   '-Xmx4g'""")
    p.add_argument("command", choices=sorted(COMMANDS),
                   help="""Command to run.""")
    p.add_argument("database", choices=sorted(DATABASES),
                   help="""Database to test.""")
    args, remaining = p.parse_known_args()
    ycsb_home = get_ycsb_home()

    # Use JAVA_HOME to find java binary if set, otherwise just use PATH.
    java = "java"
    java_home = os.getenv("JAVA_HOME")
    if java_home:
        java = os.path.join(java_home, "bin", "java")
    db_classname = DATABASES[args.database]
    command = COMMANDS[args.command]["command"]
    main_classname = COMMANDS[args.command]["main"]
    db_dir = os.path.join(ycsb_home, args.database + "-binding")
    cp = [ os.path.join(db_dir, "conf") ]
    cp.extend(find_jars(os.path.join(ycsb_home, "lib"), args.database))
    cp.extend(find_jars(os.path.join(db_dir, "lib"), args.database))
    classpath = os.pathsep.join(cp)
    if args.classpath:
        classpath = os.pathsep.join([args.classpath, classpath])

    ycsb_command = ([java] + args.jvm_args +
                    ["-cp", classpath,
                     main_classname, "-db", db_classname] + remaining)
    if command:
        ycsb_command.append(command)
    print >> sys.stderr, " ".join(ycsb_command)
    return subprocess.call(ycsb_command)


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