[MySQL Flexible Server] Add --maintenance-batch to flexible-server update

src/azure-cli/azure/cli/command_modules/mysql/_help.py

@@ -332,6 +332,8 @@
   - name: Update a flexible server's tags.
     text: az mysql flexible-server update --resource-group testGroup --name testserver --tags "k1=v1" "k2=v2"
     crafted: true
+  - name: Update the batch of the custom-managed maintenance window (existing batch is preserved when --maintenance-batch is omitted).
+    text: az mysql flexible-server update --resource-group testGroup --name testserver --maintenance-window "Fri:13:00" --maintenance-batch Batch2
   - name: Set or change key and identity for data encryption.
     text: >
       # get key identifier of the existing key

src/azure-cli/azure/cli/command_modules/mysql/_params.py

@@ -140,6 +140,13 @@ def load_arguments(self, _):    # pylint: disable=too-many-statements, too-many-
         help='The patch strategy of maintenance policy. Accepted values: Regular, VirtualCanary. Default value is Regular.'
     )
 
+    maintenance_batch_arg_type = CLIArgumentType(
+        arg_type=get_enum_type(['Default', 'Batch1', 'Batch2']),
+        options_list=['--maintenance-batch'],
+        help='The batch of the custom-managed maintenance window. Accepted values: Default, Batch1, Batch2. '
+             'Only valid with an enabled --maintenance-window; if omitted, the existing batch is preserved.'
+    )
+
     yes_arg_type = CLIArgumentType(
         options_list=['--yes', '-y'],
         action='store_true',
@@ -464,6 +471,7 @@ def load_arguments(self, _):    # pylint: disable=too-many-statements, too-many-
         c.argument('administrator_login_password', arg_type=administrator_login_password_arg_type)
         c.argument('maintenance_window', options_list=['--maintenance-window'], validator=maintenance_window_validator,
                    help='Period of time (UTC) designated for maintenance. Examples: "Sun:23:30" to schedule on Sunday, 11:30pm UTC. To set back to default pass in "Disabled".')
+        c.argument('maintenance_batch', arg_type=maintenance_batch_arg_type)
         c.argument('tags', tags_type)
         c.argument('tier', arg_type=tier_arg_type)
         c.argument('sku_name', arg_type=sku_name_arg_type)

src/azure-cli/azure/cli/command_modules/mysql/custom.py

@@ -1006,7 +1006,8 @@ def flexible_server_update_custom_func(cmd, client, instance, sku_name=None, tie
                                        high_availability=None, standby_availability_zone=None, maintenance_window=None,
                                        tags=None, replication_role=None, byok_identity=None, backup_byok_identity=None,
                                        byok_key=None, backup_byok_key=None, disable_data_encryption=False,
-                                       public_access=None, maintenance_policy_patch_strategy=None, backup_interval=None):
+                                       public_access=None, maintenance_policy_patch_strategy=None, backup_interval=None,
+                                       maintenance_batch=None):
     # validator
     location = ''.join(instance.location.lower().split())
     db_context = DbContext(
@@ -1064,14 +1065,22 @@ def flexible_server_update_custom_func(cmd, client, instance, sku_name=None, tie
     if backup_interval:
         instance.backup.backup_interval_hours = backup_interval
 
+    if maintenance_batch and not maintenance_window:
+        raise CLIError('--maintenance-batch can only be used together with an enabled --maintenance-window.')
+
     if maintenance_window:
         # if disabled is pass in reset to default values
         if maintenance_window.lower() == "disabled":
+            if maintenance_batch:
+                raise CLIError('--maintenance-batch cannot be used when disabling the maintenance window.')
             day_of_week = start_hour = start_minute = 0
             custom_window = "Disabled"
+            instance.maintenance_window.batch_of_maintenance = None
         else:
             day_of_week, start_hour, start_minute = parse_maintenance_window(maintenance_window)
             custom_window = "Enabled"
+            if maintenance_batch:
+                instance.maintenance_window.batch_of_maintenance = maintenance_batch
 
         # set values - if maintenance_window when is None when created then create a new object
         instance.maintenance_window.day_of_week = day_of_week

src/azure-cli/azure/cli/command_modules/mysql/tests/latest/test_mysql_scenario.py

@@ -2301,6 +2301,53 @@ def _test_maintenance_mgmt(self, database_engine, resource_group):
         self.assertEqual(reschedule_start_time, maintenance_rescheduled_time)
 
 
+class FlexibleServerMaintenanceBatchScenarioTest(ScenarioTest):
+
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(location=DEFAULT_LOCATION)
+    def test_mysql_flexible_server_maintenance_batch_mgmt(self, resource_group):
+        self._test_maintenance_batch_mgmt('mysql', resource_group)
+
+    def _test_maintenance_batch_mgmt(self, database_engine, resource_group):
+        server_name = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+
+        self.cmd('{} flexible-server create -g {} -n {} --public-access none --tier GeneralPurpose --sku-name {}'
+                 .format(database_engine, resource_group, server_name, DEFAULT_GENERAL_PURPOSE_SKU))
+
+        # set a custom maintenance window together with a batch
+        self.cmd('{} flexible-server update -g {} -n {} --maintenance-window "Fri:13:00" --maintenance-batch Batch1'
+                 .format(database_engine, resource_group, server_name),
+                 checks=[
+                     JMESPathCheck('maintenanceWindow.customWindow', 'Enabled'),
+                     JMESPathCheck('maintenanceWindow.dayOfWeek', 5),
+                     JMESPathCheck('maintenanceWindow.startHour', 13),
+                     JMESPathCheck('maintenanceWindow.batchOfMaintenance', 'Batch1')])
+
+        # update the window only (no --maintenance-batch): the existing batch must be preserved
+        self.cmd('{} flexible-server update -g {} -n {} --maintenance-window "Sat:14:00"'
+                 .format(database_engine, resource_group, server_name),
+                 checks=[
+                     JMESPathCheck('maintenanceWindow.dayOfWeek', 6),
+                     JMESPathCheck('maintenanceWindow.startHour', 14),
+                     JMESPathCheck('maintenanceWindow.batchOfMaintenance', 'Batch1')])
+
+        # change the batch explicitly
+        self.cmd('{} flexible-server update -g {} -n {} --maintenance-window "Sat:14:00" --maintenance-batch Batch2'
+                 .format(database_engine, resource_group, server_name),
+                 checks=[JMESPathCheck('maintenanceWindow.batchOfMaintenance', 'Batch2')])
+
+        # guardrail: --maintenance-batch without --maintenance-window is rejected
+        self.cmd('{} flexible-server update -g {} -n {} --maintenance-batch Batch2'
+                 .format(database_engine, resource_group, server_name), expect_failure=True)
+
+        # guardrail: --maintenance-batch cannot be combined with disabling the window
+        self.cmd('{} flexible-server update -g {} -n {} --maintenance-window "Disabled" --maintenance-batch Batch2'
+                 .format(database_engine, resource_group, server_name), expect_failure=True)
+
+        self.cmd('{} flexible-server delete -g {} -n {} --yes'
+                 .format(database_engine, resource_group, server_name))
+
+
 class MySQLExportTest(ScenarioTest):
     profile = None