|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# |
| 3 | +# Copyright (c) 2012 - 2020 YCSB contributors. All rights reserved. |
| 4 | +# |
| 5 | +# Licensed under the Apache License, Version 2.0 (the "License"); you |
| 6 | +# may not use this file except in compliance with the License. You |
| 7 | +# may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, software |
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
| 14 | +# implied. See the License for the specific language governing |
| 15 | +# permissions and limitations under the License. See accompanying |
| 16 | +# LICENSE file. |
| 17 | +# |
| 18 | + |
| 19 | +import errno |
| 20 | +import fnmatch |
| 21 | +import io |
| 22 | +import os |
| 23 | +import shlex |
| 24 | +import sys |
| 25 | +import subprocess |
| 26 | +import argparse |
| 27 | + |
| 28 | +BASE_URL = "https://github.com/brianfrankcooper/YCSB/tree/master/" |
| 29 | +COMMANDS = { |
| 30 | + "shell" : { |
| 31 | + "command" : "", |
| 32 | + "description" : "Interactive mode", |
| 33 | + "main" : "site.ycsb.CommandLine", |
| 34 | + }, |
| 35 | + "load" : { |
| 36 | + "command" : "-load", |
| 37 | + "description" : "Execute the load phase", |
| 38 | + "main" : "site.ycsb.Client", |
| 39 | + }, |
| 40 | + "run" : { |
| 41 | + "command" : "-t", |
| 42 | + "description" : "Execute the transaction phase", |
| 43 | + "main" : "site.ycsb.Client", |
| 44 | + }, |
| 45 | +} |
| 46 | + |
| 47 | +DATABASES = { |
| 48 | + "mongodb" : "site.ycsb.db.MongoDbClient", |
| 49 | + "mongodb-async": "site.ycsb.db.AsyncMongoDbClient", |
| 50 | + "postgrenosql" : "site.ycsb.postgrenosql.PostgreNoSQLDBClient", |
| 51 | +} |
| 52 | + |
| 53 | +OPTIONS = { |
| 54 | + "-P file" : "Specify workload file", |
| 55 | + "-p key=value" : "Override workload property", |
| 56 | + "-s" : "Print status to stderr", |
| 57 | + "-target n" : "Target ops/sec (default: unthrottled)", |
| 58 | + "-threads n" : "Number of client threads (default: 1)", |
| 59 | + "-cp path" : "Additional Java classpath entries", |
| 60 | + "-jvm-args args" : "Additional arguments to the JVM", |
| 61 | +} |
| 62 | + |
| 63 | +def usage(): |
| 64 | + output = io.StringIO() |
| 65 | + print("%s command database [options]" % sys.argv[0], file=output) |
| 66 | + |
| 67 | + print("\nCommands:", file=output) |
| 68 | + for command in sorted(COMMANDS.keys()): |
| 69 | + print(" %s %s" % (command.ljust(14), |
| 70 | + COMMANDS[command]["description"]), file=output) |
| 71 | + |
| 72 | + print("\nDatabases:", file=output) |
| 73 | + for db in sorted(DATABASES.keys()): |
| 74 | + print(" %s %s" % (db.ljust(14), BASE_URL + |
| 75 | + db.split("-")[0]), file=output) |
| 76 | + |
| 77 | + print("\nOptions:", file=output) |
| 78 | + for option in sorted(OPTIONS.keys()): |
| 79 | + print(" %s %s" % (option.ljust(14), OPTIONS[option]), file=output) |
| 80 | + |
| 81 | + print("""\nWorkload Files: |
| 82 | + There are various predefined workloads under workloads/ directory. |
| 83 | + See https://github.com/brianfrankcooper/YCSB/wiki/Core-Properties |
| 84 | + for the list of workload properties.""", file=output) |
| 85 | + |
| 86 | + return output.getvalue() |
| 87 | + |
| 88 | +def debug(message): |
| 89 | + print("[DEBUG] ", message, file=sys.stderr) |
| 90 | + |
| 91 | +def warn(message): |
| 92 | + print("[WARN] ", message, file=sys.stderr) |
| 93 | + |
| 94 | +def error(message): |
| 95 | + print("[ERROR] ", message, file=sys.stderr) |
| 96 | + |
| 97 | +def find_jars(dir, glob='*.jar'): |
| 98 | + jars = [] |
| 99 | + for (dirpath, dirnames, filenames) in os.walk(dir): |
| 100 | + for filename in fnmatch.filter(filenames, glob): |
| 101 | + jars.append(os.path.join(dirpath, filename)) |
| 102 | + return jars |
| 103 | + |
| 104 | +def get_ycsb_home(): |
| 105 | + dir = os.path.abspath(os.path.dirname(sys.argv[0])) |
| 106 | + while "LICENSE.txt" not in os.listdir(dir): |
| 107 | + dir = os.path.join(dir, os.path.pardir) |
| 108 | + return os.path.abspath(dir) |
| 109 | + |
| 110 | +def is_distribution(): |
| 111 | + # If there's a top level pom, we're a source checkout. otherwise a dist artifact |
| 112 | + return "pom.xml" not in os.listdir(get_ycsb_home()) |
| 113 | + |
| 114 | +# Run the maven dependency plugin to get the local jar paths. |
| 115 | +# presumes maven can run, so should only be run on source checkouts |
| 116 | +# will invoke the 'package' goal for the given binding in order to resolve intra-project deps |
| 117 | +# presumes maven properly handles system-specific path separators |
| 118 | +# Given module is full module name eg. 'core' or 'couchbase-binding' |
| 119 | +def get_classpath_from_maven(module): |
| 120 | + output_file = tempfile.NamedTemporaryFile(delete=True) |
| 121 | + cmd = ["mvn", "-B", "-pl", "site.ycsb:" + module, |
| 122 | + "-am", "package", "-DskipTests", |
| 123 | + "dependency:list", |
| 124 | + "-DoutputAbsoluteArtifactFilename", |
| 125 | + "-DappendOutput=false", |
| 126 | + "-DoutputFile=" + output_file.name |
| 127 | + ] |
| 128 | + debug("Running '" + " ".join(cmd) + "'") |
| 129 | + subprocess.check_call(cmd) |
| 130 | + |
| 131 | + # the output file will now contain all of the dependencies in the format: |
| 132 | + # group:artifact:type:[classifier:]version:scope:path[ -- module info] |
| 133 | + classpath_items = [] |
| 134 | + with open(output_file.name) as f: |
| 135 | + for l in f.readlines(): |
| 136 | + l = l.strip() |
| 137 | + m = re.match('(?P<group>[^:]+):(?P<artifact>[^:]+):(?P<type>[^:]+)(?::(?P<classifier>[^:]+))?:(?P<version>[^:]+):(?P<scope>[^:]+):(?P<path>.*?)( -- module .*)?$', l) |
| 138 | + if not m: |
| 139 | + continue |
| 140 | + if m.groupdict()["scope"] == "test": |
| 141 | + continue |
| 142 | + classpath_items.append(m.groupdict()["path"]) |
| 143 | + return ":".join(classpath_items) |
| 144 | + |
| 145 | + |
| 146 | +def main(): |
| 147 | + p = argparse.ArgumentParser( |
| 148 | + usage=usage(), |
| 149 | + formatter_class=argparse.RawDescriptionHelpFormatter) |
| 150 | + p.add_argument('-cp', dest='classpath', help="""Additional classpath |
| 151 | + entries, e.g. '-cp /tmp/hbase-1.0.1.1/conf'. Will be |
| 152 | + prepended to the YCSB classpath.""") |
| 153 | + p.add_argument("-jvm-args", default=[], type=shlex.split, |
| 154 | + help="""Additional arguments to pass to 'java', e.g. |
| 155 | + '-Xmx4g'""") |
| 156 | + p.add_argument("command", choices=sorted(COMMANDS), |
| 157 | + help="""Command to run.""") |
| 158 | + p.add_argument("database", choices=sorted(DATABASES), |
| 159 | + help="""Database to test.""") |
| 160 | + args, remaining = p.parse_known_args() |
| 161 | + ycsb_home = get_ycsb_home() |
| 162 | + |
| 163 | + # Use JAVA_HOME to find java binary if set, otherwise just use PATH. |
| 164 | + java = "java" |
| 165 | + java_home = os.getenv("JAVA_HOME") |
| 166 | + if java_home: |
| 167 | + java = os.path.join(java_home, "bin", "java") |
| 168 | + db_classname = DATABASES[args.database] |
| 169 | + command = COMMANDS[args.command]["command"] |
| 170 | + main_classname = COMMANDS[args.command]["main"] |
| 171 | + |
| 172 | + # Classpath set up |
| 173 | + binding = args.database.split("-")[0] |
| 174 | + |
| 175 | + if binding == "cassandra2": |
| 176 | + warn("The 'cassandra2-cql' client has been deprecated. It has been " |
| 177 | + "renamed to simply 'cassandra-cql'. This alias will be removed" |
| 178 | + " in the next YCSB release.") |
| 179 | + binding = "cassandra" |
| 180 | + |
| 181 | + if binding == "couchbase": |
| 182 | + warn("The 'couchbase' client has been deprecated. If you are using " |
| 183 | + "Couchbase 4.0+ try using the 'couchbase2' client instead.") |
| 184 | + |
| 185 | + if binding == "hbase14": |
| 186 | + warn("The 'hbase14' client has been deprecated. HBase 1.y users should " |
| 187 | + "rely on the 'hbase1' client instead.") |
| 188 | + binding = "hbase1" |
| 189 | + |
| 190 | + if binding == "arangodb3": |
| 191 | + warn("The 'arangodb3' client has been deprecated. The binding 'arangodb' " |
| 192 | + "now covers every ArangoDB version. This alias will be removed " |
| 193 | + "in the next YCSB release.") |
| 194 | + binding = "arangodb" |
| 195 | + |
| 196 | + if is_distribution(): |
| 197 | + db_dir = os.path.join(ycsb_home, binding + "-binding") |
| 198 | + # include top-level conf for when we're a binding-specific artifact. |
| 199 | + # If we add top-level conf to the general artifact, starting here |
| 200 | + # will allow binding-specific conf to override (because it's prepended) |
| 201 | + cp = [os.path.join(ycsb_home, "conf")] |
| 202 | + cp.extend(find_jars(os.path.join(ycsb_home, "lib"))) |
| 203 | + cp.extend(find_jars(os.path.join(db_dir, "lib"))) |
| 204 | + else: |
| 205 | + warn("Running against a source checkout. In order to get our runtime " |
| 206 | + "dependencies we'll have to invoke Maven. Depending on the state " |
| 207 | + "of your system, this may take ~30-45 seconds") |
| 208 | + db_location = "core" if (binding == "basic" or binding == "basicts") else binding |
| 209 | + project = "core" if (binding == "basic" or binding == "basicts") else binding + "-binding" |
| 210 | + db_dir = os.path.join(ycsb_home, db_location) |
| 211 | + # goes first so we can rely on side-effect of package |
| 212 | + maven_says = get_classpath_from_maven(project) |
| 213 | + # TODO when we have a version property, skip the glob |
| 214 | + cp = find_jars(os.path.join(db_dir, "target"), |
| 215 | + project + "*.jar") |
| 216 | + # alredy in jar:jar:jar form |
| 217 | + cp.append(maven_says) |
| 218 | + cp.insert(0, os.path.join(db_dir, "conf")) |
| 219 | + classpath = os.pathsep.join(cp) |
| 220 | + if args.classpath: |
| 221 | + classpath = os.pathsep.join([args.classpath, classpath]) |
| 222 | + |
| 223 | + ycsb_command = ([java] + args.jvm_args + |
| 224 | + ["-cp", classpath, |
| 225 | + main_classname, "-db", db_classname] + remaining) |
| 226 | + if command: |
| 227 | + ycsb_command.append(command) |
| 228 | + print(" ".join(ycsb_command), file=sys.stderr) |
| 229 | + try: |
| 230 | + return subprocess.call(ycsb_command) |
| 231 | + except OSError as e: |
| 232 | + if e.errno == errno.ENOENT: |
| 233 | + error('Command failed. Is java installed and on your PATH?') |
| 234 | + return 1 |
| 235 | + else: |
| 236 | + raise |
| 237 | + |
| 238 | +if __name__ == '__main__': |
| 239 | + sys.exit(main()) |
0 commit comments