Implement backup restore for DSQL

PiperOrigin-RevId: 845429012

Author: ScottLinnn

perfkitbenchmarker/linux_benchmarks/benchbase_benchmark.py

@@ -72,6 +72,7 @@
 _WARMUP_DURATION = benchbase.BENCHBASE_WARMUP_DURATION
 _WORKLOAD_DURATION = benchbase.BENCHBASE_WORKLOAD_DURATION
 _WAREHOUSES = benchbase.BENCHBASE_WAREHOUSES
+_RECOVERY_POINT_ARN = aws_aurora_dsql_db.AWS_AURORA_DSQL_RECOVERY_POINT_ARN
 
 
 def GetConfig(user_config: Dict[str, Any]) -> Dict[str, Any]:
@@ -87,7 +88,8 @@ def GetConfig(user_config: Dict[str, Any]) -> Dict[str, Any]:
   if _COOLDOWN_DURATION.value > 0 and _WARMUP_DURATION.value <= 0:
     raise errors.Config.InvalidValue(
         'benchbase_warmup_duration must be positive if'
-        ' benchbase_cooldown_duration is positive.')
+        ' benchbase_cooldown_duration is positive.'
+    )
   return config
 
 
@@ -107,6 +109,7 @@ 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
@@ -117,19 +120,27 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec) -> None:
     # 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
-      else 'auroradsql'
-  )
-  load_command: str = (
-      f'source /etc/profile.d/maven.sh && cd {benchbase.BENCHBASE_DIR} && mvn'
-      f' clean compile exec:java -P {profile} -Dexec.args="-b tpcc -c'
-      f' {benchbase.CONFIG_FILE_NAME} --create=true --load=true'
-      ' --execute=false"'
-  )
 
-  client_vm.RemoteCommand(load_command)
+  dsql_create_from_raw: bool = (
+      FLAGS.db_engine == sql_engine_utils.AURORA_DSQL_POSTGRES
+      and _RECOVERY_POINT_ARN.value is None
+  )
+  if (
+      dsql_create_from_raw
+      or FLAGS.db_engine == sql_engine_utils.SPANNER_POSTGRES
+  ):
+    profile: str = (
+        'postgres'
+        if FLAGS.db_engine == sql_engine_utils.SPANNER_POSTGRES
+        else 'auroradsql'
+    )
+    load_command: str = (
+        f'source /etc/profile.d/maven.sh && cd {benchbase.BENCHBASE_DIR} && mvn'
+        f' clean compile exec:java -P {profile} -Dexec.args="-b tpcc -c'
+        f' {benchbase.CONFIG_FILE_NAME} --create=true --load=true'
+        ' --execute=false"'
+    )
+    client_vm.RemoteCommand(load_command)
 
 
 def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]:

