Refactor RelationalDbSpec registration to use CLOUD and ENGINE.

PiperOrigin-RevId: 904720629

Author: bvliu

perfkitbenchmarker/configs/benchmark_config_spec.py

@@ -811,14 +811,18 @@ def Decode(self, value, component_full_name, flag_values):
     # LoadProvider is required for resources to be registered.
     _LoadProvider(relational_db_config, flag_values)
 
+    cloud = relational_db_config.get('cloud')
+    if not cloud or flag_values['cloud'].present:
+      cloud = flag_values.cloud
+
     if 'engine' in relational_db_config:
       if flag_values['db_engine'].present:
         db_spec_class = relational_db_spec.GetRelationalDbSpecClass(
-            flag_values['db_engine'].value
+            cloud, flag_values['db_engine'].value
         )
       else:
         db_spec_class = relational_db_spec.GetRelationalDbSpecClass(
-            relational_db_config['engine']
+            cloud, relational_db_config['engine']
         )
     else:
       raise errors.Config.InvalidValue(

perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py

@@ -52,7 +52,8 @@
 class AwsAuroraDsqlSpec(relational_db_spec.RelationalDbSpec):
   """Configurable options for AWS Aurora DSQL."""
 
-  SERVICE_TYPE = 'aurora-dsql'
+  CLOUD = 'AWS'
+  ENGINE = [sql_engine_utils.AURORA_DSQL_POSTGRES]
 
   @classmethod
   def _GetOptionDecoderConstructions(cls):

perfkitbenchmarker/providers/gcp/gcp_spanner.py

@@ -206,7 +206,11 @@ class SpannerAsyncOperationInProgressError(Exception):
 class SpannerSpec(relational_db_spec.RelationalDbSpec):
   """Configurable options of a Spanner instance."""
 
-  SERVICE_TYPE = 'spanner'
+  CLOUD = 'GCP'
+  ENGINE = [
+      sql_engine_utils.SPANNER_GOOGLESQL,
+      sql_engine_utils.SPANNER_POSTGRES,
+  ]
 
   spanner_instance_id: str
   spanner_database_id: str

perfkitbenchmarker/relational_db_spec.py

@@ -33,16 +33,9 @@
 _NONE_OK = {'default': None, 'none_ok': True}
 
 
-def GetRelationalDbSpecClass(engine):
-  """Get the RelationalDbSpec class corresponding to 'engine'."""
-  if engine in [
-      sql_engine_utils.SPANNER_GOOGLESQL,
-      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
+def GetRelationalDbSpecClass(cloud, engine):
+  """Get the RelationalDbSpec class corresponding to 'cloud' and 'engine'."""
+  return spec.GetSpecClass(RelationalDbSpec, CLOUD=cloud, ENGINE=engine)
 
 
 class RelationalDbSpec(freeze_restore_spec.FreezeRestoreSpec):
@@ -63,7 +56,7 @@ class RelationalDbSpec(freeze_restore_spec.FreezeRestoreSpec):
   """
 
   SPEC_TYPE = 'RelationalDbSpec'
-  SPEC_ATTRS = ['SERVICE_TYPE']
+  SPEC_ATTRS = ['CLOUD', 'ENGINE']
 
   cloud: str
   engine: str

tests/relational_db_spec_test.py

@@ -15,7 +15,9 @@
 
 import unittest
 from absl import flags
+from absl.testing import parameterized
 from perfkitbenchmarker import errors
+from perfkitbenchmarker import providers
 from perfkitbenchmarker import relational_db
 from perfkitbenchmarker import relational_db_spec
 from perfkitbenchmarker.providers.gcp import gce_virtual_machine
@@ -318,6 +320,22 @@ def testReverseOrderedDbReplicaZonesFlag(self):
         result.vm_groups['servers_replicas'].vm_spec.zone, 'us-central1-d'
     )
 
+  # pyformat: disable
+  @parameterized.named_parameters([
+      ('SpannerGoogleSql', 'GCP', 'spanner-googlesql', 'SpannerSpec'),
+      ('SpannerPostgres', 'GCP', 'spanner-postgres', 'SpannerSpec'),
+      ('AuroraDsql', 'AWS', 'aurora-dsql-postgres', 'AwsAuroraDsqlSpec'),
+      ('RdsMysql', 'AWS', 'mysql', 'RelationalDbSpec'),
+      ('CloudSqlPostgres', 'GCP', 'postgres', 'RelationalDbSpec'),
+      ('AzureFlexibleServer', 'Azure', 'flexible-server-postgres', 'RelationalDbSpec'),
+      ('AzureSqlManagedInstance', 'Azure', 'sqlserver', 'RelationalDbSpec'),
+  ])
+  # pyformat: enable
+  def testGetRelationalDbSpecClass(self, cloud, engine, expected_class_name):
+    providers.LoadProvider(cloud, True)
+    actual_class = relational_db_spec.GetRelationalDbSpecClass(cloud, engine)
+    self.assertEqual(actual_class.__name__, expected_class_name)
+
 
 if __name__ == '__main__':
   unittest.main()