{AKS} Add v2 managed cluster decorator

src/aks-preview/azext_aks_preview/_help.py

@@ -1211,6 +1211,9 @@
         - name: --node-taints
           type: string
           short-summary: The node taints for the node pool.
+        - name: --aks-custom-headers
+          type: string
+          short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2
     examples:
       - name: Reconcile the nodepool back to its current state.
         text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster

src/aks-preview/azext_aks_preview/_helpers.py

@@ -8,7 +8,7 @@
 from azure.cli.core.azclierror import InvalidArgumentValueError, ResourceNotFoundError
 from azure.core.exceptions import AzureError
 
-from azext_aks_preview._client_factory import cf_nodepool_snapshots
+from azext_aks_preview._client_factory import cf_nodepool_snapshots, cf_mc_snapshots
 from azext_aks_preview._consts import CONST_CONTAINER_NAME_MAX_LENGTH
 
 
@@ -98,3 +98,30 @@ def get_nodepool_snapshot(cli_ctx, resource_group_name, snapshot_name):
             raise ResourceNotFoundError("Snapshot '{}' not found.".format(snapshot_name))
         raise map_azure_error_to_cli_error(ex)
     return snapshot
+
+
+def get_cluster_snapshot_by_snapshot_id(cli_ctx, snapshot_id):
+    _re_mc_snapshot_resource_id = re.compile(
+        r"/subscriptions/(.*?)/resourcegroups/(.*?)/providers/microsoft.containerservice/managedclustersnapshots/(.*)",
+        flags=re.IGNORECASE,
+    )
+    snapshot_id = snapshot_id.lower()
+    match = _re_mc_snapshot_resource_id.search(snapshot_id)
+    if match:
+        resource_group_name = match.group(2)
+        snapshot_name = match.group(3)
+        return get_cluster_snapshot(cli_ctx, resource_group_name, snapshot_name)
+    raise InvalidArgumentValueError(
+        "Cannot parse snapshot name from provided resource id {}.".format(snapshot_id))
+
+
+def get_cluster_snapshot(cli_ctx, resource_group_name, snapshot_name):
+    snapshot_client = cf_mc_snapshots(cli_ctx)
+    try:
+        snapshot = snapshot_client.get(resource_group_name, snapshot_name)
+    # track 2 sdk raise exception from azure.core.exceptions
+    except AzureError as ex:
+        if "not found" in ex.message:
+            raise ResourceNotFoundError("Managed cluster snapshot '{}' not found.".format(snapshot_name))
+        raise map_azure_error_to_cli_error(ex)
+    return snapshot

src/aks-preview/azext_aks_preview/_loadbalancer.py

@@ -3,7 +3,9 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from distutils.version import StrictVersion  # pylint: disable=no-name-in-module,import-error
+from distutils.version import StrictVersion
+from types import SimpleNamespace
+
 from knack.log import get_logger
 
 logger = get_logger(__name__)
@@ -37,7 +39,10 @@ def create_load_balancer_profile(managed_outbound_ip_count, managed_outbound_ipv
                                              outbound_ip_prefixes, outbound_ports, idle_timeout):
         return None
 
-    ManagedClusterLoadBalancerProfile = models.get("ManagedClusterLoadBalancerProfile")
+    if isinstance(models, SimpleNamespace):
+        ManagedClusterLoadBalancerProfile = models.ManagedClusterLoadBalancerProfile
+    else:
+        ManagedClusterLoadBalancerProfile = models.get("ManagedClusterLoadBalancerProfile")
     profile = ManagedClusterLoadBalancerProfile()
     return configure_load_balancer_profile(managed_outbound_ip_count, managed_outbound_ipv6_count, outbound_ips,
                                            outbound_ip_prefixes, outbound_ports, idle_timeout, profile, models)
