{AKS} AKS snapshot related commands (Azure#4046)

* AKS Snapshot related commands

* add --snapshot-id to the aks create command

* create/upgrade cluster/nodepool from snapshot

* add warning at deleting snapshot

* Snapshot related recording tests

* fix some style errors

* test_aks_snapshot to ext_matrix_default.json

* some changes according to PR comments, add more command helper info and re-run the recording test

* some changes according to PR comments, add more command helper info and re-run the recording test

Co-authored-by: Charlie Li <charlili@microsoft.com>

src/aks-preview/HISTORY.md

@@ -2,6 +2,18 @@
 
 Release History
 ===============
+0.5.40
++++++
+* Add support for new snapshot commands
+  * `az aks snapshot create`
+  * `az aks snapshot delete`
+  * `az aks snapshot list`
+  * `az aks snapshot show`
+* Add --snapshot-id to creating/upgrading commands
+  * `az aks create --snapshot-id`
+  * `az aks nodepool add --snapshot-id`
+  * `az aks nodepool upgrade --snapshot-id`
+
 0.5.39
 +++++
 * Add commands for agentpool start stop feature

src/aks-preview/azcli_aks_live_test/configs/ext_matrix_default.json

@@ -40,7 +40,8 @@
             "test_aks_nodepool_add_with_workload_runtime",
             "test_aks_nodepool_add_with_gpu_instance_profile",
             "test_aks_create_with_windows_gmsa",
-            "test_aks_update_with_windows_gmsa"
+            "test_aks_update_with_windows_gmsa",
+            "test_aks_snapshot"
         ]
     }
 }

src/aks-preview/azext_aks_preview/_client_factory.py

@@ -38,6 +38,14 @@ def cf_agent_pools(cli_ctx, *_):
     return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).agent_pools
 
 
+def cf_snapshots_client(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW, subscription_id=subscription_id).snapshots
+
+
+def cf_snapshots(cli_ctx, *_):
+    return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).snapshots
+
+
 def cf_maintenance_configurations(cli_ctx, *_):
     return get_mgmt_service_client(cli_ctx, CUSTOM_MGMT_AKS_PREVIEW).maintenance_configurations
 

src/aks-preview/azext_aks_preview/_format.py

@@ -214,3 +214,28 @@ def aks_pod_identities_table_format(result):
     }""")
     # use ordered dicts so headers are predictable
     return parsed.search(result, Options(dict_cls=OrderedDict, custom_functions=_custom_functions(preview)))
+
+
+def aks_list_snapshot_table_format(results):
+    """"Format a list of snapshots as summary results for display with "-o table"."""
+    return [_aks_snapshot_table_format(r) for r in results]
+
+
+def aks_show_snapshot_table_format(result):
+    """Format a snapshot as summary results for display with "-o table"."""
+    return [_aks_snapshot_table_format(result)]
+
+
+def _aks_snapshot_table_format(result):
+    parsed = compile_jmes("""{
+        name: name,
+        location: location,
+        resourceGroup: resourceGroup,
+        nodeImageVersion: nodeImageVersion,
+        kubernetesVersion: kubernetesVersion,
+        osType: osType,
+        osSku: osSku,
+        enableFIPS: enableFIPS
+    }""")
+    # use ordered dicts so headers are predictable
+    return parsed.search(result, Options(dict_cls=OrderedDict))

src/aks-preview/azext_aks_preview/_help.py

@@ -376,6 +376,9 @@
           long-summary: |-
              You do not need to set this if you have set DNS server in the VNET used by the cluster.
              You must set or not set --gmsa-dns-server and --gmsa-root-domain-name at the same time when setting --enable-windows-gmsa.
+        - name: --snapshot-id
+          type: string
+          short-summary: The source snapshot id used to create this cluster.
     examples:
         - name: Create a Kubernetes cluster with an existing SSH public key.
           text: az aks create -g MyResourceGroup -n MyManagedCluster --ssh-key-value /path/to/publickey
@@ -429,6 +432,8 @@
           text: az aks create -g MyResourceGroup -n MyManagedCluster --load-balancer-sku Standard --network-plugin azure --windows-admin-username azure --windows-admin-password 'replacePassword1234$' --enable-windows-gmsa
         - name: Create a kubernetes cluster with enabling Windows gmsa but without setting DNS server in the vnet used by the cluster.
           text: az aks create -g MyResourceGroup -n MyManagedCluster --load-balancer-sku Standard --network-plugin azure --windows-admin-username azure --windows-admin-password 'replacePassword1234$' --enable-windows-gmsa --gmsa-dns-server "10.240.0.4" --gmsa-root-domain-name "contoso.com"
+        - name: create a kubernetes cluster with a snapshot id.
+          text: az aks create -g MyResourceGroup -n MyManagedCluster --kubernetes-version 1.20.9 --snapshot-id "/subscriptions/00000/resourceGroups/AnotherResourceGroup/providers/Microsoft.ContainerService/snapshots/mysnapshot1"
 """.format(sp_cache=AKS_SERVICE_PRINCIPAL_CACHE)
 
 helps['aks scale'] = """
