{CI} Fix extension release scripts

scripts/ci/build_ext_cmd_tree.sh

@@ -6,8 +6,6 @@ if [[ -z "$changed_content" ]]; then
     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
 
@@ -20,7 +18,8 @@ export AZURE_EXTENSION_INDEX_URL=https://raw.githubusercontent.com/Azure/azure-c
 output=$(az extension list-available --query [].name -otsv)
 # azure-cli-ml is replaced by ml
 # disable alias which relies on Jinja2 2.10
-blocklist=("azure-cli-ml" "alias")
+# disable rdbms-connect and deploy-to-azure which cause cmd tree build failures
+blocklist=("azure-cli-ml" "alias" "rdbms-connect" "deploy-to-azure")
 
 rm -f ~/.azure/extCmdTreeToUpload.json
 
@@ -40,4 +39,7 @@ for ext in $output; do
     fi
 done
 
+pip install azure-cli-core azure-cli requests azure-storage-blob==1.5.0
+# arcdata: ModuleNotFoundError: No module named 'regex._regex'
+pip install regex
 python $(cd $(dirname $0); pwd)/update_ext_cmd_tree.py $filter_exts

scripts/ci/sync_extensions.py

@@ -5,11 +5,13 @@
 # pylint: disable=line-too-long
 # pylint: disable=broad-except
 
+import json
 import os
 import re
-import json
 import subprocess
 
+from util import run_az_cmd
+
 DEFAULT_TARGET_INDEX_URL = os.getenv('AZURE_EXTENSION_TARGET_INDEX_URL')
 STORAGE_ACCOUNT = os.getenv('AZURE_EXTENSION_TARGET_STORAGE_ACCOUNT')
 STORAGE_CONTAINER = os.getenv('AZURE_EXTENSION_TARGET_STORAGE_CONTAINER')
@@ -48,7 +50,6 @@ def download_file(url, file_path):
             if chunk:  # ignore keep-alive new chunks
                 f.write(chunk)
 
-
 def _sync_wheel(ext, updated_indexes, failed_urls, overwrite, temp_dir):
     download_url = ext['downloadUrl']
     whl_file = download_url.split('/')[-1]
@@ -62,27 +63,24 @@ def _sync_wheel(ext, updated_indexes, failed_urls, overwrite, temp_dir):
     if not overwrite:
         cmd = ['az', 'storage', 'blob', 'exists', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
                f'{STORAGE_ACCOUNT}', '--name', f'{blob_name}', '--auth-mode', 'login']
-        result = subprocess.run(cmd, capture_output=True)
-        if result.stdout and json.loads(result.stdout)['exists']:
+        message = f"Checking if '{blob_name}' exists in the storage"
+        result = run_az_cmd(cmd, message=message, raise_error=True)
+        if json.loads(result.stdout)['exists']:
             print("Skipping '{}' as it already exists...".format(whl_file))
             return
 
     cmd = ['az', 'storage', 'blob', 'upload', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
            f'{STORAGE_ACCOUNT}', '--name', f'{blob_name}', '--file', f'{os.path.abspath(whl_path)}',
            '--auth-mode', 'login', '--overwrite']
-    result = subprocess.run(cmd, capture_output=True)
-    if result.returncode != 0:
-        print(f"Failed to upload '{whl_file}' to the storage account")
-        raise
+    message = f"uploading '{blob_name}' to the storage"
+    run_az_cmd(cmd, message=message, raise_error=True)
+
     cmd = ['az', 'storage', 'blob', 'url', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
            f'{STORAGE_ACCOUNT}', '--name',  f'{blob_name}', '--auth-mode', 'login']
-    result = subprocess.run(cmd, capture_output=True)
-    print(result)
-    if result.stdout and result.returncode == 0:
-        url = json.loads(result.stdout)
-    else:
-        print("Failed to get the URL for '{}'".format(whl_file))
-        raise
+    message = f"Getting the URL for '{blob_name}'"
+    result = run_az_cmd(cmd, message=message, raise_error=True)
+    url = json.loads(result.stdout)
+
     updated_index = ext
     updated_index['downloadUrl'] = url
     updated_indexes.append(updated_index)
@@ -157,10 +155,9 @@ def main():
         cmd = ['az', 'storage', 'blob', 'upload', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
                f'{STORAGE_ACCOUNT}', '--name', f'{backup_index_name}',
                '--file', f'{os.path.abspath(target_index_path)}', '--auth-mode', 'login', '--overwrite']
-        result = subprocess.run(cmd, capture_output=True)
-        if result.returncode != 0:
-            print(f"Failed to upload '{target_index_path}' to the storage account")
-            raise
+        message = f"Uploading '{backup_index_name}' to the storage"
+        run_az_cmd(cmd, message=message, raise_error=True)
+
         # start with an empty index.json to sync all extensions
         initial_index = {"extensions": {}, "formatVersion": "1"}
         open(target_index_path, 'w').write(json.dumps(initial_index, indent=4, sort_keys=True))
@@ -185,10 +182,9 @@ def main():
     cmd = ['az', 'storage', 'blob', 'upload', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
            f'{STORAGE_ACCOUNT}', '--name', f'{index_name}', '--file', f'{os.path.abspath(target_index_path)}',
            '--auth-mode', 'login', '--overwrite']
-    result = subprocess.run(cmd, capture_output=True)
-    if result.returncode != 0:
-        print(f"Failed to upload '{target_index_path}' to the storage account")
-        raise
+    message = f"Uploading '{index_name}' to the storage"
+    run_az_cmd(cmd, message=message, raise_error=True)
+
     print("\nSync finished.")
     if updated_indexes:
         print("New extensions available in:")

