diffbase, omitted

perfkitbenchmarker/providers/aws/aws_aurora_db.py

@@ -106,21 +106,31 @@ def _Create(self):
             'create-db-cluster',
             '--db-cluster-identifier=%s' % self.cluster_id,
             '--engine=%s' % self.spec.engine,
-            '--engine-version=%s' % self.spec.engine_version,
-            '--master-username=%s' % self.spec.database_username,
-            '--master-user-password=%s' % self.spec.database_password,
             '--region=%s' % self.region,
-            '--db-subnet-group-name=%s' % self.db_subnet_group_name,
-            '--vpc-security-group-ids=%s' % self.security_group_id,
-            '--availability-zones=%s' % self.spec.zones[0],
-            '--storage-type=%s' % self.storage_type,
             '--backup-retention-period=1',  # backups cannot be disabled
-            '--tags',
         ]
-        + util.MakeFormattedDefaultTags()
     )
+    if self.spec.aws_aurora_express_configuration:
+      cmd.append('--with-express-configuration')
+    else:
+      cmd.extend([
+          '--engine-version=%s' % self.spec.engine_version,
+          '--master-username=%s' % self.spec.database_username,
+          '--master-user-password=%s' % self.spec.database_password,
+          '--db-subnet-group-name=%s' % self.db_subnet_group_name,
+          '--vpc-security-group-ids=%s' % self.security_group_id,
+          '--availability-zones=%s' % self.spec.zones[0],
+          '--storage-type=%s' % self.storage_type,
+      ])
+    cmd.extend(['--tags'] + util.MakeFormattedDefaultTags())
+
+    stdout, _, _ = vm_util.IssueCommand(cmd)
+    json_output = json.loads(stdout)
+    for member in json_output.get('DBCluster', {}).get('DBClusterMembers', []):
+      self.all_instance_ids.append(member['DBInstanceIdentifier'])
 
-    vm_util.IssueCommand(cmd)
+    if self.spec.aws_aurora_express_configuration:
+      return
 
     for zone in self.zones:
       # The first instance is assumed to be writer -
@@ -161,6 +171,25 @@ def _PostCreate(self):
     super()._PostCreate()
     self._SetPrimaryAndSecondaryZones()
 
+  def _WaitUntilReady(self) -> None:
+    """Waits & retries until the resource is ready."""
+    # Poll with a shorter interval for express configuration.
+    if self.spec.aws_aurora_express_configuration:
+
+      @vm_util.Retry(
+          poll_interval=1,
+          fuzz=0,
+          timeout=self.READY_TIMEOUT,
+          retryable_exceptions=(errors.Resource.RetryableCreationError,),
+      )
+      def _InnerWaitUntilReady() -> None:
+        if not self._IsReady(poll_interval=1):
+          raise errors.Resource.RetryableCreationError('Not yet ready')
+
+      _InnerWaitUntilReady()
+    else:
+      super()._WaitUntilReady()
+
   def _UpdateClusterClass(self, instance_class: str) -> None:
     """Updates DBInstanceClass for all instances in the cluster."""
     for db_instance in self.all_instance_ids:
@@ -327,6 +356,8 @@ def GetDefaultEngineVersion(engine):
   def GetResourceMetadata(self) -> dict[str, Any]:
     metadata = super().GetResourceMetadata()
     metadata['aurora_storage_type'] = self.storage_type
+    if self.spec.aws_aurora_express_configuration:
+      metadata['aws_aurora_express_configuration'] = True
     return metadata
 
   def _GetMetricsToCollect(self) -> list[relational_db.MetricSpec]:

perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py

@@ -203,7 +203,9 @@ def _DescribeCluster(self) -> dict[str, Any] | None:
     json_output: dict[str, Any] = json.loads(stdout)
     return json_output
 
-  def _IsReady(self, timeout=aws_relational_db.IS_READY_TIMEOUT) -> bool:
+  def _IsReady(
+      self, timeout=aws_relational_db.IS_READY_TIMEOUT, poll_interval=5
+  ) -> bool:
     """Returns true if the cluster is ready."""
     if self.use_backup:
       if not self.restore_job_id:
@@ -270,7 +272,7 @@ def _GetHostname(self) -> str:
     """Returns endpoint of DSQL cluster."""
     return f'{self.cluster_id}.dsql.{self.region}.on.aws'
 
-# TODO(shuninglin): Extend PostgresCliQueryTools for DSQL.
+  # TODO(shuninglin): Extend PostgresCliQueryTools for DSQL.
   def RunSqlQuery(self, sql_query: str) -> None:
     """Runs a SQL query on the database."""
     # Local import to avoid dependency not found issue.

perfkitbenchmarker/providers/aws/aws_relational_db.py

@@ -174,14 +174,15 @@ def __init__(self, relational_db_spec):
     self.db_subnet_group_name: str = None
     self.security_group_id: str = None
 
-  def _IsReady(self, timeout=IS_READY_TIMEOUT):
+  def _IsReady(self, timeout=IS_READY_TIMEOUT, poll_interval=5):
     """Return true if the underlying resource is ready.
 
-    This method will query all of the instance every 5 seconds until
+    This method will query all of the instance every poll_interval seconds until
     its instance state is 'available', or until a timeout occurs.
 
     Args:
       timeout: timeout in seconds
+      poll_interval: time in seconds between calls
 
     Returns:
       True if the resource was ready in time, False if the wait timed out
@@ -191,7 +192,7 @@ def _IsReady(self, timeout=IS_READY_TIMEOUT):
       return False
 
     for instance_id in self.all_instance_ids:
-      if not self._IsInstanceReady(instance_id, timeout):
+      if not self._IsInstanceReady(instance_id, timeout, poll_interval):
         return False
 
     return True
@@ -424,15 +425,18 @@ def _DescribeInstance(self, instance_id):
     json_output = json.loads(stdout)
     return json_output
 
-  def _IsInstanceReady(self, instance_id, timeout=IS_READY_TIMEOUT):
+  def _IsInstanceReady(
+      self, instance_id, timeout=IS_READY_TIMEOUT, poll_interval=5
+  ):
     """Return true if the instance is ready.
 
