[AKS] aks machine update cli preview (#9346)

Author: zjpjack-github

src/aks-preview/HISTORY.rst

@@ -11,6 +11,7 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+* `az aks machine update`: Add support for updating machine tags, node taints and node labels.
 * Fix `az aks bastion` subshell defaulting to cmd on Windows when invoked from PowerShell by implementing grandparent process detection to identify the actual user shell. 
 
 19.0.0b8

src/aks-preview/azext_aks_preview/_help.py

@@ -2620,6 +2620,30 @@
          short-summary: The ipTags of the machine public IPs.
 """
 
+helps['aks machine update'] = """
+   type: command
+   short-summary: Update the specified machine in an agentpool
+   parameters:
+       - name: --cluster-name
+         type: string
+         short-summary: Name of the managed cluster.
+       - name: --nodepool-name
+         type: string
+         short-summary: Name of the agentpool of a managed cluster.
+       - name: --machine-name
+         type: string
+         short-summary: Host name of the machine.
+       - name: --tags
+         type: string
+         short-summary: The tags of the machine.
+       - name: --labels
+         type: string
+         short-summary: The labels of the machine.
+       - name: --node-taints
+         type: string
+         short-summary: The taints of the machine.
+"""
+
 helps['aks machine list'] = """
    type: command
    short-summary: List the details for all machines in an agentpool

src/aks-preview/azext_aks_preview/_params.py

@@ -2248,6 +2248,14 @@ def load_arguments(self, _):
             help="Version of Kubernetes to use for the machine.",
         )
 
+    with self.argument_context("aks machine update") as c:
+        c.argument(
+            "machine_name", help="The machine name."
+        )
+        c.argument("tags", tags_type, help="The tags to set on the machine.")
+        c.argument("node_taints", validator=validate_nodepool_taints)
+        c.argument("labels", nargs="*", help="Labels to set on the machine.")
+
     with self.argument_context("aks operation") as c:
         c.argument(
             "nodepool_name",

src/aks-preview/azext_aks_preview/commands.py

@@ -301,6 +301,7 @@ def load_command_table(self, _):
             "show", "aks_machine_show", table_transformer=aks_machine_show_table_format
         )
         g.custom_command("add", "aks_machine_add", supports_no_wait=True)
+        g.custom_command("update", "aks_machine_update", supports_no_wait=True)
 
     with self.command_group(
         "aks operation", operations_sdk, client_factory=cf_operations

src/aks-preview/azext_aks_preview/custom.py

@@ -120,6 +120,7 @@
 )
 from azext_aks_preview.machine import (
     add_machine,
+    update_machine,
 )
 from azext_aks_preview.jwtauthenticator import (
     aks_jwtauthenticator_add_internal,
@@ -2611,6 +2612,33 @@ def aks_machine_add(
     return add_machine(cmd, client, raw_parameters, no_wait)
 
 
+# pylint: disable=unused-argument
+def aks_machine_update(
+    cmd,
+    client,
+    resource_group_name,
+    cluster_name,
+    nodepool_name,
+    machine_name=None,
+    tags=None,
+    node_taints=None,
+    labels=None,
+    no_wait=False,
+):
+    existedMachine = None
+    try:
+        existedMachine = client.get(resource_group_name, cluster_name, nodepool_name, machine_name)
+    except ResourceNotFoundError:
+        raise ClientRequestError(
+            f"Machine '{machine_name}' does not exist. Please use 'az aks machine list' to get current list of machines."
+        )
+
+    if existedMachine:
+        # DO NOT MOVE: get all the original parameters and save them as a dictionary
+        raw_parameters = locals()
+        return update_machine(client, raw_parameters, existedMachine, no_wait)
+
+
 def aks_addon_list_available():
     available_addons = []
     for k, v in ADDONS.items():

src/aks-preview/azext_aks_preview/machine.py

@@ -11,6 +11,18 @@
 from azure.cli.core.util import sdk_no_wait
 
 
+def parse_key_value_list(pairs):
+    result = {}
+    if pairs is None:
+        return result
+    for pair in pairs:
+        if "=" not in pair:
+            raise ValueError(f"Invalid format '{pair}'. Expected format key=value.")
+        key, value = pair.split("=", 1)
+        result[key.strip()] = value.strip()
+    return result
+
+
 def add_machine(cmd, client, raw_parameters, no_wait):
     resource_group_name = raw_parameters.get("resource_group_name")
     cluster_name = raw_parameters.get("cluster_name")
@@ -30,6 +42,35 @@ def add_machine(cmd, client, raw_parameters, no_wait):
     )
 
 
+def update_machine(client, raw_parameters, existedMachine, no_wait):
+    resource_group_name = raw_parameters.get("resource_group_name")
+    cluster_name = raw_parameters.get("cluster_name")
+    nodepool_name = raw_parameters.get("nodepool_name")
+    machine_name = raw_parameters.get("machine_name")
+
+    updated_machine = updateMachine(raw_parameters, existedMachine)
+
+    return sdk_no_wait(
+        no_wait,
+        client.begin_create_or_update,
+        resource_group_name,
+        cluster_name,
+        nodepool_name,
+        machine_name,
+        updated_machine,
+    )
+
+
+def updateMachine(raw_parameters, existedMachine):
+    existedMachine = update_machine_tags(raw_parameters, existedMachine)
+    existedMachine.properties.kubernetes = update_machine_kubernetes_profile_taints_labels(
+        raw_parameters,
+        existedMachine
+    )
+
+    return existedMachine
+
+
 def constructMachine(cmd, raw_parameters, machine_name):
     machine_name = raw_parameters.get("machine_name")
     if machine_name is None:
@@ -115,6 +156,27 @@ def set_machine_kubernetes_profile(cmd, raw_parameters):
     return machineKubernetesProfile
 
 
+def update_machine_tags(raw_parameters, existedMachine):
+    tags = raw_parameters.get("tags")
+    if tags is not None and len(tags) != 0:
+        existedMachine.properties.tags = tags
+    return existedMachine
+
+
+def update_machine_kubernetes_profile_taints_labels(raw_parameters, existedMachine):
+    taints_raw = raw_parameters.get("node_taints")
+    if taints_raw is not None:
+        node_taints = [x.strip() for x in (taints_raw.split(",") if taints_raw else [])]
+        existedMachine.properties.kubernetes.node_taints = node_taints
+
+    labels_raw = raw_parameters.get("labels")
+    labels = parse_key_value_list(labels_raw)
+    if labels is not None and len(labels) != 0:
+        existedMachine.properties.kubernetes.node_labels = labels
+
+    return existedMachine.properties.kubernetes
+
+
 def set_machine_os_profile(cmd, raw_parameters):
     MachineOSProfile = cmd.get_models(
         "MachineOSProfile",