@@ -64,26 +69,39 @@ def configure_load_balancer_profile(managed_outbound_ip_count, managed_outbound_
         profile.managed_outbound_i_ps = None
 
         if managed_outbound_ip_count or managed_outbound_ipv6_count:
-            ManagedClusterLoadBalancerProfileManagedOutboundIPs = models.get(
-                "ManagedClusterLoadBalancerProfileManagedOutboundIPs"
-            )
+            if isinstance(models, SimpleNamespace):
+                ManagedClusterLoadBalancerProfileManagedOutboundIPs = (
+                    models.ManagedClusterLoadBalancerProfileManagedOutboundIPs
+                )
+            else:
+                ManagedClusterLoadBalancerProfileManagedOutboundIPs = models.get(
+                    "ManagedClusterLoadBalancerProfileManagedOutboundIPs"
+                )
             profile.managed_outbound_i_ps = ManagedClusterLoadBalancerProfileManagedOutboundIPs()
             if managed_outbound_ip_count:
                 profile.managed_outbound_i_ps.count = managed_outbound_ip_count
             if managed_outbound_ipv6_count:
                 profile.managed_outbound_i_ps.count_ipv6 = managed_outbound_ipv6_count
         if outbound_ip_resources:
-            ManagedClusterLoadBalancerProfileOutboundIPs = models.get(
-                "ManagedClusterLoadBalancerProfileOutboundIPs"
-            )
+            if isinstance(models, SimpleNamespace):
+                ManagedClusterLoadBalancerProfileOutboundIPs = models.ManagedClusterLoadBalancerProfileOutboundIPs
+            else:
+                ManagedClusterLoadBalancerProfileOutboundIPs = models.get(
+                    "ManagedClusterLoadBalancerProfileOutboundIPs"
+                )
             # ips -> i_ps due to track 2 naming issue
             profile.outbound_i_ps = ManagedClusterLoadBalancerProfileOutboundIPs(
                 public_i_ps=outbound_ip_resources
             )
         if outbound_ip_prefix_resources:
-            ManagedClusterLoadBalancerProfileOutboundIPPrefixes = models.get(
-                "ManagedClusterLoadBalancerProfileOutboundIPPrefixes"
-            )
+            if isinstance(models, SimpleNamespace):
+                ManagedClusterLoadBalancerProfileOutboundIPPrefixes = (
+                    models.ManagedClusterLoadBalancerProfileOutboundIPPrefixes
+                )
+            else:
+                ManagedClusterLoadBalancerProfileOutboundIPPrefixes = models.get(
+                    "ManagedClusterLoadBalancerProfileOutboundIPPrefixes"
+                )
             profile.outbound_ip_prefixes = ManagedClusterLoadBalancerProfileOutboundIPPrefixes(
                 public_ip_prefixes=outbound_ip_prefix_resources
             )
@@ -107,7 +125,10 @@ def is_load_balancer_profile_provided(managed_outbound_ip_count, managed_outboun
 def _get_load_balancer_outbound_ips(load_balancer_outbound_ips, models):
     """parse load balancer profile outbound IP ids and return an array of references to the outbound IP resources"""
     load_balancer_outbound_ip_resources = None
-    ResourceReference = models.get("ResourceReference")
+    if isinstance(models, SimpleNamespace):
+        ResourceReference = models.ResourceReference
+    else:
+        ResourceReference = models.get("ResourceReference")
     if load_balancer_outbound_ips:
         load_balancer_outbound_ip_resources = \
             [ResourceReference(id=x.strip()) for x in load_balancer_outbound_ips.split(',')]
@@ -118,7 +139,10 @@ def _get_load_balancer_outbound_ip_prefixes(load_balancer_outbound_ip_prefixes,
     """parse load balancer profile outbound IP prefix ids and return an array \
     of references to the outbound IP prefix resources"""
     load_balancer_outbound_ip_prefix_resources = None
-    ResourceReference = models.get("ResourceReference")
+    if isinstance(models, SimpleNamespace):
+        ResourceReference = models.ResourceReference
+    else:
+        ResourceReference = models.get("ResourceReference")
     if load_balancer_outbound_ip_prefixes:
         load_balancer_outbound_ip_prefix_resources = \
             [ResourceReference(id=x.strip()) for x in load_balancer_outbound_ip_prefixes.split(',')]

src/aks-preview/azext_aks_preview/_params.py

@@ -187,7 +187,6 @@ def load_arguments(self, _):
         c.argument('service_cidrs')
         c.argument('load_balancer_sku', arg_type=get_enum_type(load_balancer_skus), validator=validate_load_balancer_sku)
         c.argument('load_balancer_managed_outbound_ip_count', type=int)
-        c.argument('load_balancer_managed_outbound_ipv6_count', type=int)
         c.argument('load_balancer_outbound_ips', validator=validate_load_balancer_outbound_ips)
         c.argument('load_balancer_outbound_ip_prefixes', validator=validate_load_balancer_outbound_ip_prefixes)
         c.argument('load_balancer_outbound_ports', type=int, validator=validate_load_balancer_outbound_ports)
@@ -231,7 +230,6 @@ def load_arguments(self, _):
         c.argument('enable_msi_auth_for_monitoring', arg_type=get_three_state_flag(), is_preview=True)
         c.argument('aci_subnet_name')
         c.argument('appgw_name', arg_group='Application Gateway')
-        c.argument('appgw_subnet_prefix', arg_group='Application Gateway', deprecate_info=c.deprecate(redirect='--appgw-subnet-cidr', hide=True))
         c.argument('appgw_subnet_cidr', arg_group='Application Gateway')
         c.argument('appgw_id', arg_group='Application Gateway')
         c.argument('appgw_subnet_id', arg_group='Application Gateway')
@@ -245,6 +243,7 @@ def load_arguments(self, _):
         c.argument('node_vm_size', options_list=[
                    '--node-vm-size', '-s'], completer=get_vm_size_completion_list)
         c.argument('os_sku', arg_type=get_enum_type(node_os_skus))
+        c.argument('snapshot_id', validator=validate_snapshot_id)
         c.argument('vnet_subnet_id', validator=validate_vnet_subnet_id)
         c.argument('pod_subnet_id', validator=validate_pod_subnet_id)
         c.argument('enable_node_public_ip', action='store_true')
@@ -267,7 +266,6 @@ def load_arguments(self, _):
         c.argument('enable_encryption_at_host', arg_type=get_three_state_flag(), help='Enable EncryptionAtHost.')
         c.argument('enable_ultra_ssd', action='store_true')
         c.argument('enable_fips_image', action='store_true')
-        c.argument('snapshot_id', validator=validate_snapshot_id)
         c.argument('kubelet_config')
         c.argument('linux_os_config')
         c.argument('disable_disk_driver', arg_type=get_three_state_flag())
@@ -279,8 +277,11 @@ def load_arguments(self, _):
         # extensions
         # managed cluster
         c.argument('node_resource_group')
-        c.argument('ip_families')
         c.argument('http_proxy_config')
+        c.argument('ip_families')
+        c.argument('pod_cidrs')
+        c.argument('service_cidrs')
+        c.argument('load_balancer_managed_outbound_ipv6_count', type=int)
         c.argument('enable_pod_security_policy', action='store_true')
         c.argument('enable_pod_identity', action='store_true')
         c.argument('enable_workload_identity', arg_type=get_three_state_flag())
@@ -412,10 +413,9 @@ def load_arguments(self, _):
                        '--node-vm-size', '-s'], completer=get_vm_size_completion_list)
             c.argument('os_type')
             c.argument('os_sku', arg_type=get_enum_type(node_os_skus))
-            c.argument('vnet_subnet_id',
-                       validator=validate_vnet_subnet_id)
-            c.argument('pod_subnet_id',
-                       validator=validate_pod_subnet_id)
+            c.argument('snapshot_id', validator=validate_snapshot_id)
+            c.argument('vnet_subnet_id', validator=validate_vnet_subnet_id)
+            c.argument('pod_subnet_id', validator=validate_pod_subnet_id)
             c.argument('enable_node_public_ip', action='store_true')
             c.argument('node_public_ip_prefix_id')
             c.argument('enable_cluster_autoscaler', options_list=[
@@ -431,9 +431,9 @@ def load_arguments(self, _):
             c.argument('node_taints', validator=validate_taints)
             c.argument('node_osdisk_type', arg_type=get_enum_type(node_os_disk_types))
             c.argument('node_osdisk_size', type=int)
+            c.argument('max_surge', validator=validate_max_surge)
             c.argument('mode', arg_type=get_enum_type(node_mode_types))
             c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes))
-            c.argument('max_surge', validator=validate_max_surge)
             c.argument('max_pods', type=int, options_list=['--max-pods', '-m'])
             c.argument('node_zones', zones_type, options_list=['--node-zones'], help='(--node-zones will be deprecated) Space-separated list of availability zones where agent nodes will be placed.', deprecate_info=c.deprecate(redirect='--zones', hide='2.37.0'))
             c.argument('zones', zones_type, options_list=['--zones', '-z'], help='Space-separated list of availability zones where agent nodes will be placed.')
@@ -442,7 +442,6 @@ def load_arguments(self, _):
                        '--enable-encryption-at-host'], action='store_true')
             c.argument('enable_ultra_ssd', action='store_true')
             c.argument('enable_fips_image', action='store_true')
