Add support for flexible server PremiumSSD v2.

PiperOrigin-RevId: 853313032

Author: bvliu

perfkitbenchmarker/providers/azure/azure_flexible_server.py

@@ -22,6 +22,7 @@
 import datetime
 import json
 import logging
+import re
 import time
 from typing import Any, Tuple
 
@@ -57,6 +58,10 @@
 CREATE_AZURE_DB_TIMEOUT = 60 * 120
 
 
+class AzureAsyncOperationInProgressError(Exception):
+  """Exception raised when an Azure async operation is still in progress."""
+
+
 class AzureFlexibleServer(azure_relational_db.AzureRelationalDb):
   """An object representing an Azure Flexible server.
 
@@ -91,10 +96,6 @@ def __init__(self, relational_db_spec: Any):
       self.storage_type = (
           self.spec.db_disk_spec.disk_type or azure_disk.PREMIUM_STORAGE
       )
-      if self.storage_type == azure_disk.PREMIUM_STORAGE_V2:
-        raise errors.Config.InvalidValue(
-            'Premium Storage v2 is not supported for Postgres Flexible Server.'
-        )
 
   @staticmethod
   def GetDefaultEngineVersion(engine: str) -> str:
@@ -151,6 +152,13 @@ def _Create(self) -> None:
         self.spec.engine_version,
         '--yes',
     ]
+    if (
+        self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES
+        and self.storage_type == azure_disk.PREMIUM_STORAGE_V2
+    ):
+      # Postgres PremiumV2_LRS requires creation through REST API.
+      self._CreatePremiumSsdV2()
+      return
     if self.storage_type:
       cmd.extend([
           '--storage-type',
@@ -177,6 +185,91 @@ def _Create(self) -> None:
 
     vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT)
 
+  @vm_util.Retry(
+      timeout=CREATE_AZURE_DB_TIMEOUT,
+      retryable_exceptions=(AzureAsyncOperationInProgressError,),
+  )
+  def _WaitForAzureAsyncOperation(self, async_url: str) -> None:
+    """Polls Azure async operation URL until completion."""
+    cmd = [azure.AZURE_PATH, 'rest', '--method', 'GET', '--uri', async_url]
+    stdout, _, _ = vm_util.IssueCommand(cmd, raise_on_failure=True)
+    status_info = json.loads(stdout)
+    status = status_info.get('status')
+    if not status and 'properties' in status_info:
+      status = status_info['properties'].get('status')
+
+    if status == 'Succeeded':
+      return
+    if status == 'Failed':
+      raise errors.Resource.CreationError(
+          f'Azure async operation failed: {stdout}'
+      )
+    raise AzureAsyncOperationInProgressError(
+        f'Azure async operation status unknown: {status_info}'
+    )
+
+  def _CreatePremiumSsdV2(self) -> None:
+    """Creates a Postgres PremiumV2_LRS instance."""
+    subscription_id = util.GetSubscriptionId()
+    uri = (
+        f'https://management.azure.com/subscriptions/{subscription_id}'
+        f'/resourceGroups/{self.resource_group.name}'
+        '/providers/Microsoft.DBforPostgreSQL/flexibleServers/'
+        f'{self.instance_id}?api-version=2024-08-01'
+    )
+    body = {
+        'location': self.region,
+        'sku': {
+            'name': self.spec.db_spec.machine_type,
+            'tier': self.spec.db_tier,
+        },
+        'properties': {
+            'administratorLogin': self.spec.database_username,
+            'administratorLoginPassword': self.spec.database_password,
+            'version': self.spec.engine_version,
+            'storage': {
+                'storageSizeGB': self.spec.db_disk_spec.disk_size,
+                'type': azure_disk.PREMIUM_STORAGE_V2,
+            },
+            'backup': {
+                'backupRetentionDays': 7,
+                'geoRedundantBackup': 'Disabled',
+            },
+            'network': {'publicNetworkAccess': 'Enabled'},
+            'highAvailability': {'mode': 'ZoneRedundant'},
+        },
+    }
+    if self.spec.db_disk_spec.provisioned_iops:
+      body['properties']['storage'][
+          'iops'
+      ] = self.spec.db_disk_spec.provisioned_iops
+    if self.spec.db_disk_spec.provisioned_throughput:
+      body['properties']['storage'][
+          'throughput'
+      ] = self.spec.db_disk_spec.provisioned_throughput
+    cmd = [
+        azure.AZURE_PATH,
+        'rest',
+        '--method',
+        'PUT',
+        '--uri',
+        uri,
+        '--body',
+        json.dumps(body),
+        '--verbose',
+    ]
+    _, stderr, _ = vm_util.IssueCommand(
+        cmd, timeout=CREATE_AZURE_DB_TIMEOUT, raise_on_failure=False
+    )
+    match = re.search(r"'Azure-AsyncOperation': '([^']*)'", stderr)
+    if not match:
+      raise errors.Resource.CreationError(
+          'Could not find Azure-AsyncOperation header in az rest verbose'
+          ' output.'
+      )
+    async_operation_url = match.group(1)
+    self._WaitForAzureAsyncOperation(async_operation_url)
+
   def GetAzCommandForEngine(self) -> str:
     engine = self.spec.engine
     if engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES:

tests/providers/azure/azure_flexible_server_test.py

@@ -1,6 +1,7 @@
 import datetime
 import inspect
 import json
+import time
 import unittest
 
 from absl import flags
@@ -231,5 +232,157 @@ def testConstructMysqlWithThroughputRaises(self):
       bm_spec.ConstructRelationalDb()
 
 
+class AzureFlexibleServerPremiumV2CreateTestCase(
+    pkb_common_test_case.PkbCommonTestCase
+):
+
+  def setUp(self):
+    super().setUp()
+    FLAGS.cloud = provider_info.AZURE
+    FLAGS.run_uri = 'test_run_uri'
+    self.resource_group_patch = self.enter_context(
+        mock.patch.object(
+            azure_network, 'GetResourceGroup', return_value=mock.Mock()
+        )
+    )
+    self.enter_context(
+        mock.patch.object(util, 'GetSubscriptionId', return_value='test-sub')
+    )
+    self.enter_context(mock.patch.object(time, 'sleep'))
+
+  def testCreatePostgresPremiumV2(self):
+    mock_cmd = self.MockIssueCommand({
+        'rest --method PUT': [(
+            '',
+            "'Azure-AsyncOperation': 'https://management.azure.com/async_op'",
+            0,
+        )],
+        'rest --method GET': [
+            ('{"status": "InProgress"}', '', 0),
+            ('{"status": "Succeeded"}', '', 0),
+        ],
+    })
+    yaml_spec = inspect.cleandoc(f"""
+        sysbench:
+          relational_db:
+            cloud: {provider_info.AZURE}
+            engine: {sql_engine_utils.FLEXIBLE_SERVER_POSTGRES}
+            engine_version: '13'
+            db_tier: GeneralPurpose
+            db_spec:
+              {provider_info.AZURE}:
+                machine_type: Standard_D2ds_v4
+                zone: eastus
+            db_disk_spec:
+              {provider_info.AZURE}:
+                disk_size: 128
+                disk_type: PremiumV2_LRS
+                provisioned_iops: 1000
+                provisioned_throughput: 200
+            vm_groups:
+              clients:
+                vm_spec: *default_dual_core
+                disk_spec: *default_500_gb
+    """)
+    bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
+        yaml_spec, 'sysbench'
+    )
+    bm_spec.ConstructRelationalDb()
+
+    bm_spec.relational_db._Create()
+
+    self.assertEqual(mock_cmd.func_to_mock.call_count, 3)
+    put_call, get_call1, get_call2 = mock_cmd.func_to_mock.call_args_list
+
+    # Check az rest PUT call
+    put_cmd = ' '.join(put_call[0][0])
+    with self.subTest(name='RestApi'):
+      self.assertIn('rest --method PUT', put_cmd)
+      self.assertIn('PremiumV2_LRS', put_cmd)
+      self.assertIn('"iops": 1000', put_cmd)
+      self.assertIn('"throughput": 200', put_cmd)
+
+    # Check az rest GET calls
+    get_cmd1 = ' '.join(get_call1[0][0])
+    self.assertIn('rest --method GET', get_cmd1)
+    get_cmd2 = ' '.join(get_call2[0][0])
+    self.assertIn('rest --method GET', get_cmd2)
+
+  def testCreatePostgresPremiumV2Fails(self):
+    self.MockIssueCommand({
+        'rest --method PUT': [(
+            '',
+            "'Azure-AsyncOperation': 'https://management.azure.com/async_op'",
+            0,
+        )],
+        'rest --method GET': [('{"status": "Failed"}', '', 0)],
+    })
+    yaml_spec = inspect.cleandoc(f"""
+        sysbench:
+          relational_db:
+            cloud: {provider_info.AZURE}
+            engine: {sql_engine_utils.FLEXIBLE_SERVER_POSTGRES}
+            engine_version: '13'
+            db_tier: GeneralPurpose
+            db_spec:
+              {provider_info.AZURE}:
+                machine_type: Standard_D2ds_v4
+                zone: eastus
+            db_disk_spec:
+              {provider_info.AZURE}:
+                disk_size: 128
+                disk_type: PremiumV2_LRS
+            vm_groups:
+              clients:
+                vm_spec: *default_dual_core
+                disk_spec: *default_500_gb
+    """)
+    bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
+        yaml_spec, 'sysbench'
+    )
+    bm_spec.ConstructRelationalDb()
+
+    with self.assertRaises(errors.Resource.CreationError):
+      bm_spec.relational_db._Create()
+
+  def testCreatePostgresPremiumV2NoAsyncHeader(self):
+    self.MockIssueCommand(
+        {
+            'rest --method PUT': [(
+                '',
+                'some other stderr',
+                0,
+            )]
+        }
+    )
+    yaml_spec = inspect.cleandoc(f"""
+        sysbench:
+          relational_db:
+            cloud: {provider_info.AZURE}
+            engine: {sql_engine_utils.FLEXIBLE_SERVER_POSTGRES}
+            engine_version: '13'
+            db_tier: GeneralPurpose
+            db_spec:
+              {provider_info.AZURE}:
+                machine_type: Standard_D2ds_v4
+                zone: eastus
+            db_disk_spec:
+              {provider_info.AZURE}:
+                disk_size: 128
+                disk_type: PremiumV2_LRS
+            vm_groups:
+              clients:
+                vm_spec: *default_dual_core
+                disk_spec: *default_500_gb
+    """)
+    bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
+        yaml_spec, 'sysbench'
+    )
+    bm_spec.ConstructRelationalDb()
+
+    with self.assertRaises(errors.Resource.CreationError):
+      bm_spec.relational_db._Create()
+
+
 if __name__ == '__main__':
   unittest.main()