[MySQL] Fix #32217: az mysql flexible-server restore: Fix --no-wait flag (#32744)

Author: a0x1ab

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

@@ -843,19 +843,29 @@ def flexible_server_restore(cmd, client, resource_group_name, server_name, sourc
         else:
             parameters.network = source_server_object.network
 
-    except Exception as e:
-        raise ResourceNotFoundError(e)
-
-    resolve_poller(
-        client.begin_create(resource_group_name, server_name, parameters), cmd.cli_ctx,
-        'Restore Server')
+    except HttpResponseError as exc:
+        raise ResourceNotFoundError(exc) from exc
+
+    def _begin_network_update():
+        restore_server_object = client.get(resource_group_name, server_name)
+        restore_server_network = restore_server_object.network
+        restore_server_network.public_network_access = public_access if public_access else source_server_object.network.public_network_access
+        update_parameter = models.ServerForUpdate(network=restore_server_network)
+        return client.begin_update(resource_group_name, server_name, update_parameter)
+
+    create_poller = sdk_no_wait(no_wait, client.begin_create, resource_group_name, server_name, parameters)
+    if no_wait:
+        def _post_create_update(poller):
+            try:
+                _begin_network_update()
+            except (HttpResponseError, CLIError) as ex:
+                logger.warning('Skipping post-restore network update: %s', ex)
 
-    restore_server_object = client.get(resource_group_name, server_name)
-    restore_server_network = restore_server_object.network
-    restore_server_network.public_network_access = public_access if public_access else source_server_object.network.public_network_access
-    update_parameter = models.ServerForUpdate(network=restore_server_network)
+        create_poller.add_done_callback(_post_create_update)
+        return create_poller
 
-    return sdk_no_wait(no_wait, client.begin_update, resource_group_name, server_name, update_parameter)
+    resolve_poller(create_poller, cmd.cli_ctx, 'Restore Server')
+    return sdk_no_wait(no_wait, _begin_network_update)
 
 
 # pylint: disable=too-many-locals, too-many-statements, raise-missing-from

src/azure-cli/azure/cli/command_modules/mysql/tests/latest/recordings/test_mysql_flexible_server_restore_no_wait_mgmt.yaml

@@ -120,6 +120,12 @@ def test_mysql_flexible_server_georestore_mgmt(self, resource_group):
     def test_mysql_flexible_server_georestore_update_mgmt(self, resource_group):
         self._test_flexible_server_georestore_update_mgmt('mysql', resource_group)
 
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(location=DEFAULT_LOCATION)
+    @live_only()
+    def test_mysql_flexible_server_restore_no_wait_mgmt(self, resource_group):
+        self._test_flexible_server_restore_no_wait('mysql', resource_group)
+
     @AllowLargeResponse()
     @ResourceGroupPreparer(location=DEFAULT_LOCATION)
     def test_mysql_flexible_server_gtid_reset(self, resource_group):
@@ -698,6 +704,23 @@ def _test_flexible_server_georestore_update_mgmt(self, database_engine, resource
         self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, source_server))
         self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, target_server))
 
+    def _test_flexible_server_restore_no_wait(self, database_engine, resource_group):
+        location = DEFAULT_LOCATION
+        source_server = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+        target_server = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+
+        self.cmd('{} flexible-server create -g {} -n {} -l {} --public-access None --tier GeneralPurpose --sku-name {}'
+                 .format(database_engine, resource_group, source_server, location, DEFAULT_GENERAL_PURPOSE_SKU))
+
+        self.cmd('{} flexible-server restore -g {} --name {} --source-server {} --no-wait'
+                 .format(database_engine, resource_group, target_server, source_server))
+
+        self.cmd('{} flexible-server wait -g {} -n {} --created --interval 30 --timeout 600'
+                 .format(database_engine, resource_group, target_server))
+
+        self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, source_server))
+        self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, target_server))
+
 
     def _test_flexible_server_byok_mgmt(self, database_engine, resource_group, vault_name, backup_vault_name=None):
         key_name = self.create_random_name('rdbmskey', 32)