diff --git a/bin/ycsb b/bin/ycsb index 5380479e..5494c767 100755 --- a/bin/ycsb +++ b/bin/ycsb @@ -1,126 +1,150 @@ #!/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", "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", + "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 = [] for (dirpath, dirnames, filenames) in os.walk(dir): if dirpath.endswith("conf"): jars.append(dirpath) for filename in filenames: if filename.endswith(".jar") and \ (filename.startswith("core") or \ filename.startswith(database.split("-")[0]) or \ not "binding" in filename): 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())