Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 114 additions & 59 deletions perfkitbenchmarker/providers/azure/azure_flexible_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,67 +122,26 @@ def GetResourceMetadata(self) -> dict[str, Any]:

def _Create(self) -> None:
"""Creates the Azure database instance."""
ha_flag = ENABLE_HA if self.spec.high_availability else DISABLE_HA
# Consider adding --zone and maybe --standby-zone for parity with GCP.
cmd = [
azure.AZURE_PATH,
self.GetAzCommandForEngine(),
self.SERVER_TYPE,
'create',
'--resource-group',
self.resource_group.name,
'--name',
self.instance_id,
'--location',
self.region,
'--high-availability',
ha_flag,
'--admin-user',
self.spec.database_username,
'--admin-password',
self.spec.database_password,
'--storage-size',
str(self.spec.db_disk_spec.disk_size),
'--sku-name',
self.spec.db_spec.machine_type,
'--tier',
self.spec.db_tier,
'--version',
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',
self.storage_type,
])
if (
self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL
and self.spec.db_disk_spec.provisioned_iops is not None
):
cmd.extend([
'--auto-scale-iops',
'Disabled',
])
if self.spec.db_disk_spec.provisioned_iops:
cmd.extend([
'--iops',
str(self.spec.db_disk_spec.provisioned_iops),
])
if self.spec.db_disk_spec.provisioned_throughput:
cmd.extend([
'--throughput',
str(self.spec.db_disk_spec.provisioned_throughput),
])

vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT)
try:
self._CreatePremiumSsdV2()
return
except errors.Resource.CreationError as e:
if self._IsFallbackableError(e):
logging.warning(
'Failed to create Postgres PremiumV2_LRS instance: %s. '
'Falling back to Premium_LRS.',
e,
)
self._FallbackToPremiumLrs()
# Fall through to _CreateStandard()
else:
raise

self._CreateStandard()

@vm_util.Retry(
timeout=CREATE_AZURE_DB_TIMEOUT,
Expand Down Expand Up @@ -257,18 +216,114 @@ def _CreatePremiumSsdV2(self) -> None:
json.dumps(body),
'--verbose',
]
_, stderr, _ = vm_util.IssueCommand(
stdout, stderr, retcode = vm_util.IssueCommand(
cmd, timeout=CREATE_AZURE_DB_TIMEOUT, raise_on_failure=False
)
if retcode != 0:
raise errors.Resource.CreationError(
f'Azure REST API call failed (retcode={retcode}). '
f'Stdout: {stdout}, Stderr: {stderr}'
)
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.'
f' output. Stdout: {stdout}, Stderr: {stderr}'
)
async_operation_url = match.group(1)
self._WaitForAzureAsyncOperation(async_operation_url)

def _IsFallbackableError(self, error: Exception) -> bool:
"""Returns True if the error is fallbackable."""
error_str = str(error)
return (
'SkuNotAvailable' in error_str
or 'ServerDropping' in error_str
or 'AvailabilityZone' in error_str
or 'not enough resources' in error_str.lower()
)

def _FallbackToPremiumLrs(self) -> None:
"""Falls back to Premium_LRS storage."""
logging.info('Initiating deletion of failed instance %s', self.instance_id)
try:
self._Delete()
except Exception as e: # pylint: disable=broad-except
logging.warning(
'Failed to delete failed instance %s: %s', self.instance_id, e
)

old_instance_id = self.instance_id
self.instance_id = f'{self.instance_id}-fb'
self.storage_type = azure_disk.PREMIUM_STORAGE
logging.info(
'Falling back from %s (PremiumV2_LRS) to %s (Premium_LRS)',
old_instance_id,
self.instance_id,
)

