Adding local dns profile to az aks nodepool update

src/aks-preview/HISTORY.rst

@@ -13,6 +13,11 @@ Pending
 +++++++
 * Add support for `ManagedSystem` Agent Pool Mode.
 
+18.0.0b19
++++++++
+- Add `--localdns-config` to `az aks nodepool add` and to `az aks nodepool update` to support configuring a local DNS profile for agent pools.
+
+
 18.0.0b18
 +++++++
 * Add validation error when neither --location or --cluster and --resource-group-name are specified for az extension type list or az extension type version list
@@ -1907,4 +1912,4 @@ Pending
 +++++
 
 * new feature `enable-cluster-autoscaler`
-* default agentType is VMSS
+* default agentType is VMSS

src/aks-preview/azext_aks_preview/_help.py

@@ -2048,6 +2048,9 @@
         - name: --undrainable-node-behavior
           type: string
           short-summary: Define the behavior for undrainable nodes during upgrade. The value should be "Cordon" or "Schedule". The default value is "Schedule".
+        - name: --localdns-config
+          type: string
+          short-summary: Set the localDNS Profile for a nodepool with a JSON config file.
     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
@@ -2223,6 +2226,9 @@
         - name: --undrainable-node-behavior
           type: string
           short-summary: Define the behavior for undrainable nodes during upgrade. The value should be "Cordon" or "Schedule". The default value is "Schedule".
+        - name: --localdns-config
+          type: string
+          short-summary: Set the localDNS Profile for a nodepool with a JSON config file.
     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/_params.py

@@ -1772,6 +1772,11 @@ def load_arguments(self, _):
         )
         # virtual machines
         c.argument("vm_sizes", is_preview=True)
+        # local DNS
+        c.argument(
+            'localdns_config',
+            help='Path to a JSON file to configure the local DNS profile for a new nodepool.'
+        )
 
     with self.argument_context("aks nodepool update") as c:
         c.argument(
@@ -1864,6 +1869,11 @@ def load_arguments(self, _):
             "disable_fips_image",
             action="store_true"
         )
+        # local DNS
+        c.argument(
+            'localdns_config',
+            help='Path to a JSON file to configure the local DNS profile for an existing nodepool.',
+        )
 
     with self.argument_context("aks nodepool upgrade") as c:
         c.argument("max_surge", validator=validate_max_surge)

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -5,6 +5,7 @@
 
 import base64
 import os
+from azure.cli.core.util import get_file_json
 from types import SimpleNamespace
 from typing import Dict, TypeVar, Union, List
 
@@ -810,6 +811,30 @@ def get_disable_fips_image(self) -> bool:
         # read the original value passed by the command
         return self.raw_param.get("disable_fips_image")
 
