Add disk strategy for AWS and support temp db for AWS unmanaged SQL S…

…erver

PiperOrigin-RevId: 596101483

perfkitbenchmarker/disk_strategies.py

@@ -211,6 +211,7 @@ def SetUpDiskOnLinux(self, vm, disk_spec):
 
 class PrepareScratchDiskStrategy:
   """Strategies to prepare scratch disks."""
+  TEMPDB_DISK_LETTER = 'T'
 
   def PrepareScratchDisk(
       self,
@@ -264,7 +265,7 @@ def PrepareWindowsScratchDisk(
       scratch_disk: Union[disk.BaseDisk, disk.StripedDisk],
       disk_spec: disk.BaseDiskSpec,
   ) -> None:
-    """Helper method to format and mount scratch disk.
+    """Format and mount scratch disk.
 
     Args:
       vm: Windows Virtual Machine to create scratch disk on.
@@ -358,3 +359,85 @@ def PrepareWindowsScratchDisk(
     )
 
     vm.scratch_disks.append(scratch_disk)
+
+    if (
+        vm.OS_TYPE in os_types.WINDOWS_SQLSERVER_OS_TYPES
+        and disk_spec.disk_type != 'local'
+    ):
+      self.PrepareTempDbDisk(vm)
+
+  def GetLocalSSDNames(self) -> list[str]:
+    """Names of local ssd device when running Get-PhysicalDisk."""
+    return []
+
+  def GetLocalSSDDeviceIDs(self, vm) -> list[str]:
+    """Get all local ssd device ids."""
+    names = self.GetLocalSSDNames()
+
+    if not names:
+      # This function is not implemented in this cloud provider.
+      logging.info('Temp DB is not supported on this cloud')
+      return []
+
+    clause = ' -or '.join([f'($_.FriendlyName -eq "{name}")' for name in names])
+    clause = '{' + clause + '}'
+
+    # Select the DeviceID from the output
+    # It will be one integer per line.
+    # i.e
+    # 3
+    # 4
+    stdout, _ = vm.RemoteCommand(
+        f'Get-PhysicalDisk | where-object {clause} | Select -exp DeviceID'
+    )
+
+    # int casting will clean up the spacing and make sure we are getting
+    # correct device id.
+    return [
+        str(int(device_id)) for device_id in stdout.split('\n') if device_id
+    ]
+
+  def PrepareTempDbDisk(self, vm: 'virtual_machine.BaseVirtualMachine'):
+    """Format and mount temp db disk for sql server."""
+    # Create and then run a Diskpart script that will initialize the disks,
+    # create a volume, and then format and mount the volume.
+    # This would trigger when we detect local ssds on the VM
+    # and the disk type is not set to local ssd.
+    script = ''
+    local_ssd_disks = self.GetLocalSSDDeviceIDs(vm)
+    if not local_ssd_disks:
+      # No local ssd detected, will not move temp db.
+      return
+
+    for disk_number in local_ssd_disks:
+      # For local SSD disk, set the status to online (if it is not already),
+      # remove any formatting or partitioning on the disks, and convert
+      # it to a dynamic disk so it can be used to create a volume.
+      script += (
+          'select disk %s\n'
+          'online disk noerr\n'
+          'attributes disk clear readonly\n'
+          'clean\n'
+          'convert gpt\n'
+          'convert dynamic\n' % disk_number
+      )
+    if len(local_ssd_disks) > 1:
+      script += 'create volume stripe disk=%s\n' % ','.join(local_ssd_disks)
+    else:
+      script += 'create volume simple\n'
+
+    script += 'format fs=ntfs quick unit=64k\nassign letter={}\n'.format(
+        self.TEMPDB_DISK_LETTER.lower()
+    )
+    vm.RunDiskpartScript(script)
+
+    # Grant user permissions on the drive
+    vm.RemoteCommand(
+        'icacls {}: /grant Users:F /L'.format(self.TEMPDB_DISK_LETTER)
+    )
+    vm.RemoteCommand(
+        'icacls {}: --% /grant Users:(OI)(CI)F /L'.format(
+            self.TEMPDB_DISK_LETTER
+        )
+    )
+    vm.RemoteCommand('mkdir {}:\\TEMPDB'.format(self.TEMPDB_DISK_LETTER))

perfkitbenchmarker/providers/aws/aws_disk_strategies.py

@@ -0,0 +1,34 @@
+# Copyright 2024 PerfKitBenchmarker Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module containing strategies to prepare disks.
+
+This module abstract out the disk algorithm for formatting and creating
+scratch disks.
+"""
+
+from typing import Any
+from absl import flags
+from perfkitbenchmarker import disk_strategies
+
+FLAGS = flags.FLAGS
+
+
+virtual_machine = Any  # pylint: disable=invalid-name
+
+
+class AWSPrepareScratchDiskStrategy(disk_strategies.PrepareScratchDiskStrategy):
+  """Strategies to prepare scratch disk on AWS."""
+
+  def GetLocalSSDNames(self) -> list[str]:
+    return ['NVMe Amazon EC2 NVMe']

perfkitbenchmarker/providers/aws/aws_relational_db.py

@@ -47,9 +47,6 @@ class AWSSQLServerIAASRelationalDb(
 
   CLOUD = provider_info.AWS
 
-  def MoveSQLServerTempDb(self):
-    """Moves the SQL Server temporary database to LocalSSD."""
-
 
 class AWSPostgresIAASRelationalDb(
     postgres_iaas_relational_db.PostgresIAASRelationalDb

perfkitbenchmarker/providers/aws/aws_virtual_machine.py

@@ -32,7 +32,6 @@
 
 from absl import flags
 from perfkitbenchmarker import disk
-from perfkitbenchmarker import disk_strategies
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import flags as pkb_flags
 from perfkitbenchmarker import linux_virtual_machine
@@ -45,6 +44,7 @@
 from perfkitbenchmarker import windows_virtual_machine
 from perfkitbenchmarker.configs import option_decoders
 from perfkitbenchmarker.providers.aws import aws_disk
+from perfkitbenchmarker.providers.aws import aws_disk_strategies
 from perfkitbenchmarker.providers.aws import aws_network
 from perfkitbenchmarker.providers.aws import util
 from six.moves import range
@@ -1548,7 +1548,7 @@ def CreateScratchDisk(self, spec_index, disk_spec):
       nvme_devices = self.GetNVMEDeviceInfo()
       self.PopulateNVMEDevicePath(scratch_disk, nvme_devices)
       self.UpdateDevicePath(scratch_disk)
-    disk_strategies.PrepareScratchDiskStrategy().PrepareScratchDisk(
+    aws_disk_strategies.AWSPrepareScratchDiskStrategy().PrepareScratchDisk(
         self, scratch_disk, disk_spec
     )
 

perfkitbenchmarker/providers/gcp/gce_disk_strategies.py

@@ -16,7 +16,7 @@
 This module abstract out the disk algorithm for formatting and creating
 scratch disks.
 """
-from typing import Any, Union
+from typing import Any
 from absl import flags
 from perfkitbenchmarker import disk
 from perfkitbenchmarker import disk_strategies
@@ -281,69 +281,9 @@ def SetUpDiskOnLinux(self, vm, disk_spec):
 
 class GCEPrepareScratchDiskStrategy(disk_strategies.PrepareScratchDiskStrategy):
   """Strategies to prepare scratch disk on GCE."""
-  TEMPDB_DISK_LETTER = 'T'
-
-  def PrepareWindowsScratchDisk(
-      self,
-      vm: 'virtual_machine.BaseVirtualMachine',
-      scratch_disk: Union[disk.BaseDisk, disk.StripedDisk],
-      disk_spec: disk.BaseDiskSpec,
-  ) -> None:
-    super().PrepareWindowsScratchDisk(vm, scratch_disk, disk_spec)
-    if FLAGS.gce_num_local_ssds > 0 and FLAGS.db_disk_type != 'local':
-      self.PrepareTempDbDisk(vm)
-
-  def PrepareTempDbDisk(self, vm: 'virtual_machine.BaseVirtualMachine'):
-    """Formats and sets up disk for SQL Server TempDB."""
-    # Create and then run a Diskpart script that will initialize the disks,
-    # create a volume, and then format and mount the volume.
-    script = ''
-    stdout, _ = vm.RemoteCommand(
-        'Get-PhysicalDisk | where-object '
-        '{($_.FriendlyName -eq "Google EphemeralDisk") -or '
-        '($_.FriendlyName -eq "nvme_card")} | Select -exp DeviceID'
-    )
-    local_ssd_disks = [
-        int(device_id) for device_id in stdout.split('\n') if device_id
-    ]
-    local_ssd_disks_str = [str(d) for d in local_ssd_disks]
-
-    for disk_number in local_ssd_disks_str:
-      # For local SSD disk, set the status to online (if it is not already),
-      # remove any formatting or partitioning on the disks, and convert
-      # it to a dynamic disk so it can be used to create a volume.
-      script += (
-          'select disk %s\n'
-          'online disk noerr\n'
-          'attributes disk clear readonly\n'
-          'clean\n'
-          'convert gpt\n'
-          'convert dynamic\n' % disk_number
-      )
 
-    if local_ssd_disks:
-      if len(local_ssd_disks_str) > 1:
-        script += 'create volume stripe disk=%s\n' % ','.join(
-            local_ssd_disks_str
-        )
-      else:
-        script += 'create volume simple\n'
-      script += 'format fs=ntfs quick unit=64k\nassign letter={}\n'.format(
-          self.TEMPDB_DISK_LETTER.lower()
-      )
-    vm.RunDiskpartScript(script)
-
-    # Grant user permissions on the drive
-    if local_ssd_disks:
-      vm.RemoteCommand(
-          'icacls {}: /grant Users:F /L'.format(self.TEMPDB_DISK_LETTER)
-      )
-      vm.RemoteCommand(
-          'icacls {}: --% /grant Users:(OI)(CI)F /L'.format(
-              self.TEMPDB_DISK_LETTER
-          )
-      )
-      vm.RemoteCommand('mkdir {}:\\TEMPDB'.format(self.TEMPDB_DISK_LETTER))
+  def GetLocalSSDNames(self):
+    return ['Google EphemeralDisk', 'nvme_card']
 
 
 def _GenerateDiskNamePrefix(

perfkitbenchmarker/providers/gcp/gcp_relational_db.py

@@ -42,8 +42,6 @@
 
 FLAGS = flags.FLAGS
 
-TEMPDB_DISK_LETTER = 'T'
-
 GCP_DATABASE_VERSION_MAPPING = {
     sql_engine_utils.MYSQL: {
         '5.5': 'MYSQL_5_5',
@@ -113,52 +111,6 @@ class GCPSQLServerIAASRelationalDb(
 
   CLOUD = provider_info.GCP
 
-  def MoveSQLServerTempDb(self):
-    vms = [self.server_vm]
-    if self.spec.high_availability_type == 'AOAG' and self.replica_vms:
-      vms.append(self.replica_vms[0])
-
-    # Moves the SQL Server temporary database to LocalSSD.
-    for vm in vms:
-      stdout, _ = vm.RemoteCommand(
-          'Get-PSDrive -PSProvider FileSystem | Select Name'
-      )
-
-      drive_list = [
-          str(drive.strip().replace('\r', ''))
-          for drive in stdout.split('\n')
-          if drive
-      ]
-
-      if TEMPDB_DISK_LETTER in drive_list:
-        stdout, _ = vm.RemoteCommand(
-            'sqlcmd -h -1 -Q "SET NOCOUNT '
-            ' ON; select f.name + CASE WHEN '
-            "f.type = 1 THEN '.ldf' "
-            "ELSE '.mdf' END "
-            'FROM sys.master_files '
-            'f WHERE f.database_id'
-            " = DB_ID('tempdb');\""
-        )
-        tmp_db_files_list = [
-            str(tmp_file.strip().replace('\r', ''))
-            for tmp_file in stdout.split('\n')
-            if tmp_file
-        ]
-
-        for tmp_db_file in tmp_db_files_list:
-          tmp_db_name = tmp_db_file.split('.')[0]
-          vm.RemoteCommand(
-              'sqlcmd -Q "ALTER DATABASE tempdb '
-              'MODIFY FILE (NAME = [{}], '
-              "FILENAME = '{}:\\TEMPDB\\{}');\"".format(
-                  tmp_db_name, TEMPDB_DISK_LETTER, tmp_db_file
-              )
-          )
-
-        vm.RemoteCommand('net stop mssqlserver /y')
-        vm.RemoteCommand('net start mssqlserver')
-
 
 class GCPPostgresIAASRelationalDb(
     postgres_iaas_relational_db.PostgresIAASRelationalDb

perfkitbenchmarker/sqlserver_iaas_relational_db.py

@@ -47,6 +47,8 @@
 
 DEFAULT_ENGINE_VERSION = sql_engine_utils.SQLSERVER_ENTERPRISE
 
+TEMPDB_DISK_LETTER = "T"
+
 
 class SQLServerIAASRelationalDb(iaas_relational_db.IAASRelationalDb):
   """Object representing a IAAS relational database Service."""
@@ -103,6 +105,53 @@ def SetVms(self, vm_groups):
       self.controller_vm = vm_groups["servers"][-1]
       self.replica_vms = vm_groups["servers"][1:-1]
 
+  def MoveSQLServerTempDb(self):
+    """Moves the SQL Server temporary database to LocalSSD."""
+    vms = [self.server_vm]
+    if self.spec.high_availability_type == "AOAG" and self.replica_vms:
+      vms.append(self.replica_vms[0])
+
+    # Moves the SQL Server temporary database to LocalSSD.
+    for vm in vms:
+      stdout, _ = vm.RemoteCommand(
+          "Get-PSDrive -PSProvider FileSystem | Select Name"
+      )
+
+      drive_list = [
+          str(drive.strip().replace("\r", ""))
+          for drive in stdout.split("\n")
+          if drive
+      ]
+
+      if TEMPDB_DISK_LETTER in drive_list:
+        stdout, _ = vm.RemoteCommand(
+            'sqlcmd -h -1 -Q "SET NOCOUNT '
+            " ON; select f.name + CASE WHEN "
+            "f.type = 1 THEN '.ldf' "
+            "ELSE '.mdf' END "
+            "FROM sys.master_files "
+            "f WHERE f.database_id"
+            " = DB_ID('tempdb');\""
+        )
+        tmp_db_files_list = [
+            str(tmp_file.strip().replace("\r", ""))
+            for tmp_file in stdout.split("\n")
+            if tmp_file
+        ]
+
+        for tmp_db_file in tmp_db_files_list:
+          tmp_db_name = tmp_db_file.split(".")[0]
+          vm.RemoteCommand(
+              'sqlcmd -Q "ALTER DATABASE tempdb '
+              "MODIFY FILE (NAME = [{}], "
+              "FILENAME = '{}:\\TEMPDB\\{}');\"".format(
+                  tmp_db_name, TEMPDB_DISK_LETTER, tmp_db_file
+              )
+          )
+
+        vm.RemoteCommand("net stop mssqlserver /y")
+        vm.RemoteCommand("net start mssqlserver")
+
   def _SetupWindowsUnamangedDatabase(self):
     self.spec.database_username = "sa"
     self.spec.database_password = db_util.GenerateRandomDbPassword()
@@ -621,10 +670,6 @@ def ConfigureSQLServer(vm, username: str, password: str):
   vm.RemoteCommand("net start mssqlserver")
 
 
-def MoveSQLServerTempDb():
-  """Moves the SQL Server temporary database to LocalSSD."""
-
-
 def _TuneForSQL(vm):
   """Set TuneD settings specific to SQL Server on RedHat."""
   tune_settings = (