Support running DSQL with benchbase

PiperOrigin-RevId: 838952584

Author: ScottLinnn

perfkitbenchmarker/linux_benchmarks/benchbase_benchmark.py

@@ -24,6 +24,7 @@
 from absl import flags
 from perfkitbenchmarker import benchmark_spec as bm_spec
 from perfkitbenchmarker import configs
+from perfkitbenchmarker import errors
 from perfkitbenchmarker import sample
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker.linux_packages import benchbase
@@ -75,7 +76,8 @@ def GetConfig(user_config: Dict[str, Any]) -> Dict[str, Any]:
   Returns:
     loaded benchmark configuration
   """
-  return configs.LoadConfig(BENCHMARK_CONFIG, user_config, BENCHMARK_NAME)
+  config = configs.LoadConfig(BENCHMARK_CONFIG, user_config, BENCHMARK_NAME)
+  return config
 
 
 # TODO(shuninglin): need to implement auth logic(automatic password gen)
@@ -96,6 +98,16 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec) -> None:
 
   # Create the configuration file on the client VM
   benchbase.CreateConfigFile(client_vm)
+  if FLAGS.db_engine == sql_engine_utils.AURORA_DSQL_POSTGRES:
+    dsql: aws_aurora_dsql_db.AwsAuroraDsqlRelationalDb = (
+        benchmark_spec.relational_db
+    )
+    # Ideally we want to use endpoint from get-cluster command but it's not
+    # returning endpoint as documented. That said this hard-coded endpoint
+    # construction is also documented so should be reliable.
+    # https://docs.aws.amazon.com/aurora-dsql/latest/userguide/SECTION_authentication-token.html#authentication-token-cli
+    endpoint: str = f'{dsql.cluster_id}.dsql.{dsql.region}.on.aws'
+    benchbase.OverrideEndpoint(client_vm, endpoint)
   profile: str = (
       'postgres'
       if FLAGS.db_engine == sql_engine_utils.SPANNER_POSTGRES
@@ -107,6 +119,7 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec) -> None:
       f' {benchbase.CONFIG_FILE_NAME} --create=true --load=true'
       ' --execute=false"'
   )
+
   client_vm.RemoteCommand(load_command)
 
 

perfkitbenchmarker/linux_packages/benchbase.py

@@ -179,10 +179,10 @@ def CreateConfigFile(vm: virtual_machine.BaseVirtualMachine) -> None:
 
   if db_engine == sql_engine_utils.AURORA_DSQL_POSTGRES:
     context['db_type'] = 'AURORADSQL'
-    # For DSQL we use automatic username and password generation so comment out
-    # the username and password elements.
-    context['username_element'] = '<!--<username>admin</username>-->'
-    context['password_element'] = '<!--<password>password</password>-->'
+    # Following guide here to use automatic username and password generation:
+    # https://github.com/amazon-contributing/aurora-dsql-benchbase-benchmarking/wiki#loading-data-and-running-tpc-c-against-an-aurora-dsql-cluster
+    context['username_element'] = '<username>admin</username>'
+    context['password_element'] = '<password></password>'
   else:  # spanner by default
     context['db_type'] = 'POSTGRES'
     context['username_element'] = '<username>admin</username>'
@@ -209,6 +209,18 @@ def CreateConfigFile(vm: virtual_machine.BaseVirtualMachine) -> None:
     return
 
 
+def OverrideEndpoint(
+    vm: virtual_machine.BaseVirtualMachine, endpoint: str
+) -> None:
+  """Overrides the endpoint in the Benchbase XML configuration file on the client VM.
+
+  Args:
+    vm: The client virtual machine to create the file on.
+    endpoint: The endpoint of the database.
+  """
+  vm.RemoteCommand(f"sed -i 's/localhost/{endpoint}/g' {CONFIG_FILE_PATH}")
+
+
 def Uninstall(vm: virtual_machine.BaseVirtualMachine) -> None:
   """Uninstalls the BenchBase package on the VM by removing the directory.
 

perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py

@@ -69,6 +69,7 @@ class AwsAuroraDsqlRelationalDb(aws_relational_db.BaseAwsRelationalDb):
   def __init__(self, dsql_spec: AwsAuroraDsqlSpec):
     super().__init__(dsql_spec)
     self.cluster_id = None
+    self.assigned_name = f'pkb-{FLAGS.run_uri}'
 
   # DSQL has different format for tags:
   # https://docs.aws.amazon.com/cli/v1/reference/rds/create-db-cluster.html
@@ -77,7 +78,7 @@ def _MakeDsqlTags(self) -> list[str]:
     tags: dict[str, Any] = util.MakeDefaultTags()
     formatted_tags_list: list[str] = [
         '%s=%s' % (k, v) for k, v in sorted(tags.items())
-    ]
+    ] + ['Name=%s' % self.assigned_name]
     formatted_tags_str = ','.join(formatted_tags_list)
     return [formatted_tags_str]
 

tests/linux_packages/benchbase_test.py

@@ -100,10 +100,10 @@ def test_create_config_file_aurora_dsql(self, _):
     context = kwargs['context']
     self.assertEqual(context['db_type'], 'AURORADSQL')
     self.assertEqual(
-        context['username_element'], '<!--<username>admin</username>-->'
+        context['username_element'], '<username>admin</username>'
     )
     self.assertEqual(
-        context['password_element'], '<!--<password>password</password>-->'
+        context['password_element'], '<password></password>'
     )
     self.assertEqual(context['rate_element'], '<rate>unlimited</rate>')
     self.assertIn(

tests/providers/aws/aws_aurora_dsql_db_test.py

@@ -113,7 +113,7 @@ def testMakeDsqlTags(self):
     ):
       db = self.CreateDbFromSpec()
       tags = db._MakeDsqlTags()
-      self.assertEqual(tags, ['tag1=value1,tag2=value2'])
+      self.assertEqual(tags, ['tag1=value1,tag2=value2,Name=pkb-123'])
 
   def testDescribeCluster(self):
     """Tests that the describe cluster command is correct."""