scripts/ci/update_ext_cmd_tree.py

@@ -6,13 +6,15 @@
 import filecmp
 import json
 import os
-import subprocess
 import sys
+import subprocess
+
 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 sync_extensions import download_file
+from util import run_az_cmd
 
 STORAGE_ACCOUNT = os.getenv('AZURE_EXTENSION_CMD_TREE_STORAGE_ACCOUNT')
 STORAGE_CONTAINER = os.getenv('AZURE_EXTENSION_CMD_TREE_STORAGE_CONTAINER')
@@ -22,6 +24,36 @@
 file_name = 'extCmdTreeToUpload.json'
 
 
+def execute_command(command):
+    """Execute a shell command and return the output."""
+    try:
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.returncode == 0:
+            return result.stdout.strip()
+        else:
+            return f"Error: {result.stderr.strip()}"
+    except Exception as e:
+        return f"Exception: {str(e)}"
+
+
+def get_package_version(package_name):
+    """Get the current version of a Python package."""
+    command = ["pip", "show", package_name]
+    output = execute_command(command)
+    if "Version:" in output:
+        for line in output.splitlines():
+            if line.startswith("Version:"):
+                version = line.split(":")[1].strip()
+                print(f"{package_name} current version: {version}")
+
+
+def upgrade_package(package_name):
+    """Upgrade a Python package to the latest version."""
+    command = ["pip", "install", "--upgrade", package_name]
+    print(f"{command}")
+    return execute_command(command)
+
+
 def merge(data, key, value):
     if isinstance(value, str):
         if key in data:
@@ -35,6 +67,16 @@ def merge(data, key, value):
 
 def update_cmd_tree(ext_name):
     print(f"Processing {ext_name}")
+    if ext_name == 'ml':
+        get_package_version("azure-storage-blob")
+        upgrade_package("azure-storage-blob")
+        get_package_version("azure-storage-blob")
+        get_package_version("rpds")
+        upgrade_package("rpds")
+        get_package_version("rpds")
+        get_package_version("rpds-py")
+        upgrade_package("rpds-py")
+        get_package_version("rpds-py")
 
     ext_dir = get_extension_path(ext_name)
     ext_mod = get_extension_modname(ext_name, ext_dir=ext_dir)
@@ -83,19 +125,14 @@ def upload_cmd_tree():
     cmd = ['az', 'storage', 'blob', 'upload', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
            f'{STORAGE_ACCOUNT}', '--name', f'{blob_file_name}', '--file', f'{file_path}', '--auth-mode', 'login',
            '--overwrite']
-    result = subprocess.run(cmd, capture_output=True)
-    if result.returncode != 0:
-        print(f"Failed to upload '{blob_file_name}' to the storage account")
-        print(result)
+    message = f"Uploading '{blob_file_name}' to the storage"
+    run_az_cmd(cmd, message=message, raise_error=True)
 
     cmd = ['az', 'storage', 'blob', 'url', '--container-name', f'{STORAGE_CONTAINER}', '--account-name',
            f'{STORAGE_ACCOUNT}', '--name', f'{blob_file_name}', '--auth-mode', 'login']
-    result = subprocess.run(cmd, capture_output=True)
-    if result.stdout and result.returncode == 0:
-        url = json.loads(result.stdout)
-    else:
-        print(f"Failed to get the URL for '{blob_file_name}'")
-        raise
+    message = f"Getting the URL for '{blob_file_name}'"
+    result = run_az_cmd(cmd, message=message, raise_error=True)
+    url = json.loads(result.stdout)
 
     download_file_path = os.path.expanduser(os.path.join('~', '.azure', downloaded_file_name))
     download_file(url, download_file_path)

scripts/ci/util.py

@@ -3,15 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+import json
 import logging
 import os
 import re
 import shlex
-import json
+import subprocess
 import zipfile
 
-from subprocess import check_output
-
 logger = logging.getLogger(__name__)
 
 # copy from wheel==0.30.0
@@ -136,7 +135,7 @@ def diff_code(start, end):
 
         # If running in Travis CI, only run tests for edited extensions
         commit_range = os.environ.get('TRAVIS_COMMIT_RANGE')
-        if commit_range and not check_output(
+        if commit_range and not subprocess.check_output(
                 ['git', '--no-pager', 'diff', '--name-only', commit_range, '--', src_d_full]):
             continue
 
@@ -154,7 +153,7 @@ def diff_code(start, end):
             else:
                 cmd = cmd_tpl.format(start=start, end=end,
                                      code_dir=src_d_full)
-                if not check_output(shlex.split(cmd)):
+                if not subprocess.check_output(shlex.split(cmd)):
                     continue
 
         diff_ref.append((pkg_name, src_d_full))
@@ -163,3 +162,26 @@ def diff_code(start, end):
                    f'end: {end}, '
                    f'diff_ref: {diff_ref}.')
     return diff_ref
+
+
+def run_az_cmd(cmd, message=False, raise_error=True):
+    """
+    :param cmd: The entire command line to run.
+    :param message: A custom message to print, or True (bool) to use a default.
+    :param raise_error: Whether to raise an exception if the command fails.
+    """
+    # use default message if custom not provided
+    if message is True:
+        print(f'Running: {cmd}')
+
+    if message:
+        print(f'{message}')
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, check=True)
+        return result
+    except subprocess.CalledProcessError as ex:
+        error_message = ex.stderr if ex.stderr else str(ex)
+        if raise_error:
+            raise Exception(f"Failed to run command: {cmd}\nError: {error_message}") from ex
+        return ex