def _CreateStandard(self) -> None:
"""Creates standard Azure database instance (MySQL or Postgres Premium)."""
ha_flag = ENABLE_HA if self.spec.high_availability else DISABLE_HA
cmd = [
azure.AZURE_PATH,
self.GetAzCommandForEngine(),
self.SERVER_TYPE,
'create',
'--resource-group',
self.resource_group.name,
'--name',
self.instance_id,
'--location',
self.region,
'--high-availability',
ha_flag,
'--admin-user',
self.spec.database_username,
'--admin-password',
self.spec.database_password,
'--storage-size',
str(self.spec.db_disk_spec.disk_size),
'--sku-name',
self.spec.db_spec.machine_type,
'--tier',
self.spec.db_tier,
'--version',
self.spec.engine_version,
'--yes',
]
if self.storage_type:
cmd.extend([
'--storage-type',
self.storage_type,
])
if (
self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL
and self.spec.db_disk_spec.provisioned_iops is not None
):
cmd.extend([
'--auto-scale-iops',
'Disabled',
])

is_postgres = self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES
is_premium_v2 = self.storage_type == azure_disk.PREMIUM_STORAGE_V2

if self.spec.db_disk_spec.provisioned_iops:
if not is_postgres or is_premium_v2:
cmd.extend([
'--iops',
str(self.spec.db_disk_spec.provisioned_iops),
])
if self.spec.db_disk_spec.provisioned_throughput:
if not is_postgres or is_premium_v2:
cmd.extend([
'--throughput',
str(self.spec.db_disk_spec.provisioned_throughput),
])

vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT)

def GetAzCommandForEngine(self) -> str:
engine = self.spec.engine
if engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES:
Expand Down
74 changes: 73 additions & 1 deletion tests/providers/azure/azure_flexible_server_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ def testCreatePostgres(self):
with self.subTest(name='StorageType'):
self.assertIn('--storage-type Premium_LRS', cmd)
with self.subTest(name='Iops'):
self.assertIn('--iops 1000', cmd)
self.assertNotIn('--iops', cmd)
with self.subTest(name='Metadata'):
self.assertEqual(metadata['disk_type'], 'Premium_LRS')

Expand Down Expand Up @@ -415,6 +415,78 @@ def testCreatePostgresPremiumV2NoAsyncHeader(self):
with self.assertRaises(errors.Resource.CreationError):
bm_spec.relational_db._Create()

def testCreatePostgresPremiumV2Fallback(self):
mock_cmd = self.MockIssueCommand({
'rest --method PUT': [(
'',
"'Azure-AsyncOperation': 'https://management.azure.com/async_op'",
0,
)],
'rest --method GET': [(
'{"status": "Failed", "error": {"code": "ServerDropping"}}',
'',
0,
)],
'postgres flexible-server delete': [('', '', 0)],
'postgres flexible-server create': [('', '', 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()
db = bm_spec.relational_db
db.resource_group.name = 'az_resource'

db._Create()

self.assertEqual(mock_cmd.func_to_mock.call_count, 4)
put_call, get_call, delete_call, create_call = (
mock_cmd.func_to_mock.call_args_list
)

put_cmd = ' '.join(put_call[0][0])
self.assertIn('rest --method PUT', put_cmd)
self.assertIn('PremiumV2_LRS', put_cmd)

get_cmd = ' '.join(get_call[0][0])
self.assertIn('rest --method GET', get_cmd)

delete_cmd = ' '.join(delete_call[0][0])
self.assertIn('postgres flexible-server delete', delete_cmd)
self.assertIn('--name pkb-db-instance-test_run_uri', delete_cmd)

create_cmd = ' '.join(create_call[0][0])
self.assertIn('postgres flexible-server create', create_cmd)
self.assertIn('--name pkb-db-instance-test_run_uri-fb', create_cmd)
self.assertIn('--storage-type Premium_LRS', create_cmd)
self.assertNotIn('--iops', create_cmd)
self.assertNotIn('--throughput', create_cmd)

self.assertEqual(db.instance_id, 'pkb-db-instance-test_run_uri-fb')
self.assertEqual(db.storage_type, 'Premium_LRS')


if __name__ == '__main__':
unittest.main()
Loading