From aaf78afb3a9bf2442f944cf315d7d861dadb1a0f Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Tue, 26 Apr 2016 10:16:57 -0400 Subject: [PATCH] Add an cli option to specify the keyspace name and the ability to keep the data after the benchmark --- benchmarks/base.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/benchmarks/base.py b/benchmarks/base.py index 39912be4..b3b0fb04 100644 --- a/benchmarks/base.py +++ b/benchmarks/base.py @@ -71,10 +71,10 @@ def setup(options): session.execute(""" CREATE KEYSPACE %s WITH replication = { 'class': 'SimpleStrategy', 'replication_factor': '2' } - """ % KEYSPACE) + """ % options.keyspace) log.debug("Setting keyspace...") - session.set_keyspace(KEYSPACE) + session.set_keyspace(options.keyspace) log.debug("Creating table...") create_table_query = """ @@ -93,7 +93,8 @@ def setup(options): def teardown(options): cluster = Cluster(options.hosts, protocol_version=options.protocol_version) session = cluster.connect() - session.execute("DROP KEYSPACE " + KEYSPACE) + if not options.keep_data: + session.execute("DROP KEYSPACE " + options.keyspace) cluster.shutdown() @@ -108,7 +109,7 @@ def benchmark(thread_class): if options.protocol_version: kwargs['protocol_version'] = options.protocol_version cluster = Cluster(options.hosts, **kwargs) - session = cluster.connect(KEYSPACE) + session = cluster.connect(options.keyspace) log.debug("Sleeping for two seconds...") time.sleep(2.0) @@ -201,6 +202,10 @@ def parse_options(): help='Native protocol version to use') parser.add_option('-c', '--num-columns', type='int', dest='num_columns', default=2, help='Specify the number of columns for the schema') + parser.add_option('-k', '--keyspace', type='str', dest='keyspace', default=KEYSPACE, + help='Specify the keyspace name for the schema') + parser.add_option('--keep-data', action='store_true', dest='keep_data', default=False, + help='Keep the data after the benchmark') options, args = parser.parse_args()