Support protocol for pgbench

PiperOrigin-RevId: 696254311
GoogleCloudPlatform · Nov 13, 2024 · 249f684 · 249f684
1 parent 454927d
commit 249f684
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/perfkitbenchmarker/linux_benchmarks/pgbench_benchmark.py b/perfkitbenchmarker/linux_benchmarks/pgbench_benchmark.py
@@ -40,12 +40,25 @@
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker.linux_packages import pgbench
 
+# PGProtocal
+SIMPLE = 'simple'
+EXTENDED = 'extended'
+PREPARED = 'prepared'
+
 flags.DEFINE_integer(
     'pgbench_scale_factor',
     1,
     'scale factor used to fill the database',
     lower_bound=1,
 )
+
+flags.DEFINE_enum(
+    'pgbench_protocol',
+    None,
+    [SIMPLE, EXTENDED, PREPARED],
+    'Protocol to use for pgbench',
+)
+
 flags.DEFINE_integer(
     'pgbench_seconds_per_test',
     10,
@@ -167,6 +180,7 @@ def UpdateBenchmarkSpecWithRunStageFlags(benchmark_spec):
   benchmark_spec.seconds_to_pause = FLAGS.pgbench_seconds_to_pause_before_steps
   benchmark_spec.client_counts = FLAGS.pgbench_client_counts
   benchmark_spec.job_counts = FLAGS.pgbench_job_counts
+  benchmark_spec.pgbench_protocol = FLAGS.pgbench_protocol
 
 
 def GetDbSize(relational_db, db_name):
@@ -358,6 +372,7 @@ def Run(benchmark_spec):
       benchmark_spec.job_counts,
       benchmark_spec.seconds_to_pause,
       benchmark_spec.seconds_per_test,
+      benchmark_spec.pgbench_protocol,
       common_metadata,
   )
   return []

diff --git a/perfkitbenchmarker/linux_packages/pgbench.py b/perfkitbenchmarker/linux_packages/pgbench.py
@@ -95,6 +95,7 @@ def RunPgBench(
     job_counts,
     seconds_to_pause,
     seconds_per_test,
+    protocol,
     metadata,
     file=None,
     path=None,
@@ -110,6 +111,7 @@ def RunPgBench(
     job_counts: Number of job
     seconds_to_pause: Seconds to pause between test
     seconds_per_test: Seconds per test
+    protocol: Protocol to use for the benchmark
     metadata: Metadata of the benchmark
     file: Filename of the benchmark
     path: File path of the benchmar.
@@ -136,6 +138,8 @@ def RunPgBench(
     metadata['pgbench_file'] = file
 
   extended_protocol = ''
+  if protocol:
+    extended_protocol = ' --protocol=' + protocol
   if relational_db.engine_type == sql_engine_utils.SPANNER_POSTGRES:
     extended_protocol = ' --protocol=extended'