Audapure/wo cli improvements

src/workload-orchestration/HISTORY.rst

@@ -2,6 +2,11 @@
 
 Release History
 ===============
+5.2.1
+++++++
+* **Auto-install of required az CLI dependencies** — ``az workload-orchestration cluster init`` (and any flow that invokes ``target_prepare``) now performs a pre-flight check for the ``connectedk8s``, ``k8s-extension``, and ``customlocation`` az CLI extensions, and installs any that are missing. Previously, missing dependencies surfaced as opaque ``command not recognized`` errors deep inside the cluster onboarding flow.
+* **Cleaner output for ``az workload-orchestration context capability-add``** — removed redundant ``Adding N: <names>`` log line that duplicated information already shown in the ``✓ Done (N total capabilities)`` summary.
+
 5.2.0
 ++++++
 * **CLI Onboarding Simplification** — reduces onboarding from 11 commands to 4:

src/workload-orchestration/azext_workload_orchestration/common/context.py

@@ -288,13 +288,10 @@ def capability_add(cli_ctx, resource_group, context_name, name=None,
         return ctx
 
     merged = existing + added
-    names_str = ", ".join(c["name"] for c in added)
-    _log(f"Adding {len(added)}: {names_str}")
 
     updated = _patch_context_capabilities(
         cli_ctx, sub_id, resource_group, context_name, merged
     )
-    _log(f"\u2713 Done ({len(merged)} total capabilities)")
     return updated
 
 
@@ -344,13 +341,9 @@ def capability_remove(cli_ctx, resource_group, context_name, name=None,
                 "Use --yes to confirm removal in non-interactive sessions."
             ) from exc
 
-    names_str = ", ".join(c["name"] for c in to_remove)
-    _log(f"Removing {len(to_remove)}: {names_str}")
-
     updated = _patch_context_capabilities(
         cli_ctx, sub_id, resource_group, context_name, remaining
     )
-    _log(f"\u2713 Done ({len(remaining)} total capabilities)")
     return updated
 
 

src/workload-orchestration/azext_workload_orchestration/common/hierarchy.py

@@ -277,8 +277,10 @@ def _create_sg_level(  # pylint: disable=too-many-arguments
 
     # 1. Create ServiceGroup
     _eprint(f"{parent_prefix}{connector}{name} ({level})")
+    sg_url = f"{get_arm_endpoint(cmd)}{sg_id}"
+    sg_existed = _arm_get(cmd, sg_url, SERVICE_GROUP_API_VERSION) is not None
     try:
-        _arm_put(cmd, f"{get_arm_endpoint(cmd)}{sg_id}", {
+        _arm_put(cmd, sg_url, {
             "properties": {
                 "displayName": name,
                 "parent": {"resourceId": parent_id},
@@ -359,9 +361,11 @@ def _create_sg_level(  # pylint: disable=too-many-arguments
 
     children = node.get("children")
     has_children = children is not None
+    sg_label = "(reused) " if sg_existed else "(created) "
     site_label = "(reused) " if existing_sg_site else ""
     config_label = "(reused) " if config_reused else ""
     ref_label = "(reused) " if existing_ref else ""
+    _eprint(f"{child_prefix}├── ServiceGroup '{name}' {sg_label}✓")
     _eprint(f"{child_prefix}├── Site '{effective_site_name}' {site_label}✓")
     _eprint(f"{child_prefix}├── Configuration '{config_name}' {config_label}✓")
     if has_children:

src/workload-orchestration/azext_workload_orchestration/common/target.py

@@ -40,6 +40,7 @@
 )
 from azext_workload_orchestration.common.utils import (
     _eprint,
+    ensure_required_cli_extensions,
     invoke_cli_command,
 )
 
@@ -84,6 +85,17 @@ def target_prepare(
 
     step_results = {}
 
+    # Step 0: Ensure required az CLI extensions are installed
+    # (connectedk8s, k8s-extension, customlocation are called via invoke_cli_command
+    # and would fail with opaque errors if missing)
+    try:
+        ensure_required_cli_extensions()
+        step_results["cli-extensions"] = "Ready"
+    except Exception as exc:
+        step_results["cli-extensions"] = f"FAILED: {exc}"
+        _print_failure_hint(step_results)
+        raise
+
     try:
         connected_cluster_id = _preflight_checks(cmd, cluster_name, resource_group)
         step_results["preflight"] = "Passed"

src/workload-orchestration/azext_workload_orchestration/common/utils.py

@@ -156,3 +156,74 @@ def print_step(step_num, total, message, status=""):
         _eprint(f"{connector} {message} {status}")
     else:
         _eprint(f"{connector} {message}...")
+
+
+# ---------------------------------------------------------------------------
+# CLI extension dependency check
+# ---------------------------------------------------------------------------
+
+# az CLI extensions that workload-orchestration calls at runtime via
+# invoke_cli_command. These MUST be installed before `cluster init` runs,
+# otherwise sub-command invocations fail with opaque "command not recognized"
+# errors. Mirrors the azext_vme.utils.check_and_add_cli_extension pattern.
+REQUIRED_CLI_EXTENSIONS = [
+    "connectedk8s",     # `connectedk8s show` in _preflight_checks
+    "k8s-extension",    # `k8s-extension list/create` for aio-certmgr + wo-extension
+    "customlocation",   # `customlocation create` for Step 4
+]
+
+
+def check_and_add_cli_extension(extension_name):
+    """Check if an az CLI extension is installed; install it if missing.
+
+    Uses subprocess (not invoke_cli_command) so that `az extension add`
+    runs in its own CLI process — needed because in-process invoke does
+    not pick up newly installed extensions in the same Python process.
+
+    Raises CLIInternalError if the install fails.
+    """
+    import shutil
+    import subprocess
+
+    az = shutil.which("az") or "az"
+
+    # Check if installed
+    try:
+        result = subprocess.run(
+            [az, "extension", "list",
+             "--query", f"[?name=='{extension_name}'].name",
+             "-o", "tsv"],
+            capture_output=True, text=True, check=True, encoding="utf-8"
+        )
+        if extension_name in (result.stdout or "").strip():
+            logger.debug("az cli extension '%s' already installed", extension_name)
+            return
+    except subprocess.CalledProcessError as exc:
+        raise CLIInternalError(
+            f"Failed to check for az cli extension '{extension_name}': {exc.stderr or exc}"
+        ) from None
+
+    _eprint(f"  ├── Installing required az cli extension: {extension_name}...")
+    try:
+        subprocess.run(
+            [az, "extension", "add", "--name", extension_name, "--yes"],
+            capture_output=True, text=True, check=True, encoding="utf-8"
+        )
+        _eprint(f"  ├── Installed az cli extension: {extension_name} ✓")
+    except subprocess.CalledProcessError as exc:
+        raise CLIInternalError(
+            f"Failed to install required az cli extension '{extension_name}': "
+            f"{exc.stderr or exc}\n"
+            f"Please install manually: az extension add --name {extension_name}"
+        ) from None
+
+
+def ensure_required_cli_extensions(extension_names=None):
+    """Ensure all required az CLI extensions are installed.
+
+    Defaults to REQUIRED_CLI_EXTENSIONS. Idempotent — skips already-installed
+    extensions. Fails fast with a clear message if any install fails.
+    """
+    extensions = extension_names if extension_names is not None else REQUIRED_CLI_EXTENSIONS
+    for ext in extensions:
+        check_and_add_cli_extension(ext)

src/workload-orchestration/setup.py

@@ -10,7 +10,7 @@
 
 
 # HISTORY.rst entry.
-VERSION = '5.2.0'
+VERSION = '5.2.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers