Adding localdns profile to az aks nodepool cli

Author: nirnaymsft

src/azure-cli/azure/cli/command_modules/acs/_help.py

@@ -1806,6 +1806,9 @@
   - name: --gateway-prefix-size
     type: int
     short-summary: The size of Public IPPrefix attached to the Gateway-mode node pool. The node pool must be in Gateway mode.
+  - 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
@@ -1955,6 +1958,9 @@
   - name: --if-none-match
     type: string
     short-summary: Set to '*' to allow a new node pool to be created, but to prevent updating an existing node pool. Other values will be ignored.
+  - 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/azure-cli/azure/cli/command_modules/acs/_helpers.py

@@ -321,3 +321,20 @@ def _get_id_from_shared_control_plane_identity(shared_identity) -> int:
             )
         )
     return 0
+
+def process_dns_overrides(overrides_dict, target_dict, build_override_func):
+    """Helper function to safely process DNS overrides with null checks.
+    Processes DNS override dictionaries from LocalDNS configuration,
+    filtering out null values and applying the build function to valid entries.
+    :param overrides_dict: Dictionary containing DNS overrides (can be None)
+    :param target_dict: Target dictionary to populate with processed overrides
+    :param build_override_func: Function to build override objects from dict values
+    """
+    if not isinstance(overrides_dict, dict):
+        raise InvalidArgumentValueError(
+            f"Expected a dictionary for DNS overrides, but got {type(overrides_dict).__name__}: {overrides_dict}"
+        )
+    if overrides_dict is not None:
+        for key, value in overrides_dict.items():
+            if value is not None:
+                target_dict[key] = build_override_func(value)

src/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -972,6 +972,7 @@ def load_arguments(self, _):
         c.argument("if_none_match")
         c.argument('gpu_driver', arg_type=get_enum_type(gpu_driver_install_modes))
         c.argument("gateway_prefix_size", type=int, validator=validate_gateway_prefix_size)
+        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', resource_type=ResourceType.MGMT_CONTAINERSERVICE, operation_group='agent_pools') as c:
         c.argument('enable_cluster_autoscaler', options_list=[
@@ -1003,6 +1004,7 @@ def load_arguments(self, _):
         c.argument('disable_secure_boot', action='store_true')
         c.argument("if_match")
         c.argument("if_none_match")
+        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 upgrade') as c:
         c.argument('max_surge', validator=validate_max_surge)

src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py

@@ -32,7 +32,7 @@
     DecoratorEarlyExitException,
     DecoratorMode,
 )
-from azure.cli.command_modules.acs._helpers import get_snapshot_by_snapshot_id, safe_list_get
+from azure.cli.command_modules.acs._helpers import get_snapshot_by_snapshot_id, safe_list_get, process_dns_overrides
 from azure.cli.command_modules.acs._validators import extract_comma_separated_string
 from azure.cli.command_modules.acs.base_decorator import BaseAKSContext, BaseAKSModels, BaseAKSParamDict
 from azure.cli.core import AzCommandsLoader
@@ -63,7 +63,6 @@
 # TODO:
 # 1. Add extra type checking for getter functions
 
-
 # pylint: disable=too-few-public-methods
 class AKSAgentPoolModels(BaseAKSModels):
     """Store the models used in aks agentpool series of commands.
@@ -88,15 +87,13 @@ def __choose_agentpool_model_by_agentpool_decorator_mode(self):
             return self.ManagedClusterAgentPoolProfile
         return self.AgentPool
 
-
 # pylint: disable=too-few-public-methods
 class AKSAgentPoolParamDict(BaseAKSParamDict):
     """Store the original parameters passed in by aks agentpool series of commands as an internal dictionary.
 
     Only expose the "get" method externally to obtain parameter values, while recording usage.
     """
 
-
 # pylint: disable=too-many-public-methods
 class AKSAgentPoolContext(BaseAKSContext):
     """Implement getter functions for all parameters in aks_agentpool_add and aks_agentpool_update.
@@ -1685,7 +1682,102 @@ def get_gateway_prefix_size(self) -> Union[int, None]:
         """
         return self.raw_param.get('gateway_prefix_size')
 
+    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
+
+    def build_localdns_profile(self, agentpool: AgentPool) -> AgentPool:
+        """Build local DNS profile for the AgentPool object if provided via --localdns-config."""
+        localdns_profile = self.get_localdns_profile()
+        kube_dns_overrides, vnet_dns_overrides = None, None
+
+        if localdns_profile is not None:
+            def find_keys_case_insensitive(dictionary, target_keys):
+                """Find multiple keys case-insensitively and return a dict mapping target_key -> actual_key"""
+                result = {}
+                lowered_keys = {key.lower(): key for key in dictionary.keys()}
+                for target_key in target_keys:
+                    lowered_target = target_key.lower()
+                    if lowered_target in lowered_keys:
+                        result[target_key] = lowered_keys[lowered_target]
+                    else:
+                        result[target_key] = None
+                return result
+
+            def build_override(override_dict):
+                if not isinstance(override_dict, dict):
+                    raise InvalidArgumentValueError(
+                        f"Expected a dictionary for DNS override settings,"
+                        f" but got {type(override_dict).__name__}: {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
+            key_mappings = find_keys_case_insensitive(localdns_profile, ["kubeDNSOverrides", "vnetDNSOverrides"])
+            actual_kube_key = key_mappings["kubeDNSOverrides"]
+            if actual_kube_key:
+                logger.debug("Found kubeDNSOverrides key as: %s", actual_kube_key)
+                kube_dns_overrides = {}
+                process_dns_overrides(
+                    localdns_profile.get(actual_kube_key),
+                    kube_dns_overrides,
+                    build_override
+                )
+
+            actual_vnet_key = key_mappings["vnetDNSOverrides"]
+            if actual_vnet_key:
+                logger.debug("Found vnetDNSOverrides key as: %s", actual_vnet_key)
+                vnet_dns_overrides = {}
+                process_dns_overrides(
+                    localdns_profile.get(actual_vnet_key),
+                    vnet_dns_overrides,
+                    build_override
+                )
+
+            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
+    
 class AKSAgentPoolAddDecorator:
     def __init__(
         self,
@@ -2186,6 +2278,8 @@ def construct_agentpool_profile_default(self, bypass_restore_defaults: bool = Fa
         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)
         # restore defaults
         if not bypass_restore_defaults:
             agentpool = self._restore_defaults_in_agentpool(agentpool)
@@ -2227,6 +2321,10 @@ def add_agentpool(self, agentpool: AgentPool) -> AgentPool:
             headers=self.context.get_aks_custom_headers(),
         )
 
+    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)
+        return self.context.build_localdns_profile(agentpool)
 
 class AKSAgentPoolUpdateDecorator:
     def __init__(
@@ -2513,11 +2611,12 @@ def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -
         agentpool = self.update_os_sku(agentpool)
         # update fips image
         agentpool = self.update_fips_image(agentpool)
-
         # update vtpm
         agentpool = self.update_vtpm(agentpool)
         # update secure boot
         agentpool = self.update_secure_boot(agentpool)
+        # update local DNS profile
+        agentpool = self.update_localdns_profile(agentpool)
         return agentpool
 
     def update_agentpool(self, agentpool: AgentPool) -> AgentPool:
@@ -2553,3 +2652,8 @@ def update_agentpool(self, agentpool: AgentPool) -> AgentPool:
             if_none_match=self.context.get_if_none_match(),
             headers=self.context.get_aks_custom_headers(),
         )
+
+    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)
+        return self.context.build_localdns_profile(agentpool)