perfkitbenchmarker/providers/aws/aws_aurora_dsql_db.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Managed relational database provisioning and teardown for AWS Aurora DSQL."""
 
+import functools
 import json
 from typing import Any
 
@@ -27,11 +28,16 @@
 from perfkitbenchmarker.providers.aws import util
 
 
-# TODO(shuninglin): Add cluster creation from a backup.
 # TODO(shuninglin): Add reaper for this new resource.
 
 FLAGS = flags.FLAGS
 
+AWS_AURORA_DSQL_RECOVERY_POINT_ARN = flags.DEFINE_string(
+    'aws_aurora_dsql_recovery_point_arn',
+    None,
+    'The ARN of the recovery point to restore AWS Aurora DSQL cluster from. If '
+    'not provided, a new cluster is created from scratch.',
+)
 DEFAULT_AURORA_DSQL_POSTGRES_VERSION = '16.2'
 
 _MAP_ENGINE_TO_DEFAULT_VERSION = {
@@ -69,7 +75,15 @@ class AwsAuroraDsqlRelationalDb(aws_relational_db.BaseAwsRelationalDb):
   def __init__(self, dsql_spec: AwsAuroraDsqlSpec):
     super().__init__(dsql_spec)
     self.cluster_id = None
+    self.cluster_arn = None
     self.assigned_name = f'pkb-{FLAGS.run_uri}'
+    self.use_backup = bool(AWS_AURORA_DSQL_RECOVERY_POINT_ARN.value)
+    self.restore_job_id = None
+
+  @functools.cached_property
+  def account_id(self) -> str:
+    """Returns the AWS account ID."""
+    return util.GetAccount()
 
   # DSQL has different format for tags:
   # https://docs.aws.amazon.com/cli/v1/reference/rds/create-db-cluster.html
@@ -83,6 +97,64 @@ def _MakeDsqlTags(self) -> list[str]:
     return [formatted_tags_str]
 
   def _Create(self) -> None:
+    """Creates AWS Aurora DSQL cluster, from backup if recovery point ARN is provided."""
+    if not self.use_backup:
+      self._CreateRawCluster()
+      return
+    if self.restore_job_id:
+      logging.info(
+          'Restore job %s already exists. Skipping creation.',
+          self.restore_job_id,
+      )
+      return
+    cmd = util.AWS_PREFIX + [
+        'backup',
+        'start-restore-job',
+        '--recovery-point-arn',
+        AWS_AURORA_DSQL_RECOVERY_POINT_ARN.value,
+        '--iam-role-arn',
+        (
+            f'arn:aws:iam::{self.account_id}:role/service-role/'
+            'AWSBackupDefaultServiceRole'
+        ),
+        '--metadata',
+        '{"regionalConfig": "[{\\"region\\": \\"%s\\",'
+        ' \\"isDeletionProtectionEnabled\\": false}]"}'
+        % self.region,
+    ]
+    stdout, _, _ = vm_util.IssueCommand(cmd)
+    response = json.loads(stdout)
+    self.restore_job_id = response['RestoreJobId']
+    if self.restore_job_id:
+      # Mark created so we don't try to create it again on a retry.
+      self.created = True
+
+  def _DescribeRestoreJob(self, job_id: str) -> dict[str, Any]:
+    """Describes the restore job."""
+    cmd = util.AWS_PREFIX + [
+        'backup',
+        'describe-restore-job',
+        '--restore-job-id',
+        job_id,
+    ]
+    stdout, _, _ = vm_util.IssueCommand(cmd)
+    return json.loads(stdout)
+
+  def _AddTagsToCluster(self, cluster_arn: str) -> None:
+    """Adds tags to the DSQL cluster."""
+    cmd = (
+        util.AWS_PREFIX
+        + [
+            'dsql',
+            'tag-resource',
+            '--resource-arn=%s' % cluster_arn,
+            '--tags',
+        ]
+        + self._MakeDsqlTags()
+    )
+    vm_util.IssueCommand(cmd)
+
+  def _CreateRawCluster(self) -> None:
     """Creates the AWS Aurora DSQL instance.
 
     Raises:
@@ -108,6 +180,9 @@ def _Create(self) -> None:
     self.cluster_id = response['identifier']
 
   def _DescribeCluster(self) -> dict[str, Any] | None:
