Support setting up temp db for Azure. Azure uses D drive for temp disk.

PiperOrigin-RevId: 596052324

perfkitbenchmarker/disk_strategies.py

@@ -230,7 +230,7 @@ def PrepareLinuxScratchDisk(
       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: Linux Virtual Machine to create scratch disk on.

perfkitbenchmarker/providers/azure/azure_disk_strategies.py

@@ -0,0 +1,43 @@
+# 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
+from perfkitbenchmarker.providers.azure import azure_disk
+
+FLAGS = flags.FLAGS
+
+
+virtual_machine = Any  # pylint: disable=invalid-name
+
+
+class AzurePrepareScratchDiskStrategy(
+    disk_strategies.PrepareScratchDiskStrategy
+):
+  """Strategies to prepare scratch disk on AWS."""
+
+  def GetLocalSSDNames(self) -> list[str]:
+    # This is only for Ls-Series local ssd
+    # Temp disk is automatically setup on Azure and default to disk D.
+    return ['Microsoft NVMe Direct Disk']
+
+  def PrepareTempDbDisk(self, vm: 'virtual_machine.BaseVirtualMachine'):
+    if azure_disk.HasTempDrive(vm.machine_type):
+      vm.RemoteCommand('mkdir D:\\TEMPDB')

perfkitbenchmarker/providers/azure/azure_relational_db.py

@@ -62,9 +62,7 @@ class AzureSQLServerIAASRelationalDb(
   """A AWS IAAS database resource."""
 
   CLOUD = provider_info.AZURE
-
-  def MoveSQLServerTempDb(self):
-    """Moves the SQL Server temporary database to LocalSSD."""
+  TEMPDB_DISK_LETTER = 'D'
 
 
 class AzurePostgresIAASRelationalDb(

perfkitbenchmarker/providers/azure/azure_virtual_machine.py

@@ -24,7 +24,6 @@
 operate on the VM: boot, shutdown, etc.
 """
 
-
 import abc
 import collections
 import itertools
@@ -39,7 +38,6 @@
 from absl import flags
 from perfkitbenchmarker import custom_virtual_machine_spec
 from perfkitbenchmarker import disk
-from perfkitbenchmarker import disk_strategies
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import linux_virtual_machine
 from perfkitbenchmarker import os_types
@@ -52,6 +50,7 @@
 from perfkitbenchmarker.configs import option_decoders
 from perfkitbenchmarker.providers import azure
 from perfkitbenchmarker.providers.azure import azure_disk
+from perfkitbenchmarker.providers.azure import azure_disk_strategies
 from perfkitbenchmarker.providers.azure import azure_network
 from perfkitbenchmarker.providers.azure import util
 from six.moves import range
@@ -1056,7 +1055,7 @@ def CreateScratchDisk(self, _, disk_spec):
       disks.append(data_disk)
 
     scratch_disk = self._CreateScratchDiskFromDisks(disk_spec, disks)
-    disk_strategies.PrepareScratchDiskStrategy().PrepareScratchDisk(
+    azure_disk_strategies.AzurePrepareScratchDiskStrategy().PrepareScratchDisk(
         self, scratch_disk, disk_spec
     )
 

perfkitbenchmarker/sqlserver_iaas_relational_db.py

@@ -47,11 +47,10 @@
 
 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."""
+  TEMPDB_DISK_LETTER = "T"
 
   ENGINE = sql_engine_utils.SQLSERVER
 
@@ -123,7 +122,7 @@ def MoveSQLServerTempDb(self):
           if drive
       ]
 
-      if TEMPDB_DISK_LETTER in drive_list:
+      if self.TEMPDB_DISK_LETTER in drive_list:
         stdout, _ = vm.RemoteCommand(
             'sqlcmd -h -1 -Q "SET NOCOUNT '
             " ON; select f.name + CASE WHEN "
@@ -145,7 +144,7 @@ def MoveSQLServerTempDb(self):
               'sqlcmd -Q "ALTER DATABASE tempdb '
               "MODIFY FILE (NAME = [{}], "
               "FILENAME = '{}:\\TEMPDB\\{}');\"".format(
-                  tmp_db_name, TEMPDB_DISK_LETTER, tmp_db_file
+                  tmp_db_name, self.TEMPDB_DISK_LETTER, tmp_db_file
               )
           )