Build extension command tree (Azure#2078)

scripts/ci/build_ext_cmd_tree.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+changed_content=$(git --no-pager diff --diff-filter=ACMRT HEAD~$AZURE_EXTENSION_COMMIT_NUM -- src/index.json)
+if [[ -z "$changed_content" ]]; then
+    echo "index.json not modified. End task."
+    exit 0
+fi
+
+pip install azure-cli-core azure-cli requests
+pip install azure-storage-blob==1.5.0
+echo "Listing Available Extensions:"
+az extension list-available -otable
+
+# turn off telemetry as it crowds output
+export AZURE_CORE_COLLECT_TELEMETRY=False
+
+# wait for the index.json to be synced in storage account
+# Remove this when we can support using customized index.json
+sleep 360
+
+output=$(az extension list-available --query [].name -otsv)
+# azure-cli-iot-ext is the deprecated old versions of the renamed azure-iot extension
+blocklist=("azure-cli-iot-ext")
+
+rm -f ~/.azure/extCmdTreeToUpload.json
+
+filter_exts=""
+for ext in $output; do
+    ext=${ext%$'\r'}   # Remove a trailing newline when running on Windows.
+    if [[ " ${blocklist[@]} " =~ " ${ext} " ]]; then
+        continue
+    fi
+    filter_exts="${filter_exts} ${ext}"
+    echo "Adding extension:" $ext
+    az extension add -n $ext
+    if [ $? != 0 ]
+    then
+        echo "Failed to load:" $ext
+        exit 1
+    fi
+done
+
+python $(cd $(dirname $0); pwd)/update_ext_cmd_tree.py $filter_exts

scripts/ci/sync_extensions.py

@@ -26,7 +26,7 @@ def _get_updated_extension_filenames():
     return added_ext_filenames, deleted_ext_filenames
 
 
-def _download_file(url, file_path):
+def download_file(url, file_path):
     import requests
     count = 3
     the_ex = None
@@ -54,7 +54,7 @@ def _sync_wheel(ext, updated_indexes, failed_urls, client, overwrite, temp_dir):
     whl_file = download_url.split('/')[-1]
     whl_path = os.path.join(temp_dir, whl_file)
     try:
-        _download_file(download_url, whl_path)
+        download_file(download_url, whl_path)
     except Exception:
         failed_urls.append(download_url)
         return
@@ -119,7 +119,7 @@ def main():
     target_index = DEFAULT_TARGET_INDEX_URL
     os.mkdir(os.path.join(temp_dir, 'target'))
     target_index_path = os.path.join(temp_dir, 'target', 'index.json')
-    _download_file(target_index, target_index_path)
+    download_file(target_index, target_index_path)
 
     client = BlockBlobService(account_name=STORAGE_ACCOUNT, account_key=STORAGE_ACCOUNT_KEY)
     updated_indexes = []

scripts/ci/update_ext_cmd_tree.py

@@ -0,0 +1,91 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import filecmp
+import os
+import sys
+from azure.cli.core import get_default_cli
+from azure.cli.core._session import Session
+from azure.cli.core.commands import _load_extension_command_loader
+from azure.cli.core.extension import get_extension_modname, get_extension_path
+from azure.storage.blob import BlockBlobService
+from sync_extensions import download_file
+
+STORAGE_ACCOUNT_KEY = os.getenv('AZURE_EXTENSION_CMD_TREE_STORAGE_ACCOUNT_KEY')
+STORAGE_ACCOUNT = os.getenv('AZURE_EXTENSION_CMD_TREE_STORAGE_ACCOUNT')
+STORAGE_CONTAINER = os.getenv('AZURE_EXTENSION_CMD_TREE_STORAGE_CONTAINER')
+
+az_cli = get_default_cli()
+file_name = 'extCmdTreeToUpload.json'
+
+
+def merge(data, key, value):
+    if isinstance(value, str):
+        if key in data:
+            raise Exception(f"Key: {key} already exists. 2 extensions cannot have the same command!")
+        data[key] = value
+    else:
+        data.setdefault(key, {})
+        for k, v in value.items():
+            merge(data[key], k, v)
+
+
+def update_cmd_tree(ext_name):
+    print(f"Processing {ext_name}")
+
+    ext_dir = get_extension_path(ext_name)
+    ext_mod = get_extension_modname(ext_name, ext_dir=ext_dir)
+
+    invoker = az_cli.invocation_cls(cli_ctx=az_cli, commands_loader_cls=az_cli.commands_loader_cls,
+                                    parser_cls=az_cli.parser_cls, help_cls=az_cli.help_cls)
+    az_cli.invocation = invoker
+
+    sys.path.append(ext_dir)
+    extension_command_table, _ = _load_extension_command_loader(invoker.commands_loader,
+                                                                "", ext_mod)
+
+    EXT_CMD_TREE_TO_UPLOAD = Session()
+    EXT_CMD_TREE_TO_UPLOAD.load(os.path.expanduser(os.path.join('~', '.azure', file_name)))
+    root = {}
+    for cmd_name, _ in extension_command_table.items():
+        parts = cmd_name.split()
+        parent = root
+        for i, part in enumerate(parts):
+            if part in parent:
+                pass
+            elif i == len(parts) - 1:
+                parent[part] = ext_name
+            else:
+                parent[part] = {}
+            parent = parent[part]
+    print(root)
+    for k, v in root.items():
+        merge(EXT_CMD_TREE_TO_UPLOAD.data, k, v)
+    EXT_CMD_TREE_TO_UPLOAD.save_with_retry()
+
+
+def upload_cmd_tree():
+    blob_file_name = 'extensionCommandTree.json'
+    downloaded_file_name = 'extCmdTreeDownloaded.json'
+    file_path = os.path.expanduser(os.path.join('~', '.azure', file_name))
+
+    client = BlockBlobService(account_name=STORAGE_ACCOUNT, account_key=STORAGE_ACCOUNT_KEY)
+    client.create_blob_from_path(container_name=STORAGE_CONTAINER, blob_name=blob_file_name,
+                                 file_path=file_path)
+
+    url = client.make_blob_url(container_name=STORAGE_CONTAINER, blob_name=blob_file_name)
+
+    download_file_path = os.path.expanduser(os.path.join('~', '.azure', downloaded_file_name))
+    download_file(url, download_file_path)
+    if filecmp.cmp(file_path, download_file_path):
+        print("extensionCommandTree.json uploaded successfully. URL: {}".format(url))
+    else:
+        raise Exception("Failed to update extensionCommandTree.json in the storage account")
+
+
+if __name__ == '__main__':
+    for ext in sys.argv[1:]:
+        update_cmd_tree(ext)
+    upload_cmd_tree()