-            c.argument('snapshot_id', validator=validate_snapshot_id)
             c.argument('kubelet_config')
             c.argument('linux_os_config')
             c.argument('aks_custom_headers')
@@ -482,9 +481,9 @@ def load_arguments(self, _):
         c.argument('labels', nargs='*', validator=validate_nodepool_labels)
         c.argument('tags', tags_type)
         c.argument('node_taints', validator=validate_taints)
+        c.argument('max_surge', validator=validate_max_surge)
         c.argument('mode', arg_type=get_enum_type(node_mode_types))
         c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes))
-        c.argument('max_surge', validator=validate_max_surge)
 
     with self.argument_context('aks addon show') as c:
         c.argument('addon', options_list=[

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -51,6 +51,7 @@ def __init__(
         agentpool_decorator_mode: AgentPoolDecoratorMode,
     ):
         super().__init__(cmd, raw_parameters, models, decorator_mode, agentpool_decorator_mode)
+        # used to store external functions
         self.__external_functions = None
 
     @property
@@ -252,7 +253,8 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self.set_up_motd(agentpool)
         # set up gpu profiles
         agentpool = self.set_up_gpu_propertes(agentpool)
-        # restore defaults
+
+        # DO NOT MOVE: keep this at the bottom, restore defaults
         agentpool = self._restore_defaults_in_agentpool(agentpool)
         return agentpool
 

src/aks-preview/azext_aks_preview/custom.py

@@ -1613,46 +1613,46 @@ def aks_agentpool_add(cmd,      # pylint: disable=unused-argument,too-many-local
                       resource_group_name,
                       cluster_name,
                       nodepool_name,
-                      tags=None,
                       kubernetes_version=None,
-                      node_zones=None,
-                      zones=None,
-                      enable_node_public_ip=False,
-                      node_public_ip_prefix_id=None,
                       node_vm_size=None,
-                      node_osdisk_type=None,
-                      node_osdisk_size=0,
-                      node_count=3,
-                      vnet_subnet_id=None,
-                      pod_subnet_id=None,
-                      ppg=None,
-                      max_pods=0,
                       os_type=None,
                       os_sku=None,
-                      enable_fips_image=False,
+                      vnet_subnet_id=None,
+                      pod_subnet_id=None,
+                      enable_node_public_ip=False,
+                      node_public_ip_prefix_id=None,
+                      enable_cluster_autoscaler=False,
                       min_count=None,
                       max_count=None,
-                      enable_cluster_autoscaler=False,
-                      scale_down_mode=CONST_SCALE_DOWN_MODE_DELETE,
-                      node_taints=None,
+                      node_count=3,
                       priority=CONST_SCALE_SET_PRIORITY_REGULAR,
                       eviction_policy=CONST_SPOT_EVICTION_POLICY_DELETE,
                       spot_max_price=float('nan'),
                       labels=None,
+                      tags=None,
+                      node_taints=None,
+                      node_osdisk_type=None,
+                      node_osdisk_size=0,
                       max_surge=None,
                       mode="User",
-                      aks_custom_headers=None,
-                      kubelet_config=None,
-                      linux_os_config=None,
+                      scale_down_mode=CONST_SCALE_DOWN_MODE_DELETE,
+                      max_pods=0,
+                      node_zones=None,
+                      zones=None,
+                      ppg=None,
                       enable_encryption_at_host=False,
                       enable_ultra_ssd=False,
-                      workload_runtime=None,
-                      gpu_instance_profile=None,
+                      enable_fips_image=False,
+                      kubelet_config=None,
+                      linux_os_config=None,
                       snapshot_id=None,
                       host_group_id=None,
                       crg_id=None,
                       message_of_the_day=None,
-                      no_wait=False):
+                      workload_runtime=None,
+                      gpu_instance_profile=None,
+                      no_wait=False,
+                      aks_custom_headers=None,):
     instances = client.list(resource_group_name, cluster_name)
     for agentpool_profile in instances:
         if agentpool_profile.name == nodepool_name:
@@ -1841,17 +1841,18 @@ def aks_agentpool_update(cmd,   # pylint: disable=unused-argument
                          resource_group_name,
                          cluster_name,
                          nodepool_name,
-                         tags=None,
                          enable_cluster_autoscaler=False,
                          disable_cluster_autoscaler=False,
                          update_cluster_autoscaler=False,
-                         scale_down_mode=None,
                          min_count=None, max_count=None,
-                         max_surge=None,
-                         mode=None,
                          labels=None,
+                         tags=None,
                          node_taints=None,
-                         no_wait=False):
+                         max_surge=None,
+                         mode=None,
+                         scale_down_mode=None,
+                         no_wait=False,
+                         aks_custom_headers=None):
 
     update_autoscaler = enable_cluster_autoscaler + \
         disable_cluster_autoscaler + update_cluster_autoscaler