[RDBMS] az postgres flexible-server upgrade: Add server capability API check to the --version parameter

Author: Matthew Boentoro

src/azure-cli/azure/cli/command_modules/rdbms/_flexible_server_location_capabilities_util.py

@@ -7,7 +7,7 @@
 from azure.cli.core.azclierror import InvalidArgumentValueError
 from azure.core.paging import ItemPaged
 from ._client_factory import cf_postgres_flexible_location_capabilities, cf_postgres_flexible_server_capabilities
-
+from collections import defaultdict
 
 def get_postgres_location_capability_info(cmd, location):
     list_location_capability_client = cf_postgres_flexible_location_capabilities(cmd.cli_ctx, '_')
@@ -98,13 +98,18 @@ def _postgres_parse_list_capability(result):
     versions = set()
     for version in result[0].supported_server_versions:
         versions.add(version.name)
+    
+    supported_server_versions = defaultdict(list)
+    for version in result[0].supported_server_versions:
+        supported_server_versions[version.name] = version.supported_versions_to_upgrade
 
     return {
         'sku_info': tiers_dict,
         'single_az': single_az,
         'geo_backup_supported': geo_backup_supported,
         'zones': zones,
         'server_versions': versions,
+        'supported_server_versions': supported_server_versions,
         'index_tuning_supported': index_tuning_supported
     }
 

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

@@ -459,7 +459,6 @@ def _flexible_server_params(command_group):
         )
 
         pg_version_upgrade_arg_type = CLIArgumentType(
-            arg_type=get_enum_type(['12', '13', '14', '15', '16']),
             options_list=['--version', '-v'],
             help='Server major version.'
         )

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_custom_common.py

@@ -16,6 +16,7 @@
 from ._client_factory import cf_postgres_flexible_replica
 from ._flexible_server_util import run_subprocess, \
     fill_action_template, get_git_root_dir, resolve_poller, GITHUB_ACTION_PATH
+from ._flexible_server_location_capabilities_util import get_postgres_server_capability_info
 from .validators import validate_public_access_server, validate_resource_group, check_resource_group, validate_citus_cluster
 
 logger = get_logger(__name__)
@@ -303,6 +304,20 @@ def flexible_server_version_upgrade(cmd, client, resource_group_name, server_nam
                        "the option to select version 12 will be removed in the near future. "
                        "We recommend selecting PostgreSQL 13 or a later version for "
                        "all future operations.")
+    
+    list_server_capability_info = get_postgres_server_capability_info(cmd, resource_group_name, server_name)
+    eligible_versions = list_server_capability_info['supported_server_versions'][str(current_version)]
+    if version not in eligible_versions:
+        error_message = ""
+        if len(eligible_versions) > 0:
+            error_message = "The version to upgrade to must be one of the supported versions for the current server version from this list: {}. ".format(eligible_versions)
+        else:
+            error_message = "Your server is up-to-date with the latest version. "
+
+        if "17" in eligible_versions:
+            error_message += "Note that version 17 is still in preview."
+        raise CLIError(error_message)
+    
 
     replica_operations_client = cf_postgres_flexible_replica(cmd.cli_ctx, '_')
     version_mapped = version