+    def get_localdns_config(self):
+        return self.raw_param.get("localdns_config")
+
+    def get_localdns_profile(self):
+        """
+        Returns the local DNS profile dict if set, or None.
+        Only supports loading from --localdns-config (JSON file).
+        Assumes the input is always a string filename.
+        """
+        config = self.get_localdns_config()
+        if config:
+            if not isinstance(config, str) or not os.path.isfile(config):
+                raise InvalidArgumentValueError(
+                    f"{config} is not a valid file, or not accessible."
+                )
+            profile = get_file_json(config)
+            if not isinstance(profile, dict):
+                raise InvalidArgumentValueError(
+                    f"Error reading local DNS config from {config}. "
+                    "Please provide a valid JSON file."
+                )
+            return profile
+        return None
+
 
 class AKSPreviewAgentPoolAddDecorator(AKSAgentPoolAddDecorator):
     def __init__(
@@ -1099,6 +1124,50 @@ def set_up_managed_system_mode(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def set_up_localdns_profile(self, agentpool: AgentPool) -> AgentPool:
+        """Set up local DNS profile for the AgentPool object if provided via --localdns-config."""
+        self._ensure_agentpool(agentpool)
+        localdns_profile = self.context.get_localdns_profile()
+        if localdns_profile is not None:
+            kube_dns_overrides = {}
+            vnet_dns_overrides = {}
+
+            def build_override(override_dict):
+                camel_to_snake_case = {
+                    "queryLogging": "query_logging",
+                    "protocol": "protocol",
+                    "forwardDestination": "forward_destination",
+                    "forwardPolicy": "forward_policy",
+                    "maxConcurrent": "max_concurrent",
+                    "cacheDurationInSeconds": "cache_duration_in_seconds",
+                    "serveStaleDurationInSeconds": "serve_stale_duration_in_seconds",
+                    "serveStale": "serve_stale",
+                }
+                valid_keys = set(camel_to_snake_case.values())
+                filtered = {}
+                for k, v in override_dict.items():
+                    if k in camel_to_snake_case:
+                        filtered[camel_to_snake_case[k]] = v
+                    elif k in valid_keys:
+                        filtered[k] = v
+                return self.models.LocalDNSOverride(**filtered)
+
+            # Build kubeDNSOverrides and vnetDNSOverrides from the localdns_profile
+            kube_overrides = localdns_profile.get("kubeDNSOverrides")
+            for key, value in kube_overrides.items():
+                kube_dns_overrides[key] = build_override(value)
+
+            vnet_overrides = localdns_profile.get("vnetDNSOverrides")
+            for key, value in vnet_overrides.items():
+                vnet_dns_overrides[key] = build_override(value)
+
+            agentpool.local_dns_profile = self.models.LocalDNSProfile(
+                mode=localdns_profile.get("mode"),
+                kube_dns_overrides=kube_dns_overrides,
+                vnet_dns_overrides=vnet_dns_overrides,
+            )
+        return agentpool
+
     def construct_agentpool_profile_preview(self) -> AgentPool:
         """The overall controller used to construct the preview AgentPool profile.
 
@@ -1151,6 +1220,8 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self.set_up_agentpool_gateway_profile(agentpool)
         # set up virtual machines profile
         agentpool = self.set_up_virtual_machines_profile(agentpool)
+        # set up local DNS profile
+        agentpool = self.set_up_localdns_profile(agentpool)
         # DO NOT MOVE: keep this at the bottom, restore defaults
         agentpool = self._restore_defaults_in_agentpool(agentpool)
         return agentpool
@@ -1343,6 +1414,50 @@ def update_fips_image(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def update_localdns_profile(self, agentpool: AgentPool) -> AgentPool:
+        """Update local DNS profile for the AgentPool object if provided via --localdns-config."""
+        self._ensure_agentpool(agentpool)
+        localdns_profile = self.context.get_localdns_profile()
+        if localdns_profile is not None:
+            kube_dns_overrides = {}
+            vnet_dns_overrides = {}
+
+            def build_override(override_dict):
+                camel_to_snake_case = {
+                    "queryLogging": "query_logging",
+                    "protocol": "protocol",
+                    "forwardDestination": "forward_destination",
+                    "forwardPolicy": "forward_policy",
+                    "maxConcurrent": "max_concurrent",
+                    "cacheDurationInSeconds": "cache_duration_in_seconds",
+                    "serveStaleDurationInSeconds": "serve_stale_duration_in_seconds",
+                    "serveStale": "serve_stale",
+                }
+                valid_keys = set(camel_to_snake_case.values())
+                filtered = {}
+                for k, v in override_dict.items():
+                    if k in camel_to_snake_case:
+                        filtered[camel_to_snake_case[k]] = v
+                    elif k in valid_keys:
+                        filtered[k] = v
+                return self.models.LocalDNSOverride(**filtered)
+
+            # Build kubeDNSOverrides and vnetDNSOverrides from the localdns_profile
+            kube_overrides = localdns_profile.get("kubeDNSOverrides")
+            for key, value in kube_overrides.items():
+                kube_dns_overrides[key] = build_override(value)
+
+            vnet_overrides = localdns_profile.get("vnetDNSOverrides")
+            for key, value in vnet_overrides.items():
+                vnet_dns_overrides[key] = build_override(value)
+
+            agentpool.local_dns_profile = self.models.LocalDNSProfile(
+                mode=localdns_profile.get("mode"),
+                kube_dns_overrides=kube_dns_overrides,
+                vnet_dns_overrides=vnet_dns_overrides,
+            )
+        return agentpool
+
     def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -> AgentPool:
         """The overall controller used to update the preview AgentPool profile.
 
@@ -1387,6 +1502,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update ssh access
         agentpool = self.update_ssh_access(agentpool)
 
+        # update local DNS profile
+        agentpool = self.update_localdns_profile(agentpool)
+
         return agentpool
 
     def update_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/custom.py

@@ -1419,6 +1419,8 @@ def aks_agentpool_add(
     gateway_prefix_size=None,
     # virtualmachines
     vm_sizes=None,
+    # local DNS
+    localdns_config=None,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()
@@ -1486,6 +1488,8 @@ def aks_agentpool_update(
     if_none_match=None,
     enable_fips_image=False,
     disable_fips_image=False,
+    # local DNS
+    localdns_config=None,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()

src/aks-preview/azext_aks_preview/tests/latest/data/localdnsconfig.json

@@ -0,0 +1,47 @@
+{
+  "mode": "Required",
+  "kubeDNSOverrides": {
+    ".": {
+      "cacheDurationInSeconds": 3600,
+      "forwardDestination": "ClusterCoreDNS",
+      "forwardPolicy": "Sequential",
+      "maxConcurrent": 1000,
+      "protocol": "PreferUDP",
+      "queryLogging": "Error",
+      "serveStale": "Verify",
+      "serveStaleDurationInSeconds": 3600
+    },
+    "cluster.local": {
+      "cacheDurationInSeconds": 3600,
+      "forwardDestination": "ClusterCoreDNS",
+      "forwardPolicy": "Sequential",
+      "maxConcurrent": 1000,
+      "protocol": "ForceTCP",
+      "queryLogging": "Error",
+      "serveStale": "Immediate",
+      "serveStaleDurationInSeconds": 3600
+    }
+  },
+  "vnetDNSOverrides": {
+    ".": {
+      "cacheDurationInSeconds": 3600,
+      "forwardDestination": "VnetDNS",
+      "forwardPolicy": "Sequential",
+      "maxConcurrent": 1000,
+      "protocol": "PreferUDP",
+      "queryLogging": "Error",
+      "serveStale": "Verify",
+      "serveStaleDurationInSeconds": 3600
+    },
+    "cluster.local": {
+      "cacheDurationInSeconds": 3600,
+      "forwardDestination": "ClusterCoreDNS",
+      "forwardPolicy": "Sequential",
+      "maxConcurrent": 1000,
+      "protocol": "ForceTCP",
+      "queryLogging": "Error",
+      "serveStale": "Immediate",
+      "serveStaleDurationInSeconds": 3600
+    }
+  }
+}