Skip to content

Commit e04c009

Browse files
bvliucopybara-github
bvliu
authored and
copybara-github
committed
Add support for specifying storage type, IOPS, and throughput for Azure SQL databases.
PiperOrigin-RevId: 846591475
1 parent 33d4fe7 commit e04c009

File tree

2 files changed

+200
-2
lines changed

2 files changed

+200
-2
lines changed

‎perfkitbenchmarker/providers/azure/azure_flexible_server.py‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,27 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14-
"""Azure Flexible server provisioning and teardown."""
14+
"""Azure Flexible server provisioning and teardown.
15+
16+
Postgres Flexible Server is created with PremiumV2_LRS storage type.
17+
18+
MySQL Flexible Server is created with provisioned IOPS unless omitted, in which
19+
case it is created with auto-scale IOPS.
20+
"""
1521

1622
import datetime
1723
import logging
1824
import time
19-
from typing import Tuple
25+
from typing import Any, Tuple
2026

2127
from absl import flags
28+
from perfkitbenchmarker import errors
2229
from perfkitbenchmarker import provider_info
2330
from perfkitbenchmarker import relational_db
2431
from perfkitbenchmarker import sql_engine_utils
2532
from perfkitbenchmarker import vm_util
2633
from perfkitbenchmarker.providers import azure
34+
from perfkitbenchmarker.providers.azure import azure_disk
2735
from perfkitbenchmarker.providers.azure import azure_relational_db
2836

2937
DEFAULT_DATABASE_NAME = 'database'
@@ -63,6 +71,25 @@ class AzureFlexibleServer(azure_relational_db.AzureRelationalDb):
6371
sql_engine_utils.FLEXIBLE_SERVER_MYSQL,
6472
]
6573

74+
def __init__(self, relational_db_spec: Any):
75+
super().__init__(relational_db_spec)
76+
if (
77+
self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL
78+
and self.spec.db_disk_spec.provisioned_throughput
79+
):
80+
raise errors.Config.InvalidValue(
81+
'Provisioned throughput is not supported for MySQL Flexible Server.'
82+
)
83+
self.storage_type = None
84+
if self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES:
85+
self.storage_type = (
86+
self.spec.db_disk_spec.disk_type or azure_disk.PREMIUM_STORAGE
87+
)
88+
if self.storage_type == azure_disk.PREMIUM_STORAGE_V2:
89+
raise errors.Config.InvalidValue(
90+
'Premium Storage v2 is not supported for Postgres Flexible Server.'
91+
)
92+
6693
@staticmethod
6794
def GetDefaultEngineVersion(engine: str) -> str:
6895
"""Returns the default version of a given database engine."""
@@ -75,9 +102,22 @@ def GetDefaultEngineVersion(engine: str) -> str:
75102
'Unsupported engine {}'.format(engine)
76103
)
77104

105+
def GetResourceMetadata(self) -> dict[str, Any]:
106+
metadata = super().GetResourceMetadata()
107+
if self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL:
108+
metadata['autoscale_iops'] = (
109+
'Enabled'
110+
if self.spec.db_disk_spec.provisioned_iops is None
111+
else 'Disabled'
112+
)
113+
if self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_POSTGRES:
114+
metadata['disk_type'] = self.storage_type
115+
return metadata
116+
78117
def _Create(self) -> None:
79118
"""Creates the Azure database instance."""
80119
ha_flag = ENABLE_HA if self.spec.high_availability else DISABLE_HA
120+
# Consider adding --zone and maybe --standby-zone for parity with GCP.
81121
cmd = [
82122
azure.AZURE_PATH,
83123
self.GetAzCommandForEngine(),
@@ -105,6 +145,29 @@ def _Create(self) -> None:
105145
self.spec.engine_version,
106146
'--yes',
107147
]
148+
if self.storage_type:
149+
cmd.extend([
150+
'--storage-type',
151+
self.storage_type,
152+
])
153+
if (
154+
self.spec.engine == sql_engine_utils.FLEXIBLE_SERVER_MYSQL
155+
and self.spec.db_disk_spec.provisioned_iops is not None
156+
):
157+
cmd.extend([
158+
'--auto-scale-iops',
159+
'Disabled',
160+
])
161+
if self.spec.db_disk_spec.provisioned_iops:
162+
cmd.extend([
163+
'--iops',
164+
str(self.spec.db_disk_spec.provisioned_iops),
165+
])
166+
if self.spec.db_disk_spec.provisioned_throughput:
167+
cmd.extend([
168+
'--throughput',
169+
str(self.spec.db_disk_spec.provisioned_throughput),
170+
])
108171

109172
vm_util.IssueCommand(cmd, timeout=CREATE_AZURE_DB_TIMEOUT)
110173

