[PostgreSQL] az postgres flexible-server restore: Add --sku-name and --tier arguments to allow changing compute during point-in-time restore

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

@@ -342,6 +342,10 @@
         --source-server /subscriptions/{sourceSubscriptionId}/resourceGroups/{sourceResourceGroup}/providers/Microsoft.DBforPostgreSQL/flexibleServers/{sourceServer}
   - name: Restore 'testserver' to current point-in-time as a new server 'testservernew' using Premium SSD v2 Disks by setting storage type to "PremiumV2_LRS"
     text: az postgres flexible-server restore --resource-group testgroup --name testservernew --source-server testserver --storage-type PremiumV2_LRS
+  - name: Restore 'testserver' to current point-in-time as a new server 'testservernew' with a different compute size.
+    text: az postgres flexible-server restore --resource-group testgroup --name testservernew --source-server testserver --sku-name Standard_D4s_v3
+  - name: Restore 'testserver' to current point-in-time as a new server 'testservernew' with a different compute tier and compute size.
+    text: az postgres flexible-server restore --resource-group testgroup --name testservernew --source-server testserver --tier MemoryOptimized --sku-name Standard_E2ds_v4
 """
 
 helps['postgres flexible-server maintenance-event'] = """

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

@@ -461,6 +461,8 @@ def _flexible_server_params(command_group):
         with self.argument_context('{} flexible-server restore'.format(command_group)) as c:
             c.argument('restore_point_in_time', arg_type=restore_point_in_time_arg_type)
             c.argument('source_server', arg_type=source_server_arg_type)
+            c.argument('sku_name', arg_type=sku_name_arg_type)
+            c.argument('tier', arg_type=tier_arg_type)
             c.argument('vnet', arg_type=vnet_arg_type)
             c.argument('subnet', arg_type=subnet_arg_type)
             c.argument('private_dns_zone_arguments', private_dns_zone_arguments_arg_type)

src/azure-cli/azure/cli/command_modules/postgresql/commands/custom_commands.py

@@ -41,6 +41,8 @@
     check_resource_group,
     pg_arguments_validator,
     pg_byok_validator,
+    pg_restore_sku_validator,
+    pg_restore_tier_validator,
     pg_restore_validator,
     resolve_private_dns_zone_id,
     validate_and_format_restore_point_in_time,
@@ -328,7 +330,7 @@ def flexible_server_restore(cmd, client,
                             private_dns_zone_arguments=None, geo_redundant_backup=None,
                             byok_identity=None, byok_key=None, backup_byok_identity=None, backup_byok_key=None,
                             federated_client_id=None, backup_federated_client_id=None,
-                            storage_type=None, yes=False):
+                            storage_type=None, sku_name=None, tier=None, yes=False):
 
     server_name = server_name.lower()
 
@@ -368,7 +370,23 @@ def flexible_server_restore(cmd, client,
                           federated_client_id=federated_client_id,
                           backup_federated_client_id=backup_federated_client_id)
 
-        pg_restore_validator(source_server_object.sku.tier, storage_type=storage_type)
+        if sku_name or tier:
+            sku_info = get_postgres_location_capability_info(cmd, location)['sku_info']
+            if tier:
+                pg_restore_tier_validator(tier, source_server_object.sku.tier, sku_info)
+            else:
+                tier = source_server_object.sku.tier
+            if sku_name:
+                pg_restore_sku_validator(sku_name, sku_info, tier)
+            else:
+                sku_name = get_postgres_default_sku(sku_info, tier)
+                logger.warning('--sku-name was not specified. The restored server will use the default compute '
+                               'size \'%s\' for the %s tier.', sku_name, tier)
+        else:
+            tier = source_server_object.sku.tier
+            sku_name = source_server_object.sku.name
+
+        pg_restore_validator(tier, storage_type=storage_type)
         storage = postgresql_flexibleservers.models.Storage(type=storage_type if source_server_object.storage.type != "PremiumV2_LRS" else None)
 
         parameters = postgresql_flexibleservers.models.Server(
@@ -377,6 +395,7 @@ def flexible_server_restore(cmd, client,
             source_server_resource_id=source_server_id,  # this should be the source server name, not id
             create_mode="PointInTimeRestore",
             availability_zone=zone,
+            sku=postgresql_flexibleservers.models.Sku(name=sku_name, tier=tier),
             storage=storage
         )
 
@@ -404,6 +423,9 @@ def flexible_server_restore(cmd, client,
                                                                                              federated_client_id=federated_client_id,
                                                                                              backup_federated_client_id=backup_federated_client_id)
 
+    # Let argument validation errors surface as-is instead of being masked as "not found".
+    except CLIError:
+        raise
     except Exception as e:
         raise ResourceNotFoundError(e)