-    This method will query the instance every 5 seconds until
+    This method will query the instance every poll_interval seconds until
     its instance state is 'available', or until a timeout occurs.
 
     Args:
       instance_id: string of the instance to check is ready
       timeout: timeout in seconds
+      poll_interval: time in seconds between calls
 
     Returns:
       True if the resource was ready in time, False if the wait timed out
@@ -476,7 +480,7 @@ def _IsInstanceReady(self, instance_id, timeout=IS_READY_TIMEOUT):
               'Error attempting to read stdout. Creation failure.'
           )
           return False
-      time.sleep(5)
+      time.sleep(poll_interval)
 
     return True
 

perfkitbenchmarker/providers/aws/flags.py

@@ -278,6 +278,11 @@
     'The time to sleep before collecting Aurora metrics. By default this is a'
     ' long time in order to collect accurate VolumeBytesUsed metrics.',
 )
+AWS_AURORA_EXPRESS_CONFIGURATION = flags.DEFINE_boolean(
+    'aws_aurora_express_configuration',
+    False,
+    'Whether to use express configuration for Aurora cluster creation.',
+)
 AWS_EC2_INSTANCE_PROFILE = flags.DEFINE_string(
     'aws_ec2_instance_profile',
     None,

perfkitbenchmarker/resources/aws_relational_db_spec.py

@@ -35,6 +35,7 @@ class AwsRelationalDbSpec(relational_db_spec.RelationalDbSpec):
   CLOUD = 'AWS'
 
   aws_rds_dedicated_log_volume: bool
+  aws_aurora_express_configuration: bool
 
   def __init__(self, component_full_name, flag_values=None, **kwargs):
     super().__init__(component_full_name, flag_values=flag_values, **kwargs)
@@ -65,6 +66,10 @@ def _GetOptionDecoderConstructions(cls):
             option_decoders.BooleanDecoder,
             {'default': False},
         ),
+        'aws_aurora_express_configuration': (
+            option_decoders.BooleanDecoder,
+            {'default': False},
+        ),
     })
     return result
 
@@ -76,3 +81,7 @@ def _ApplyFlags(cls, config_values, flag_values):
       config_values['aws_rds_dedicated_log_volume'] = (
           flag_values.aws_rds_dedicated_log_volume
       )
+    if flag_values['aws_aurora_express_configuration'].present:
+      config_values['aws_aurora_express_configuration'] = (
+          flag_values.aws_aurora_express_configuration
+      )

tests/providers/aws/aws_relational_db_test.py

@@ -120,7 +120,7 @@ def setUp(self):
     FLAGS['use_managed_db'].parse(True)
 
   @contextlib.contextmanager
-  def _PatchCriticalObjects(self, stdout='', stderr='', return_code=0):
+  def _PatchCriticalObjects(self, stdout='{}', stderr='', return_code=0):
     """A context manager that patches a few critical objects with mocks."""
     retval = (stdout, stderr, return_code)
     with mock.patch(
@@ -268,6 +268,7 @@ def CreateAuroraMockSpec(self, additional_spec_items=None):
         'delete_on_freeze_error': False,
         'load_machine_type': None,
         'aws_rds_dedicated_log_volume': False,
+        'aws_aurora_express_configuration': False,
     }
     if additional_spec_items:
       spec_dict.update(additional_spec_items)
@@ -299,13 +300,14 @@ def testCreateAurora(self):
             (
                 'aws --output json rds create-db-cluster'
                 ' --db-cluster-identifier=pkb-db-cluster-123'
-                ' --engine=aurora-postgresql --engine-version=9.6.2'
+                ' --engine=aurora-postgresql --region=us-east-1'
+                ' --backup-retention-period=1 --engine-version=9.6.2'
                 ' --master-username=fakeusername'
-                ' --master-user-password=fakepassword --region=us-east-1'
+                ' --master-user-password=fakepassword'
                 ' --db-subnet-group-name=fake_db_subnet'
                 ' --vpc-security-group-ids=fake_security_group_id'
                 ' --availability-zones=us-east-1a --storage-type=aurora'
-                ' --backup-retention-period=1 --tags'
+                ' --tags'
             ),
             (
                 'aws --output json rds create-db-instance'
@@ -328,6 +330,23 @@ def testCreateAurora(self):
         ],
     )
 
+  def testCreateAuroraExpress(self):
+    command_strings = self.CreateAurora(
+        {'aws_aurora_express_configuration': True}
+    )
+    self.assertListEqual(
+        command_strings,
+        [
+            (
+                'aws --output json rds create-db-cluster'
+                ' --db-cluster-identifier=pkb-db-cluster-123'
+                ' --engine=aurora-postgresql --region=us-east-1'
+                ' --backup-retention-period=1 --with-express-configuration'
+                ' --tags'
+            ),
+        ],
+    )
+
   def testNoHighAvailability(self):
     spec_dict = {
         'multi_az': False,