code style improvement

PiperOrigin-RevId: 835279942

Author: ScottLinnn

perfkitbenchmarker/linux_benchmarks/benchbase_benchmark.py

@@ -26,28 +26,22 @@
 from perfkitbenchmarker import configs
 from perfkitbenchmarker import sample
 from perfkitbenchmarker.linux_packages import benchbase
+# Needed in order to register spec:
+from perfkitbenchmarker.providers.aws import aws_aurora_dsql_db  # pylint: disable=unused-import
 
 BENCHMARK_NAME: str = 'benchbase'
 BENCHMARK_CONFIG: str = """
 benchbase:
   description: Runs Benchbase benchmark.
   relational_db:
-    cloud: GCP
-    engine: spanner-postgres
+    cloud: AWS
+    engine: aurora-dsql-postgres
     db_spec:
       GCP:
         machine_type: db-n1-standard-16
         zone: us-central1-f
       AWS:
-        machine_type: db.m4.4xlarge
         zone: us-east-1a
-    db_disk_spec:
-      GCP:
-        disk_size: 2048
-        disk_type: pd-ssd
-      AWS: # not actually used by it's a required spec so fill random values
-        disk_size: 6144
-        disk_type: gp2
     vm_groups:
       clients:
         vm_spec:

perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py

@@ -19,8 +19,10 @@
 from absl import flags
 from absl import logging
 from perfkitbenchmarker import errors
+from perfkitbenchmarker import relational_db_spec
 from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker import vm_util
+from perfkitbenchmarker.configs import spec
 from perfkitbenchmarker.providers.aws import aws_relational_db
 from perfkitbenchmarker.providers.aws import util
 
@@ -40,6 +42,23 @@
     sql_engine_utils.AURORA_DSQL_POSTGRES,
 ]
 
+_NONE_OK = {'default': None, 'none_ok': True}
+
+
+class AwsAuroraDsqlSpec(relational_db_spec.RelationalDbSpec):
+  """Configurable options for AWS Aurora DSQL."""
+
+  SERVICE_TYPE = 'aurora-dsql'
+  db_disk_spec = None
+
+  @classmethod
+  def _GetOptionDecoderConstructions(cls):
+    result = super()._GetOptionDecoderConstructions()
+    result.update({
+        'db_disk_spec': (spec.PerCloudConfigDecoder, _NONE_OK),
+    })
+    return result
+
 
 class AwsAuroraDsqlRelationalDb(aws_relational_db.BaseAwsRelationalDb):
   """Implements the aurora DSQL database for AWS."""
@@ -48,8 +67,8 @@ class AwsAuroraDsqlRelationalDb(aws_relational_db.BaseAwsRelationalDb):
   IS_MANAGED = True
   ENGINE = _AURORA_DSQL_ENGINES
 
-  def __init__(self, relational_db_spec):
-    super().__init__(relational_db_spec)
+  def __init__(self, dsql_spec: AwsAuroraDsqlSpec):
+    super().__init__(dsql_spec)
     self.cluster_id = None
 
   # DSQL has different format for tags:
@@ -162,6 +181,8 @@ def GetDefaultEngineVersion(engine) -> str:
     return _MAP_ENGINE_TO_DEFAULT_VERSION[engine]
 
   def GetResourceMetadata(self) -> dict[str, Any]:
-    metadata = super().GetResourceMetadata()
-    metadata['dsql_cluster_id'] = self.cluster_id
+    """Returns useful metadata about the instance."""
+    metadata = {
+        'dsql_cluster_id': self.cluster_id,
+    }
     return metadata

perfkitbenchmarker/relational_db_spec.py

@@ -39,6 +39,8 @@ def GetRelationalDbSpecClass(engine):
       sql_engine_utils.SPANNER_POSTGRES,
   ]:
     return spec.GetSpecClass(RelationalDbSpec, SERVICE_TYPE='spanner')
+  if engine == sql_engine_utils.AURORA_DSQL_POSTGRES:
+    return spec.GetSpecClass(RelationalDbSpec, SERVICE_TYPE='aurora-dsql')
   return RelationalDbSpec
 
 

tests/providers/aws/aws_aurora_dsql_db_test.py

@@ -67,7 +67,6 @@ def CreateMockSpec(self):
         'database_name': 'fakedbname',
         'database_password': 'fakepassword',
         'db_spec': mock.MagicMock(),
-        'db_disk_spec': mock.MagicMock(disk_type='gp2'),
         'enable_freeze_restore': False,
         'create_on_restore_error': False,
         'delete_on_freeze_error': False,