+    if not self.cluster_id:
+      logging.info('Cluster id is not set.')
+      return None
     cmd = util.AWS_PREFIX + [
         'dsql',
         'get-cluster',
@@ -122,8 +197,29 @@ def _DescribeCluster(self) -> dict[str, Any] | None:
 
   def _IsReady(self, timeout=aws_relational_db.IS_READY_TIMEOUT) -> bool:
     """Returns true if the cluster is ready."""
-    json_output = self._DescribeCluster()
-    return bool(json_output and json_output['status'] == 'ACTIVE')
+    if self.use_backup:
+      if not self.restore_job_id:
+        return False
+      job_description = self._DescribeRestoreJob(self.restore_job_id)
+      status = job_description['Status']
+      if status == 'COMPLETED':
+        self.cluster_id = job_description['CreatedResourceArn'].split('/')[-1]
+        self.cluster_arn = job_description['CreatedResourceArn']
+        return True
+      if status in ['ABORTED', 'FAILED']:
+        raise errors.Resource.CreationError(
+            f'Restore job {self.restore_job_id} failed with status {status}'
+        )
+      return False
+    else:
+      json_output = self._DescribeCluster()
+      return bool(json_output and json_output['status'] == 'ACTIVE')
+
+  def _PostCreate(self) -> None:
+    """Add tags if we are restoring from backup."""
+    super()._PostCreate()
+    if self.use_backup:
+      self._AddTagsToCluster(self.cluster_arn)
 
   def _Exists(self) -> bool:
     """Returns true if the underlying cluster exists."""

tests/linux_benchmarks/benchbase_benchmark_test.py

@@ -21,6 +21,7 @@
 import mock
 from perfkitbenchmarker import benchmark_spec
 from perfkitbenchmarker import configs
+from perfkitbenchmarker import sql_engine_utils
 from perfkitbenchmarker.linux_benchmarks import benchbase_benchmark
 from perfkitbenchmarker.linux_packages import benchbase
 from tests import pkb_common_test_case
@@ -39,9 +40,17 @@ def setUp(self):
     self.mock_benchmark_spec.vms = [self.mock_vm]
     self.mock_benchmark_spec.relational_db = mock.Mock()
     self.mock_benchmark_spec.relational_db.GetResourceMetadata.return_value = {}
+    self.mock_benchmark_spec.relational_db.cluster_id = 'cluster'
+    self.mock_benchmark_spec.relational_db.region = 'region'
     self.mock_load_config = self.enter_context(
         mock.patch.object(configs, 'LoadConfig', autospec=True)
     )
+    self.mock_create_config = self.enter_context(
+        mock.patch.object(benchbase, 'CreateConfigFile', autospec=True)
+    )
+    self.mock_override_endpoint = self.enter_context(
+        mock.patch.object(benchbase, 'OverrideEndpoint', autospec=True)
+    )
 
   def test_get_config(self):
     user_config = {'key': 'value'}
@@ -52,13 +61,44 @@ def test_get_config(self):
         benchbase_benchmark.BENCHMARK_NAME,
     )
 
-  @flagsaver.flagsaver(db_engine='spanner-postgres')
-  @mock.patch.object(benchbase, 'CreateConfigFile', autospec=True)
-  def test_prepare(self, mock_create_config):
+  @flagsaver.flagsaver(db_engine=sql_engine_utils.SPANNER_POSTGRES)
+  def test_prepare_spanner_postgres_loads(self):
+    benchbase_benchmark.Prepare(self.mock_benchmark_spec)
+
+    self.mock_vm.Install.assert_called_once_with('benchbase')
+    self.mock_create_config.assert_called_once_with(self.mock_vm)
+    self.mock_override_endpoint.assert_not_called()
+    self.mock_vm.RemoteCommand.assert_called_once()
+    self.assertIn(
+        '--create=true --load=true', self.mock_vm.RemoteCommand.call_args[0][0]
+    )
+    self.assertIn('-P postgres', self.mock_vm.RemoteCommand.call_args[0][0])
+
+  @flagsaver.flagsaver(
+      db_engine=sql_engine_utils.AURORA_DSQL_POSTGRES,
+      aws_aurora_dsql_recovery_point_arn=None,
+  )
+  def test_prepare_dsql_raw_loads(self):
     benchbase_benchmark.Prepare(self.mock_benchmark_spec)
+    self.mock_vm.Install.assert_called_once_with('benchbase')
+    self.mock_create_config.assert_called_once_with(self.mock_vm)
+    self.mock_override_endpoint.assert_called_once()
+    self.mock_vm.RemoteCommand.assert_called_once()
+    self.assertIn(
+        '--create=true --load=true', self.mock_vm.RemoteCommand.call_args[0][0]
+    )
+    self.assertIn('-P auroradsql', self.mock_vm.RemoteCommand.call_args[0][0])
 
+  @flagsaver.flagsaver(
+      db_engine=sql_engine_utils.AURORA_DSQL_POSTGRES,
+      aws_aurora_dsql_recovery_point_arn='arn',
+  )
+  def test_prepare_dsql_restore_skips_loading(self):
+    benchbase_benchmark.Prepare(self.mock_benchmark_spec)
     self.mock_vm.Install.assert_called_once_with('benchbase')
-    mock_create_config.assert_called_once_with(self.mock_vm)
+    self.mock_create_config.assert_called_once_with(self.mock_vm)
+    self.mock_override_endpoint.assert_called_once()
+    self.mock_vm.RemoteCommand.assert_not_called()
 
   @mock.patch('time.sleep')
   @mock.patch.object(benchbase, 'ParseResults', autospec=True)