Rollback cli

Author: indusridhar

src/aks-preview/HISTORY.rst

@@ -9,6 +9,11 @@ If there is no rush to release a new version, please just add a description of t
 
 To release a new version, please select a new version number (usually plus 1 to last patch version, X.Y.Z -> Major.Minor.Patch, more details in `\doc <https://semver.org/>`_), and then add a new section named as the new version number in this file, the content should include the new modifications and everything from the *Pending* section. Finally, update the `VERSION` variable in `setup.py` with this new version number.
 
+18.0.0b45
++++++++
+* `az aks nodepool get-rollback-versions`: Add new command to get available rollback versions for a nodepool.
+* `az aks nodepool rollback`: Add new command to rollback a nodepool to a previously used configuration (N-1).
+
 Pending
 +++++++
 * `az aks update`: Add new parameter `--kms-infrastructure-encryption` to enable KMS infrastructure encryption on an existing cluster.

src/aks-preview/azext_aks_preview/_format.py

@@ -153,6 +153,22 @@ 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("""{
+            kubernetesVersion: orchestrator_version,
+            nodeImageVersion: node_image_version,
+            timestamp: timestamp
+        }""")
+        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

@@ -2346,6 +2346,48 @@
     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 a previously used configuration.
+long-summary: |
+    Rollback an agent pool to a previously used Kubernetes version or node image version. 
+    If no specific versions are provided, the nodepool will be rolled back to the most recently used configuration.
+parameters:
+  - name: --kubernetes-version
+    short-summary: Target Kubernetes version for rollback.
+    long-summary: If not specified, uses the most recent available Kubernetes version.
+  - name: --node-image-version
+    short-summary: Target node image version for rollback.
+    long-summary: If not specified, uses the most recent available node image version.
+  - name: --aks-custom-headers
+    short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2.
+  - name: --if-match
+    short-summary: The revision of the resource being updated. This should match the current revision.
+  - name: --if-none-match
+    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 recent configuration.
+    text: az aks nodepool rollback --resource-group MyResourceGroup --cluster-name MyManagedCluster --nodepool-name MyNodePool
+    crafted: true
+  - name: Rollback a nodepool to a specific Kubernetes version.
+    text: az aks nodepool rollback --resource-group MyResourceGroup --cluster-name MyManagedCluster --nodepool-name MyNodePool --kubernetes-version 1.28.5
+    crafted: true
+  - name: Rollback a nodepool to specific Kubernetes and node image versions.
+    text: az aks nodepool rollback --resource-group MyResourceGroup --cluster-name MyManagedCluster --nodepool-name MyNodePool --kubernetes-version 1.28.5 --node-image-version AKSUbuntu-1804gen2containerd-2024.01.15
+    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

@@ -2030,6 +2030,22 @@ 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(
+            "kubernetes_version",
+            help="Target Kubernetes version for rollback. If not specified, uses the most recent available version."
+        )
+        c.argument(
+            "node_image_version", 
+            help="Target node image version for rollback. If not specified, uses the most recent available version."
+        )
+        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,
@@ -275,6 +276,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

@@ -1777,6 +1777,71 @@ 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,
+                          kubernetes_version=None,
+                          node_image_version=None,
+                          aks_custom_headers=None,
+                          if_match=None,
+                          if_none_match=None,
+                          no_wait=False):
+    """Rollback a nodepool to N-1 (previously used configuration)."""
+    
+    # Get the current agent pool
+    current_agentpool = client.get(resource_group_name, cluster_name, nodepool_name)
+    
+    # Get upgrade profile to get recently used versions
+    upgrade_profile = client.get_upgrade_profile(resource_group_name, cluster_name, nodepool_name)
+    
+    # Check if rollback versions are available
+    if not upgrade_profile.recently_used_versions:
+        raise CLIError("No rollback versions are available for this nodepool.")
+    
+    # If no specific versions are provided, use the most recent (N-1) version
+    if not kubernetes_version and not node_image_version:
+        most_recent = upgrade_profile.recently_used_versions[0]
+        kubernetes_version = most_recent.orchestrator_version
+        node_image_version = most_recent.node_image_version
+    
+    # Update the agent pool configuration with rollback versions
+    if kubernetes_version:
+        current_agentpool.orchestrator_version = kubernetes_version
+    if node_image_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,