Add nodepool rollback and get-rollback-versions commands

Author: indusridhar

src/aks-preview/HISTORY.rst

@@ -11,6 +11,8 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+* `az aks nodepool get-rollback-versions`: New command to get available rollback versions for an agent pool.
+* `az aks nodepool rollback`: New command to rollback an agent pool to the most recently used configuration (N-1).
 
 19.0.0b20
 +++++++

src/aks-preview/azext_aks_preview/_format.py

@@ -153,6 +153,18 @@ def aks_agentpool_list_table_format(results):
     return [_aks_agentpool_table_format(r) for r in results]
 
 
+def aks_agentpool_rollback_versions_table_format(results):
+    """Format rollback versions for display with "-o table"."""
+    if not results:
+        return []
+
+    def _format_rollback_version(result):
+        parsed = compile_jmes("""{\n            kubernetesVersion: orchestrator_version,\n            nodeImageVersion: node_image_version,\n            timestamp: timestamp\n        }""")
+        return parsed.search(result, Options(dict_cls=OrderedDict))
+
+    return [_format_rollback_version(r) for r in results]
+
+
 def aks_list_table_format(results):
     """"Format a list of managed clusters as summary results for display with "-o table"."""
     return [_aks_table_format(r) for r in results]

src/aks-preview/azext_aks_preview/_help.py

@@ -2466,6 +2466,40 @@
     crafted: true
 """
 
+helps['aks nodepool get-rollback-versions'] = """
+type: command
+short-summary: Get the available rollback versions for an agent pool of the managed Kubernetes cluster.
+long-summary: |
+    Get the list of historically used Kubernetes and node image versions that can be used for rollback operations.
+examples:
+  - name: Get the available rollback versions for an agent pool.
+    text: az aks nodepool get-rollback-versions --resource-group MyResourceGroup --cluster-name MyManagedCluster --nodepool-name MyNodePool
+    crafted: true
+"""
+
+helps['aks nodepool rollback'] = """
+type: command
+short-summary: Rollback an agent pool to the most recently used configuration (N-1).
+long-summary: |
+    Rollback an agent pool to the most recently used version based on rollback history.
+    This will rollback both the Kubernetes version and node image version to their most recent previous state.
+    For downgrades to older versions (N-2 or earlier), use a separate downgrade operation.
+parameters:
+  - name: --aks-custom-headers
+    type: string
+    short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2.
+  - name: --if-match
+    type: string
+    short-summary: The revision of the resource being updated. This should match the current revision.
+  - name: --if-none-match
+    type: string
+    short-summary: Set to '*' to allow a new resource to be created, but to prevent updating an existing resource.
+examples:
+  - name: Rollback a nodepool to the most recently used version.
+    text: az aks nodepool rollback --resource-group MyResourceGroup --cluster-name MyManagedCluster --nodepool-name MyNodePool
+    crafted: true
+"""
+
 helps['aks nodepool stop'] = """
     type: command
     short-summary: Stop running agent pool in the managed Kubernetes cluster.

src/aks-preview/azext_aks_preview/_params.py

@@ -2199,6 +2199,14 @@ def load_arguments(self, _):
     with self.argument_context("aks nodepool manual-scale delete") as c:
         c.argument("current_vm_sizes", is_preview=True)
 
+    with self.argument_context("aks nodepool get-rollback-versions") as c:
+        pass  # Uses common nodepool parameters
+
+    with self.argument_context("aks nodepool rollback") as c:
+        c.argument("aks_custom_headers", nargs="*")
+        c.argument("if_match")
+        c.argument("if_none_match")
+
     with self.argument_context("aks machine") as c:
         c.argument("cluster_name", help="The cluster name.")
         c.argument(

src/aks-preview/azext_aks_preview/commands.py

@@ -26,6 +26,7 @@
     aks_addon_show_table_format,
     aks_agentpool_list_table_format,
     aks_agentpool_show_table_format,
+    aks_agentpool_rollback_versions_table_format,
     aks_machine_list_table_format,
     aks_machine_show_table_format,
     aks_operation_show_table_format,
@@ -274,6 +275,12 @@ def load_command_table(self, _):
         g.custom_command("update", "aks_agentpool_update", supports_no_wait=True)
         g.custom_command("delete", "aks_agentpool_delete", supports_no_wait=True)
         g.custom_command("get-upgrades", "aks_agentpool_get_upgrade_profile")
+        g.custom_command(
+            "get-rollback-versions",
+            "aks_agentpool_get_rollback_versions",
+            table_transformer=aks_agentpool_rollback_versions_table_format
+        )
+        g.custom_command("rollback", "aks_agentpool_rollback", supports_no_wait=True)
         g.custom_command("stop", "aks_agentpool_stop", supports_no_wait=True)
         g.custom_command("start", "aks_agentpool_start", supports_no_wait=True)
         g.custom_command(

src/aks-preview/azext_aks_preview/custom.py

@@ -2222,6 +2222,114 @@ def aks_agentpool_get_upgrade_profile(cmd,   # pylint: disable=unused-argument
     return client.get_upgrade_profile(resource_group_name, cluster_name, nodepool_name)
 
 
+def aks_agentpool_get_rollback_versions(cmd,   # pylint: disable=unused-argument
+                                        client,
+                                        resource_group_name,
+                                        cluster_name,
+                                        nodepool_name):
+    """Get rollback versions for a nodepool."""
+    upgrade_profile = client.get_upgrade_profile(resource_group_name, cluster_name, nodepool_name)
+    return upgrade_profile.recently_used_versions
+
+
+def aks_agentpool_rollback(cmd,   # pylint: disable=unused-argument
+                           client,
+                           resource_group_name,
+                           cluster_name,
+                           nodepool_name,
+                           aks_custom_headers=None,
+                           if_match=None,
+                           if_none_match=None,
+                           no_wait=False):
+    """Rollback a nodepool to the most recent previous version configuration."""
+    from azext_aks_preview._client_factory import cf_managed_clusters
+
+    # Warn users when auto-upgrade is enabled
+    if cmd and getattr(cmd, "cli_ctx", None):
+        try:
+            managed_clusters_client = cf_managed_clusters(cmd.cli_ctx)
+            managed_cluster = managed_clusters_client.get(resource_group_name, cluster_name)
+            auto_upgrade_profile = getattr(managed_cluster, "auto_upgrade_profile", None)
+
+            upgrade_channel = getattr(auto_upgrade_profile, "upgrade_channel", None) if auto_upgrade_profile else None
+            node_os_upgrade_channel = (
+                getattr(auto_upgrade_profile, "node_os_upgrade_channel", None)
+                if auto_upgrade_profile
+                else None
+            )
+
+            upgrade_channel_enabled = upgrade_channel and str(upgrade_channel).lower() != "none"
+            node_os_channel_enabled = node_os_upgrade_channel and str(node_os_upgrade_channel).lower() not in [
+                "none",
+                "unmanaged",
+            ]
+
+            if upgrade_channel_enabled or node_os_channel_enabled:
+                logger.warning(
+                    "Auto-upgrade is enabled on cluster '%s' (upgradeChannel=%s, nodeOSUpgradeChannel=%s). "
+                    "Rollback will not succeed until auto-upgrade is disabled. Please disable auto-upgrade to roll back the node pool.",
+                    cluster_name,
+                    upgrade_channel or "none",
+                    node_os_upgrade_channel or "Unmanaged",
+                )
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.debug("Unable to retrieve auto-upgrade configuration before rollback: %s", ex)
+
+    logger.info("Fetching the most recent rollback version...")
+
+    # Get upgrade profile to retrieve recently used versions
+    upgrade_profile = client.get_upgrade_profile(resource_group_name, cluster_name, nodepool_name)
+
+    if not upgrade_profile.recently_used_versions or len(upgrade_profile.recently_used_versions) == 0:
+        raise CLIError(
+            "No rollback versions available. The nodepool must have been upgraded at least once "
+            "to have rollback history available."
+        )
+
+    # Sort by timestamp (most recent first) and get the most recent version
+    sorted_versions = sorted(
+        upgrade_profile.recently_used_versions,
+        key=lambda v: v.timestamp if v.timestamp else datetime.datetime.min,
+        reverse=True
+    )
+    most_recent = sorted_versions[0]
+
+    kubernetes_version = most_recent.orchestrator_version
+    node_image_version = most_recent.node_image_version
+
+    logger.info(
+        "Rolling back to the most recent version: "
+        "Kubernetes version: %s, Node image version: %s (timestamp: %s)",
+        kubernetes_version, node_image_version, most_recent.timestamp
+    )
+
+    # Get the current agent pool
+    current_agentpool = client.get(resource_group_name, cluster_name, nodepool_name)
+
+    # Update the agent pool configuration with rollback versions
+    current_agentpool.orchestrator_version = kubernetes_version
+    current_agentpool.node_image_version = node_image_version
+
+    # Set custom headers if provided
+    headers = get_aks_custom_headers(aks_custom_headers)
+    if if_match:
+        headers['If-Match'] = if_match
+    if if_none_match:
+        headers['If-None-Match'] = if_none_match
+
+    # Perform the rollback by updating the agent pool
+    # Server-side will validate the versions
+    return sdk_no_wait(
+        no_wait,
+        client.begin_create_or_update,
+        resource_group_name,
+        cluster_name,
+        nodepool_name,
+        current_agentpool,
+        headers=headers if headers else None
+    )
+
+
 def aks_agentpool_stop(cmd,   # pylint: disable=unused-argument
                        client,
                        resource_group_name,