From 71a74d8295d69d0e54a2d165d29c880dab54bb73 Mon Sep 17 00:00:00 2001 From: Brandon Liu Date: Wed, 29 Apr 2026 08:50:52 -0700 Subject: [PATCH] diffbase, omitted PiperOrigin-RevId: 907612924 --- .../providers/aws/aws_aurora_db.py | 51 +++++++++++++++---- .../providers/aws/aws_aurora_dsql_db.py | 6 ++- .../providers/aws/aws_relational_db.py | 16 +++--- perfkitbenchmarker/providers/aws/flags.py | 5 ++ .../resources/aws_relational_db_spec.py | 9 ++++ tests/providers/aws/aws_relational_db_test.py | 27 ++++++++-- 6 files changed, 92 insertions(+), 22 deletions(-) diff --git a/perfkitbenchmarker/providers/aws/aws_aurora_db.py b/perfkitbenchmarker/providers/aws/aws_aurora_db.py index b1d0798796..190e3dbdc1 100644 --- a/perfkitbenchmarker/providers/aws/aws_aurora_db.py +++ b/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]: diff --git a/perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py b/perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py index 6deb69c711..e1604652db 100644 --- a/perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py +++ b/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. diff --git a/perfkitbenchmarker/providers/aws/aws_relational_db.py b/perfkitbenchmarker/providers/aws/aws_relational_db.py index 4bde116c1d..9e75be24c1 100644 --- a/perfkitbenchmarker/providers/aws/aws_relational_db.py +++ b/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 diff --git a/perfkitbenchmarker/providers/aws/flags.py b/perfkitbenchmarker/providers/aws/flags.py index 197dd358f2..897ae0b37a 100644 --- a/perfkitbenchmarker/providers/aws/flags.py +++ b/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, diff --git a/perfkitbenchmarker/resources/aws_relational_db_spec.py b/perfkitbenchmarker/resources/aws_relational_db_spec.py index 165dd9f159..4282979965 100644 --- a/perfkitbenchmarker/resources/aws_relational_db_spec.py +++ b/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 + ) diff --git a/tests/providers/aws/aws_relational_db_test.py b/tests/providers/aws/aws_relational_db_test.py index ac550eeffe..8e5087b23e 100644 --- a/tests/providers/aws/aws_relational_db_test.py +++ b/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,