Add read benchmark support

aboudreault · aboudreault · commit 86a793f0510d · 2016-04-26T12:42:24.000-04:00
diff --git a/benchmarks/base.py b/benchmarks/base.py
@@ -131,28 +131,30 @@ def benchmark(thread_class):
         log.debug("Sleeping for two seconds...")
         time.sleep(2.0)
 
-        # Generate the INSERT query
-        insert_query = "INSERT INTO {} (thekey".format(TABLE)
-        for i in range(options.num_columns):
-            insert_query += ", col{}".format(i)
-
-        insert_query += ") VALUES ('{key}'"
 
-        for i in range(options.num_columns):
-            insert_query += ", {}".format(COLUMN_VALUES[options.column_type])
-        insert_query += ")"
+        # Generate the query
+        if options.read:
+            query = "SELECT * FROM {}  WHERE thekey = '{{key}}'".format(TABLE)
+        else:
+            query = "INSERT INTO {} (thekey".format(TABLE)
+            for i in range(options.num_columns):
+                query += ", col{}".format(i)
 
-        values = None
+            query += ") VALUES ('{key}'"
+            for i in range(options.num_columns):
+                query += ", {}".format(COLUMN_VALUES[options.column_type])
+            query += ")"
 
+        values = None  # we don't use that anymore. Keeping it in case we go back to prepared statements.
         per_thread = options.num_ops // options.threads
         threads = []
 
-        log.debug("Beginning inserts...")
+        log.debug("Beginning {}...".format('reads' if options.read else 'inserts'))
         start = time.time()
         try:
             for i in range(options.threads):
                 thread = thread_class(
-                    i, session, insert_query, values, per_thread,
+                    i, session, query, values, per_thread,
                     cluster.protocol_version, options.profile)
                 thread.daemon = True
                 threads.append(thread)
@@ -225,6 +227,8 @@ def parse_options():
                       help='Keep the data after the benchmark')
     parser.add_option('--column-type', type='str', dest='column_type', default='text',
                       help='Specify the column type for the schema (supported: int, text, float, uuid, timestamp)')
+    parser.add_option('--read', action='store_true', dest='read', default=False,
+                      help='Read mode')
 
 
     options, args = parser.parse_args()
@@ -269,6 +273,9 @@ def start_profile(self):
         if self.profiler:
             self.profiler.enable()
 
+    def run_query(self, key, **kwargs):
+        return self.session.execute_async(self.query.format(key=key), **kwargs)
+
     def finish_profile(self):
         if self.profiler:
             self.profiler.disable()
diff --git a/benchmarks/callback_full_pipeline.py b/benchmarks/callback_full_pipeline.py
@@ -44,7 +44,7 @@ def insert_next(self, previous_result=sentinel):
         i = next(self.num_started)
         if  i <= self.num_queries:
             key = "{}-{}".format(self.thread_num, i)
-            future = self.session.execute_async(self.query.format(key=key), timeout=None)
+            future = self.run_query(key, timeout=None)
             future.add_callbacks(self.insert_next, self.insert_next)
 
     def run(self):
diff --git a/benchmarks/future_batches.py b/benchmarks/future_batches.py
@@ -36,7 +36,7 @@ def run(self):
                         break
 
             key = "{}-{}".format(self.thread_num, i)
-            future = self.session.execute_async(self.query.format(key=key))
+            future = self.run_query(key)
             futures.put_nowait(future)
 
         while True:
diff --git a/benchmarks/future_full_pipeline.py b/benchmarks/future_full_pipeline.py
@@ -32,7 +32,7 @@ def run(self):
                 old_future.result()
 
             key = "{}-{}".format(self.thread_num, i)
-            future = self.session.execute_async(self.query.format(key=key))
+            future = self.run_query(key)
             futures.put_nowait(future)
 
         while True:
diff --git a/benchmarks/future_full_throttle.py b/benchmarks/future_full_throttle.py
@@ -27,7 +27,7 @@ def run(self):
 
         for i in range(self.num_queries):
             key = "{}-{}".format(self.thread_num, i)
-            future = self.session.execute_async(self.query.format(key=key))
+            future = self.run_query(key)
             futures.append(future)
 
         for future in futures:
