{AKS} Support reset default value for loadbalancer profile and natgateway profile

Author: MartinForReal

src/aks-preview/HISTORY.rst

@@ -51,6 +51,7 @@ Pending
 * Add `--enable-ai-toolchain-operator` to `az aks create` and `az aks update`.
 * Add `--disable-ai-toolchain-operator` to the `az aks update` command.
 * Refactor azure service mesh related code to meet cli style requirements.
+* Support reset default value for loadbalancer profile and natgateway profile
 
 1.0.0b4
 +++++++

src/aks-preview/azext_aks_preview/_loadbalancer.py

@@ -111,8 +111,8 @@ def configure_load_balancer_profile(
     """configure a load balancer with customer supplied values"""
     if any(
         [
-            managed_outbound_ip_count,
-            managed_outbound_ipv6_count,
+            managed_outbound_ip_count is not None,
+            managed_outbound_ipv6_count is not None,
             outbound_ips,
             outbound_ip_prefixes,
         ]
@@ -152,7 +152,7 @@ def configure_load_balancer_profile(
             )
         else:
             profile.outbound_ip_prefixes = None
-        if managed_outbound_ip_count or managed_outbound_ipv6_count:
+        if managed_outbound_ip_count is not None or managed_outbound_ipv6_count is not None:
             if profile.managed_outbound_i_ps is None:
                 if isinstance(models, SimpleNamespace):
                     ManagedClusterLoadBalancerProfileManagedOutboundIPs = (
@@ -165,14 +165,14 @@ def configure_load_balancer_profile(
                 profile.managed_outbound_i_ps = (
                     ManagedClusterLoadBalancerProfileManagedOutboundIPs()
                 )
-            if managed_outbound_ip_count:
+            if managed_outbound_ip_count is not None:
                 profile.managed_outbound_i_ps.count = managed_outbound_ip_count
-            if managed_outbound_ipv6_count:
+            if managed_outbound_ipv6_count is not None:
                 profile.managed_outbound_i_ps.count_ipv6 = managed_outbound_ipv6_count
         else:
             profile.managed_outbound_i_ps = None
 
-    if outbound_ports:
+    if outbound_ports is not None:
         profile.allocated_outbound_ports = outbound_ports
     if idle_timeout:
         profile.idle_timeout_in_minutes = idle_timeout
@@ -191,11 +191,11 @@ def is_load_balancer_profile_provided(
 ):
     return any(
         [
-            managed_outbound_ip_count,
-            managed_outbound_ipv6_count,
+            managed_outbound_ip_count is not None,
+            managed_outbound_ipv6_count is not None,
             outbound_ips,
             ip_prefixes,
-            outbound_ports,
+            outbound_ports is not None,
             idle_timeout,
         ]
     )

src/aks-preview/azext_aks_preview/_natgateway.py

@@ -0,0 +1,42 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from types import SimpleNamespace
+
+
+def create_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, models: SimpleNamespace):
+    """parse and build NAT gateway profile"""
+    if not is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
+        return None
+
+    profile = models.ManagedClusterNATGatewayProfile()
+    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models)
+
+
+def update_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models: SimpleNamespace):
+    """parse and update an existing NAT gateway profile"""
+    if not is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
+        return profile
+    if not profile:
+        profile = models.ManagedClusterNATGatewayProfile()
+    return configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models)
+
+
+def is_nat_gateway_profile_provided(managed_outbound_ip_count, idle_timeout):
+    return any([managed_outbound_ip_count is not None, idle_timeout])
+
+
+def configure_nat_gateway_profile(managed_outbound_ip_count, idle_timeout, profile, models: SimpleNamespace):
+    """configure a NAT Gateway with customer supplied values"""
+    if managed_outbound_ip_count is not None:
+        ManagedClusterManagedOutboundIPProfile = models.ManagedClusterManagedOutboundIPProfile
+        profile.managed_outbound_ip_profile = ManagedClusterManagedOutboundIPProfile(
+            count=managed_outbound_ip_count
+        )
+
+    if idle_timeout:
+        profile.idle_timeout_in_minutes = idle_timeout
+
+    return profile

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -43,6 +43,10 @@
 from azext_aks_preview._loadbalancer import (
     update_load_balancer_profile as _update_load_balancer_profile,
 )
