Merge pull request brianfrankcooper#276 from cmccoy/cleanup-ycsb-script

Clean up bin/ycsb.

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())