feat: Szupport az aks loadbalancer rebalance-nodes command

Author: nilo19

src/aks-preview/HISTORY.rst

@@ -11,7 +11,11 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+
+13.0.0b9
++++++++
 * Vendor new SDK and bump API version to 2024-10-02-preview.
+* Support `az aks loadbalancer show/list/add/update/delete/rebalance` commands.
 
 13.0.0b8
 +++++++

src/aks-preview/azext_aks_preview/_params.py

@@ -2323,7 +2323,32 @@ def load_arguments(self, _):
             ),
         )
 
-    # Use same parameters for update command
+    with self.argument_context("aks loadbalancer rebalance-nodes") as c:
+        c.argument(
+            "resource_group_name",
+            options_list=["--resource-group", "-g"],
+            help="Name of resource group.",
+            id_part="resource_group",
+            configured_default="aks",
+        )
+        c.argument(
+            "cluster_name",
+            options_list=["--name", "-n"],
+            help="Name of the managed cluster.",
+        )
+        c.argument(
+            "load_balancer_names",
+            options_list=["--load-balancer-names", "--lb-names"],
+            nargs="+",
+            help=(
+                "Space-separated list of load balancer names to rebalance. "
+                "If not specified, all load balancers will be rebalanced."
+            ),
+        )
+        c.argument(
+            "no_wait", help="Do not wait for the long-running operation to finish."
+        )
+
     with self.argument_context("aks loadbalancer update") as c:
         c.argument(
             "name",

src/aks-preview/azext_aks_preview/commands.py

@@ -207,6 +207,7 @@ def load_command_table(self, _):
         g.custom_command("add", "aks_loadbalancer_add")
         g.custom_command("update", "aks_loadbalancer_update")
         g.custom_command("delete", "aks_loadbalancer_delete")
+        g.custom_command("rebalance-nodes", "aks_loadbalancer_rebalance_nodes")
 
     # AKS addon commands
     with self.command_group(
@@ -261,8 +262,12 @@ def load_command_table(self, _):
         "aks nodepool manual-scale", managed_clusters_sdk, client_factory=cf_agent_pools
     ) as g:
         g.custom_command("add", "aks_agentpool_manual_scale_add", supports_no_wait=True)
-        g.custom_command("update", "aks_agentpool_manual_scale_update", supports_no_wait=True)
-        g.custom_command("delete", "aks_agentpool_manual_scale_delete", supports_no_wait=True)
+        g.custom_command(
+            "update", "aks_agentpool_manual_scale_update", supports_no_wait=True
+        )
+        g.custom_command(
+            "delete", "aks_agentpool_manual_scale_delete", supports_no_wait=True
+        )
 
     with self.command_group(
         "aks machine", machines_sdk, client_factory=cf_machines
@@ -278,10 +283,14 @@ def load_command_table(self, _):
         "aks operation", operations_sdk, client_factory=cf_operations
     ) as g:
         g.custom_show_command(
-            "show", "aks_operation_show", table_transformer=aks_operation_show_table_format
+            "show",
+            "aks_operation_show",
+            table_transformer=aks_operation_show_table_format,
         )
         g.custom_command(
-            "show-latest", "aks_operation_show_latest", table_transformer=aks_operation_show_table_format
+            "show-latest",
+            "aks_operation_show_latest",
+            table_transformer=aks_operation_show_table_format,
         )
 
     # AKS draft commands
@@ -400,7 +409,7 @@ def load_command_table(self, _):
             "aks_mesh_disable",
             supports_no_wait=True,
             confirmation="Existing Azure Service Mesh Profile values will be reset.\n"
-            + "Are you sure you want to perform this operation?"
+            + "Are you sure you want to perform this operation?",
         )
         g.custom_command(
             "enable-ingress-gateway",

src/aks-preview/azext_aks_preview/custom.py

@@ -4180,3 +4180,44 @@ def aks_loadbalancer_show(cmd, client, resource_group_name, cluster_name, name):
     :type name: str
     """
     return client.get(resource_group_name, cluster_name, name)
+
+
+# pylint: disable=unused-argument
+def aks_loadbalancer_rebalance_nodes(
+    cmd,
+    client,
+    resource_group_name,
+    cluster_name,
+    load_balancer_names=None,
+):
+    """Rebalance nodes across specific load balancers.
+
+    :param cmd: Command context
+    :param client: AKS client
+    :param resource_group_name: Name of resource group.
+    :type resource_group_name: str
+    :param cluster_name: Name of the managed cluster.
+    :type cluster_name: str
+    :param load_balancer_names: Names of load balancers to rebalance.
+        If not specified, all load balancers will be rebalanced.
+    :type load_balancer_names: list[str]
+    :param no_wait: Do not wait for the long-running operation to finish.
+    :type no_wait: bool
+    :return: The result of the rebalance operation
+    """
+    from azext_aks_preview.loadbalancerconfiguration import (
+        aks_loadbalancer_rebalance_internal,
+    )
+    from azext_aks_preview._client_factory import cf_load_balancers
+
+    # Get the load balancers client
+    load_balancers_client = cf_load_balancers(cmd.cli_ctx)
+
+    # Prepare parameters for the internal function
+    parameters = {
+        "resource_group_name": resource_group_name,
+        "cluster_name": cluster_name,
+        "load_balancer_names": load_balancer_names,
+    }
+
+    return aks_loadbalancer_rebalance_internal(cmd, load_balancers_client, parameters)