+from azext_aks_preview._natgateway import create_nat_gateway_profile
+from azext_aks_preview._natgateway import (
+    update_nat_gateway_profile as _update_nat_gateway_profile
+)
 from azext_aks_preview._podidentity import (
     _fill_defaults_for_pod_identity_profile,
     _is_pod_identity_addon_enabled,
@@ -2739,6 +2743,13 @@ def set_up_network_profile(self, mc: ManagedCluster) -> ManagedCluster:
                 models=self.models.load_balancer_models,
             )
 
+        if self.context.get_nat_gateway_managed_outbound_ip_count() is not None:
+            network_profile.nat_gateway_profile = create_nat_gateway_profile(
+                self.context.get_nat_gateway_managed_outbound_ip_count(),
+                self.context.get_nat_gateway_idle_timeout(),
+                models=self.models.nat_gateway_models,
+            )
+
         network_profile.network_plugin_mode = self.context.get_network_plugin_mode()
 
         if self.context.get_enable_cilium_dataplane():
@@ -3574,6 +3585,10 @@ def get_special_parameter_default_value_pairs_list(self) -> List[Tuple[Any, Any]
             (self.context.get_nodepool_labels(), None),
             (self.context.get_nodepool_taints(), None),
             (self.context.raw_param.get("upgrade_settings"), None),
+            (self.context.get_load_balancer_managed_outbound_ip_count(), None),
+            (self.context.get_load_balancer_managed_outbound_ipv6_count(), None),
+            (self.context.get_load_balancer_outbound_ports(), None),
+            (self.context.get_nat_gateway_managed_outbound_ip_count(), None),
         ]
 
     def check_raw_parameters(self):
@@ -3601,7 +3616,6 @@ def check_raw_parameters(self):
             if pair[0] != pair[1]:
                 is_different_from_special_default = True
                 break
-
         if is_changed or is_different_from_special_default:
             return
 
@@ -3817,6 +3831,29 @@ def update_load_balancer_profile(self, mc: ManagedCluster) -> ManagedCluster:
             )
         return mc
 
+    def update_nat_gateway_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update nat gateway profile for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        if not mc.network_profile:
+            raise UnknownError(
+                "Unexpectedly get an empty network profile in the process of updating nat gateway profile."
+            )
+        outbound_type = self.context.get_outbound_type()
+        if outbound_type and outbound_type != CONST_OUTBOUND_TYPE_MANAGED_NAT_GATEWAY:
+            mc.network_profile.nat_gateway_profile = None
+        else:
+            mc.network_profile.nat_gateway_profile = _update_nat_gateway_profile(
+                self.context.get_nat_gateway_managed_outbound_ip_count(),
+                self.context.get_nat_gateway_idle_timeout(),
+                mc.network_profile.nat_gateway_profile,
+                models=self.models.nat_gateway_models,
+            )
+        return mc
+
     def update_outbound_type_in_network_profile(self, mc: ManagedCluster) -> ManagedCluster:
         """Update outbound type of network profile for the ManagedCluster object.
         :return: the ManagedCluster object
@@ -4657,6 +4694,8 @@ def update_mc_profile_preview(self) -> ManagedCluster:
         mc = self.update_outbound_type_in_network_profile(mc)
         # update loadbalancer profile
         mc = self.update_load_balancer_profile(mc)
+        # update natgateway profile
+        mc = self.update_nat_gateway_profile(mc)
         # update kube proxy config
         mc = self.update_kube_proxy_config(mc)
         # update custom ca trust certificates