@@ -992,6 +997,9 @@
         - name: --gpu-instance-profile
           type: string
           short-summary: GPU instance profile to partition multi-gpu Nvidia GPUs.
+        - name: --snapshot-id
+          type: string
+          short-summary: The source snapshot id used to create this nodepool.
     examples:
         - name: Create a nodepool in an existing AKS cluster with ephemeral os enabled.
           text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --node-osdisk-type Ephemeral --node-osdisk-size 48
@@ -1001,6 +1009,8 @@
           text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster  --os-sku Ubuntu
         - name: Create a nodepool which can run wasm workloads.
           text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster  --workload-runtime WasmWasi
+        - name: create a kubernetes cluster with a snapshot id.
+          text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --kubernetes-version 1.20.9 --snapshot-id "/subscriptions/00000/resourceGroups/AnotherResourceGroup/providers/Microsoft.ContainerService/snapshots/mysnapshot1"
 """
 
 helps['aks nodepool scale'] = """
@@ -1028,6 +1038,9 @@
         - name: --aks-custom-headers
           type: string
           short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2
+        - name: --snapshot-id
+          type: string
+          short-summary: The source snapshot id used to upgrade this nodepool.
 """
 
 helps['aks nodepool update'] = """
@@ -1457,3 +1470,44 @@
     type: command
     short-summary: List egress endpoints that are required or recommended to be whitelisted for a cluster.
 """
+
+helps['aks snapshot'] = """
+    type: group
+    short-summary: Commands to manage snapshots.
+"""
+
+helps['aks snapshot show'] = """
+    type: command
+    short-summary: Show the details of a snapshot.
+"""
+
+helps['aks snapshot list'] = """
+    type: command
+    short-summary: List snapshots.
+"""
+
+helps['aks snapshot create'] = """
+    type: command
+    short-summary: Create a snapshot of a node pool.
+    parameters:
+        - name: --nodepool-id
+          type: string
+          short-summary: The source nodepool id from which to create this snapshot.
+        - name: --tags
+          type: string
+          short-summary: The tags of the snapshot.
+        - name: --aks-custom-headers
+          type: string
+          short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2
+
+    examples:
+        - name: Create a snapshot.
+          text: az aks snapshot create -g MyResourceGroup -n snapshot1 --nodepool-id "/subscriptions/00000/resourceGroups/AnotherResourceGroup/providers/Microsoft.ContainerService/managedClusters/akscluster1/agentPools/nodepool1"
+        - name: Create a snapshot with custom tags.
+          text: az aks snapshot create -g MyResourceGroup -n snapshot1 --nodepool-id "/subscriptions/00000/resourceGroups/AnotherResourceGroup/providers/Microsoft.ContainerService/managedClusters/akscluster1/agentPools/nodepool1" --tags "foo=bar" "key1=val1"
+"""
+
+helps['aks snapshot delete'] = """
+    type: command
+    short-summary: Delete a snapshot.
+"""

src/aks-preview/azext_aks_preview/_params.py

@@ -18,7 +18,7 @@
 from ._validators import (
     validate_cluster_autoscaler_profile, validate_create_parameters, validate_k8s_version, validate_linux_host_name,
     validate_ssh_key, validate_nodes_count, validate_ip_ranges,
-    validate_nodepool_name, validate_vm_set_type, validate_load_balancer_sku,
+    validate_nodepool_name, validate_vm_set_type, validate_load_balancer_sku, validate_nodepool_id, validate_snapshot_id,
     validate_load_balancer_outbound_ips, validate_load_balancer_outbound_ip_prefixes, validate_nat_gateway_managed_outbound_ip_count,
     validate_taints, validate_priority, validate_eviction_policy, validate_spot_max_price, validate_acr, validate_user,
     validate_load_balancer_outbound_ports, validate_load_balancer_idle_timeout, validate_nat_gateway_idle_timeout, validate_nodepool_tags, validate_addon,
@@ -153,6 +153,7 @@ def load_arguments(self, _):
         c.argument('gmsa_root_domain_name', options_list=['--gmsa-root-domain-name'])
         c.argument('yes', options_list=['--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')
         c.argument('workload_runtime', arg_type=get_enum_type(workload_runtimes), default=CONST_WORKLOAD_RUNTIME_OCI_CONTAINER)
+        c.argument('snapshot_id', type=str, validator=validate_snapshot_id, is_preview=True)
 
     with self.argument_context('aks update') as c:
         c.argument('enable_cluster_autoscaler', options_list=["--enable-cluster-autoscaler", "-e"], action='store_true')
@@ -251,6 +252,7 @@ def load_arguments(self, _):
             c.argument('enable_ultra_ssd', action='store_true')
             c.argument('workload_runtime', arg_type=get_enum_type(workload_runtimes), default=CONST_WORKLOAD_RUNTIME_OCI_CONTAINER)
             c.argument('gpu_instance_profile', arg_type=get_enum_type(gpu_instance_profiles))
+            c.argument('snapshot_id', type=str, validator=validate_snapshot_id, is_preview=True)
 
     for scope in ['aks nodepool show', 'aks nodepool delete', 'aks nodepool scale', 'aks nodepool upgrade', 'aks nodepool update']:
         with self.argument_context(scope) as c:
@@ -259,6 +261,7 @@ def load_arguments(self, _):
     with self.argument_context('aks nodepool upgrade') as c:
         c.argument('max_surge', type=str, validator=validate_max_surge)
         c.argument('aks_custom_headers')
+        c.argument('snapshot_id', type=str, validator=validate_snapshot_id, is_preview=True)
 
     with self.argument_context('aks nodepool update') as c:
         c.argument('enable_cluster_autoscaler', options_list=["--enable-cluster-autoscaler", "-e"], action='store_true')
@@ -382,6 +385,18 @@ def load_arguments(self, _):
                    help='pod labels in key=value [key=value ...].',
                    validator=validate_pod_identity_pod_labels)
 
+    for scope in ['aks snapshot create']:
+        with self.argument_context(scope) as c:
+            c.argument('snapshot_name', type=str, options_list=['--name', '-n'], required=True, validator=validate_linux_host_name, help='The snapshot name.')
+            c.argument('tags', tags_type)
+            c.argument('nodepool_id', type=str, required=True, validator=validate_nodepool_id, help='The nodepool id.')
+            c.argument('aks_custom_headers')
+
+    for scope in ['aks snapshot show', 'aks snapshot delete']:
+        with self.argument_context(scope) as c:
+            c.argument('snapshot_name', type=str, options_list=['--name', '-n'], required=True, validator=validate_linux_host_name, help='The snapshot name.')
+            c.argument('yes', options_list=['--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')
+
 
 def _get_default_install_location(exe_name):
     system = platform.system()

src/aks-preview/azext_aks_preview/_validators.py

@@ -504,3 +504,16 @@ def validate_assign_kubelet_identity(namespace):
         from msrestazure.tools import is_valid_resource_id
         if not is_valid_resource_id(namespace.assign_kubelet_identity):
             raise CLIError("--assign-kubelet-identity is not a valid Azure resource ID.")
+
+
+def validate_nodepool_id(namespace):
+    from msrestazure.tools import is_valid_resource_id
+    if not is_valid_resource_id(namespace.nodepool_id):
+        raise InvalidArgumentValueError("--nodepool-id is not a valid Azure resource ID.")
+
+
+def validate_snapshot_id(namespace):
+    if namespace.snapshot_id:
+        from msrestazure.tools import is_valid_resource_id
+        if not is_valid_resource_id(namespace.snapshot_id):
+            raise InvalidArgumentValueError("--snapshot-id is not a valid Azure resource ID.")

src/aks-preview/azext_aks_preview/commands.py

@@ -9,6 +9,7 @@
 from ._client_factory import cf_maintenance_configurations
 from ._client_factory import cf_container_services
 from ._client_factory import cf_agent_pools
+from ._client_factory import cf_snapshots
 from ._format import aks_show_table_format
 from ._format import aks_addon_list_available_table_format, aks_addon_list_table_format, aks_addon_show_table_format
 from ._format import aks_agentpool_show_table_format
@@ -17,6 +18,8 @@
 from ._format import aks_upgrades_table_format
 from ._format import aks_pod_identities_table_format
 from ._format import aks_pod_identity_exceptions_table_format
+from ._format import aks_show_snapshot_table_format
+from ._format import aks_list_snapshot_table_format
 
 
 def load_command_table(self, _):
@@ -47,6 +50,12 @@ def load_command_table(self, _):
         client_factory=cf_maintenance_configurations
     )
 
+    snapshot_sdk = CliCommandType(
+        operations_tmpl='azext_aks_preview.vendored_sdks.azure_mgmt_preview_aks.'
+                        'operations._snapshots_operations#SnapshotsOperations.{}',
+        client_factory=cf_snapshots
+    )
+
     # AKS managed cluster commands
     with self.command_group('aks', managed_clusters_sdk, client_factory=cf_managed_clusters) as g:
         g.custom_command('kollect', 'aks_kollect')
@@ -130,3 +139,10 @@ def load_command_table(self, _):
     # AKS egress commands
     with self.command_group('aks egress-endpoints', managed_clusters_sdk, client_factory=cf_managed_clusters) as g:
         g.custom_command('list', 'aks_egress_endpoints_list')
+
+    # AKS snapshot commands
+    with self.command_group('aks snapshot', snapshot_sdk, client_factory=cf_snapshots) as g:
+        g.custom_command('list', 'aks_snapshot_list', table_transformer=aks_list_snapshot_table_format)
+        g.custom_show_command('show', 'aks_snapshot_show', table_transformer=aks_show_snapshot_table_format)
+        g.custom_command('create', 'aks_snapshot_create', supports_no_wait=True)
+        g.custom_command('delete', 'aks_snapshot_delete', supports_no_wait=True)