diff --git a/bin/ycsb b/bin/ycsb index 5380479e37e1198591d435a0dcd3a8c30c40ab85..5494c7671c71c6ba2076a61820aeed4f91931b7f 100755 --- a/bin/ycsb +++ b/bin/ycsb @@ -1,6 +1,9 @@ #!/usr/bin/env python +import argparse +import io import os +import shlex import sys import subprocess @@ -43,40 +46,46 @@ DATABASES = { "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", + "redis" : "com.yahoo.ycsb.db.RedisClient", "voldemort" : "com.yahoo.ycsb.db.VoldemortClient", "couchbase" : "com.yahoo.ycsb.db.CouchbaseClient" } 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)", + "-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(): - print "Usage: %s command database [options]" % sys.argv[0] + output = io.BytesIO() + print >> output, "%s command database [options]" % sys.argv[0] - print "\nCommands:" + print >> output, "\nCommands:" for command in sorted(COMMANDS.keys()): - print " %s %s" % (command.ljust(13), COMMANDS[command]["description"]) + print >> output, " %s %s" % (command.ljust(14), + COMMANDS[command]["description"]) - print "\nDatabases:" + print >> output, "\nDatabases:" for db in sorted(DATABASES.keys()): - print " %s %s" % (db.ljust(13), BASE_URL + db.split("-")[0]) + print >> output, " %s %s" % (db.ljust(14), BASE_URL + + db.split("-")[0]) - print "\nOptions:" + print >> output, "\nOptions:" for option in sorted(OPTIONS.keys()): - print " %s %s" % (option.ljust(13), OPTIONS[option]) + print >> output, " %s %s" % (option.ljust(14), OPTIONS[option]) - print """\nWorkload Files: + 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.""" - sys.exit(1) + return output.getvalue() + def find_jars(dir, database): jars = [] @@ -91,36 +100,51 @@ def find_jars(dir, database): 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) -if len(sys.argv) < 3: - usage() -if sys.argv[1] not in COMMANDS: - print "ERROR: Command '%s' not found" % sys.argv[1] - usage() -if sys.argv[2] not in DATABASES: - print "ERROR: Database '%s' not found" % sys.argv[2] - usage() - -ycsb_home = get_ycsb_home() -command = COMMANDS[sys.argv[1]]["command"] -database = sys.argv[2] -db_classname = DATABASES[database] -options = sys.argv[3:] -java_home = os.environ["JAVA_HOME"] - -if java_home: - ycsb_command = [java_home + "/bin/java", "-cp", os.pathsep.join(find_jars(ycsb_home, database)), \ - COMMANDS[sys.argv[1]]["main"], "-db", db_classname] + options -else: - ycsb_command = ["java", "-cp", os.pathsep.join(find_jars(ycsb_home, database)), \ - COMMANDS[sys.argv[1]]["main"], "-db", db_classname] + options - -if command: - ycsb_command.append(command) -print " ".join(ycsb_command) -subprocess.call(ycsb_command) + +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"] + classpath = os.pathsep.join(find_jars(ycsb_home, args.database)) + 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())