[POSTGRESQL] az postgres flexible-server upgrade: Introduced --validate-only param for PVC

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

@@ -121,6 +121,10 @@ def cf_postgres_flexible_maintenance_events(cli_ctx, _):
     return get_postgresql_flexible_management_client(cli_ctx).maintenance_events
 
 
+def cf_postgres_flexible_major_version_upgrade_precheck(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).major_version_upgrade_precheck
+
+
 def resource_client_factory(cli_ctx, subscription_id=None):
     return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id)
 

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

@@ -937,6 +937,8 @@
 examples:
   - name: Upgrade server 'testserver' to PostgreSQL major version 18.
     text: az postgres flexible-server upgrade -g testgroup -n testserver -v 18
+  - name: Run pre-upgrade validation for server targeting a later PostgreSQL major version without performing the upgrade.
+    text: az postgres flexible-server upgrade -g testgroup -n testserver -v 18 --validate-only
 """
 
 helps['postgres flexible-server identity'] = """

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

@@ -268,6 +268,13 @@ def _flexible_server_params(command_group):
             help='Server major version.'
         )
 
+        pg_version_upgrade_validate_arg_type = CLIArgumentType(
+            options_list=['--validate-only'],
+            action='store_true',
+            help='Run a pre-upgrade validation against the server without performing major version upgrade on resource. '
+                 'Returns the prevalidation check result.'
+        )
+
         private_dns_zone_arguments_arg_type = CLIArgumentType(
             options_list=['--private-dns-zone'],
             help='This parameter only applies for a server with private access and is required when using --vnet or --subnet. '
@@ -532,6 +539,7 @@ def _flexible_server_params(command_group):
 
         with self.argument_context('{} flexible-server upgrade'.format(command_group)) as c:
             c.argument('version', arg_type=pg_version_upgrade_arg_type)
+            c.argument('validate', arg_type=pg_version_upgrade_validate_arg_type)
             c.argument('yes', arg_type=yes_arg_type)
 
         with self.argument_context('{} flexible-server restart'.format(command_group)) as c:

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

@@ -4,18 +4,24 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=unused-argument, line-too-long
-from azure.cli.core.util import user_confirmation
+from knack.log import get_logger
 from knack.util import CLIError
-from .._client_factory import cf_postgres_flexible_replica
+from azure.cli.core.util import user_confirmation
+from .._client_factory import cf_postgres_flexible_replica, cf_postgres_flexible_major_version_upgrade_precheck
 from ..utils._flexible_server_location_capabilities_util import get_postgres_server_capability_info
 from ..utils._flexible_server_util import resolve_poller
 from ..utils.validators import pg_version_validator, validate_citus_cluster, validate_resource_group
 
+logger = get_logger(__name__)
+
 
-def flexible_server_version_upgrade(cmd, client, resource_group_name, server_name, version, yes=None):
+def flexible_server_version_upgrade(cmd, client, resource_group_name, server_name, version, validate=None, yes=None):
     validate_resource_group(resource_group_name)
     validate_citus_cluster(cmd, resource_group_name, server_name)
 
+    if validate:
+        return _flexible_server_version_upgrade_validate(cmd, client, resource_group_name, server_name, version)
+
     if not yes:
         user_confirmation(
             "Upgrading major version in server {} is irreversible. The action you're about to take can't be undone. "
@@ -67,3 +73,51 @@ def flexible_server_version_upgrade(cmd, client, resource_group_name, server_nam
             parameters=parameters),
         cmd.cli_ctx, 'Upgrading server {} to major version {}'.format(server_name, version)
     )
+
+
+def _flexible_server_version_upgrade_validate(cmd, client, resource_group_name, server_name, version):
+    body = {
+        'targetVersion': version
+    }
+
+    start_response = resolve_poller(
+        client.begin_start_major_version_upgrade_precheck(
+            resource_group_name=resource_group_name,
+            server_name=server_name,
+            body=body),
+        cmd.cli_ctx,
+        'Starting major version upgrade precheck for server {} targeting version {}'.format(server_name, version)
+    )
+
+    precheck_validation_id = _get_attr_or_key(start_response, 'name')
+    if not precheck_validation_id:
+        raise CLIError('Failed to retrieve precheck validation id from the upgrade precheck response.')
+
+    precheck_client = cf_postgres_flexible_major_version_upgrade_precheck(cmd.cli_ctx, '_')
+
+    return precheck_client.get(
+        resource_group_name=resource_group_name,
+        server_name=server_name,
+        precheck_validation_id=precheck_validation_id)
+
+
+def _get_attr_or_key(obj, name):
+    if obj is None:
+        return None
+    value = getattr(obj, name, None)
+    if value is not None:
+        return value
+    if isinstance(obj, dict):
+        return obj.get(name)
+    properties = getattr(obj, 'properties', None)
+    if properties is not None:
+        return _get_attr_or_key(properties, name)
+    return None
+
+
+def _get_status(obj):
+    status = _get_attr_or_key(obj, 'status')
+    if status is None:
+        return None
+    # Some SDK enums expose .value
+    return getattr(status, 'value', status)

src/azure-cli/azure/cli/command_modules/postgresql/tests/latest/test_postgres_flexible_commands_upgrade.py

@@ -6,10 +6,13 @@
 from azure.cli.testsdk.scenario_tests import AllowLargeResponse
 from azure.cli.testsdk import (
     JMESPathCheck,
+    JMESPathCheckExists,
     ResourceGroupPreparer,
     ScenarioTest)
 from .constants import SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH, DEFAULT_LOCATION
 
+PRECHECK_LOCATION = 'eastus2euap'
+
 
 class PostgreSQLFlexibleServerUpgradeMgmtScenarioTest(ScenarioTest):
 
@@ -35,6 +38,43 @@ def _test_flexible_server_upgrade_mgmt(self, resource_group):
         result = self.cmd('postgres flexible-server upgrade -g {} -n {} --version {} --yes'.format(resource_group, server_name, new_version)).get_output_in_json()
         self.assertTrue(result['version'].startswith(new_version))
 
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(location=PRECHECK_LOCATION)
+    def test_postgres_flexible_server_upgrade_validate_only_mgmt(self, resource_group):
+        self._test_flexible_server_upgrade_validate_only_mgmt(resource_group)
+
+    def _test_flexible_server_upgrade_validate_only_mgmt(self, resource_group):
+        server_name = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+        current_version = '15'
+        target_version = '18'
+        location = PRECHECK_LOCATION
+
+        # Create server in the precheck-enabled region
+        self.cmd('postgres flexible-server create -g {} -n {} --tier GeneralPurpose --location {} '
+                 '--version {} --public-access none --yes'.format(
+                     resource_group, server_name, location, current_version))
+
+        self.cmd('postgres flexible-server show -g {} -n {}'.format(resource_group, server_name),
+                 checks=[JMESPathCheck('version', current_version),
+                         JMESPathCheck('location', 'East US 2 EUAP')])
+
+        # Run pre-upgrade validation. With --validate-only we poll until the
+        # precheck reaches a terminal status, so the returned status must be
+        # one of Succeeded / Failed / Canceled (never still Validating).
+        result = self.cmd(
+            'postgres flexible-server upgrade -g {} -n {} --version {} --validate-only'.format(
+                resource_group, server_name, target_version),
+            checks=[
+                JMESPathCheckExists('name'),
+                JMESPathCheckExists('properties.status'),
+            ]).get_output_in_json()
+
+        self.assertIn(result['properties']['status'], ['Succeeded', 'Failed', 'Canceled'])
+
+        # Server version must not change as a result of the precheck.
+        self.cmd('postgres flexible-server show -g {} -n {}'.format(resource_group, server_name),
+                 checks=[JMESPathCheck('version', current_version)])
+
     @AllowLargeResponse()
     @ResourceGroupPreparer(location=DEFAULT_LOCATION)
     def test_postgres_flexible_server_create_premiumv2_lrs_version_below_14(self, resource_group):