‎tests/providers/azure/azure_flexible_server_test.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
import inspect
2+
import unittest
3+
4+
from absl import flags
5+
import mock
6+
from perfkitbenchmarker import errors
7+
from perfkitbenchmarker import provider_info
8+
from perfkitbenchmarker import sql_engine_utils
9+
from perfkitbenchmarker import vm_util
10+
from perfkitbenchmarker.providers.azure import azure_network
11+
from tests import pkb_common_test_case
12+
13+
FLAGS = flags.FLAGS
14+
15+
16+
class AzureFlexibleServerCreateTestCase(pkb_common_test_case.PkbCommonTestCase):
17+
18+
def setUp(self):
19+
super().setUp()
20+
FLAGS.cloud = provider_info.AZURE
21+
FLAGS.run_uri = 'test_run_uri'
22+
mock_resource_group = mock.Mock()
23+
self.resource_group_patch = self.enter_context(
24+
mock.patch.object(azure_network, 'GetResourceGroup')
25+
)
26+
self.resource_group_patch.return_value = mock_resource_group
27+
mock_resource_group.name = 'az_resource'
28+
self.mock_command = self.enter_context(
29+
mock.patch.object(vm_util, 'IssueCommand')
30+
)
31+
self.mock_command.return_value = (None, '', None)
32+
33+
def testCreatePostgres(self):
34+
yaml_spec = inspect.cleandoc(f"""
35+
sysbench:
36+
relational_db:
37+
cloud: {provider_info.AZURE}
38+
engine: {sql_engine_utils.FLEXIBLE_SERVER_POSTGRES}
39+
engine_version: '13'
40+
db_tier: GeneralPurpose
41+
db_spec:
42+
{provider_info.AZURE}:
43+
machine_type: Standard_D2ds_v4
44+
zone: eastus
45+
db_disk_spec:
46+
{provider_info.AZURE}:
47+
disk_size: 128
48+
provisioned_iops: 1000
49+
vm_groups:
50+
clients:
51+
vm_spec: *default_dual_core
52+
disk_spec: *default_500_gb
53+
""")
54+
bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
55+
yaml_spec, 'sysbench'
56+
)
57+
bm_spec.ConstructRelationalDb()
58+
59+
bm_spec.relational_db._Create()
60+
metadata = bm_spec.relational_db.GetResourceMetadata()
61+
62+
cmd = ' '.join(self.mock_command.call_args[0][0])
63+
with self.subTest(name='StorageType'):
64+
self.assertIn('--storage-type Premium_LRS', cmd)
65+
with self.subTest(name='Iops'):
66+
self.assertIn('--iops 1000', cmd)
67+
with self.subTest(name='Metadata'):
68+
self.assertEqual(metadata['disk_type'], 'Premium_LRS')
69+
70+
def testCreateMysql(self):
71+
yaml_spec = inspect.cleandoc(f"""
72+
sysbench:
73+
relational_db:
74+
cloud: {provider_info.AZURE}
75+
engine: {sql_engine_utils.FLEXIBLE_SERVER_MYSQL}
76+
engine_version: '8.0'
77+
db_tier: GeneralPurpose
78+
db_spec:
79+
{provider_info.AZURE}:
80+
machine_type: Standard_D2ds_v4
81+
zone: eastus
82+
db_disk_spec:
83+
{provider_info.AZURE}:
84+
disk_size: 128
85+
provisioned_iops: 1000
86+
vm_groups:
87+
clients:
88+
vm_spec: *default_dual_core
89+
disk_spec: *default_500_gb
90+
""")
91+
bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
92+
yaml_spec, 'sysbench'
93+
)
94+
bm_spec.ConstructRelationalDb()
95+
96+
bm_spec.relational_db._Create()
97+
metadata = bm_spec.relational_db.GetResourceMetadata()
98+
99+
cmd = ' '.join(self.mock_command.call_args[0][0])
100+
with self.subTest(name='AutoScaleIops'):
101+
self.assertIn('--auto-scale-iops Disabled', cmd)
102+
with self.subTest(name='Metadata'):
103+
self.assertEqual(metadata['autoscale_iops'], 'Disabled')
104+
105+
def testConstructMysqlWithThroughputRaises(self):
106+
yaml_spec = inspect.cleandoc(f"""
107+
sysbench:
108+
relational_db:
109+
cloud: {provider_info.AZURE}
110+
engine: {sql_engine_utils.FLEXIBLE_SERVER_MYSQL}
111+
engine_version: '8.0'
112+
db_tier: GeneralPurpose
113+
db_spec:
114+
{provider_info.AZURE}:
115+
machine_type: Standard_D2ds_v4
116+
zone: eastus
117+
db_disk_spec:
118+
{provider_info.AZURE}:
119+
disk_size: 128
120+
provisioned_iops: 1000
121+
provisioned_throughput: 200
122+
vm_groups:
123+
clients:
124+
vm_spec: *default_dual_core
125+
disk_spec: *default_500_gb
126+
""")
127+
bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
128+
yaml_spec, 'sysbench'
129+
)
130+
with self.assertRaises(errors.Config.InvalidValue):
131+
bm_spec.ConstructRelationalDb()
132+
133+
134+
if __name__ == '__main__':
135+
unittest.main()

0 commit comments

Comments
 (0)