[Extension] Support automatically installing an extension if the extension of a command is not installed

src/azure-cli-core/azure/cli/core/__init__.py

@@ -373,18 +373,6 @@ def _get_extension_suppressions(mod_loaders):
                             res.append(sup)
             return res
 
-        def _roughly_parse_command(args):
-            # Roughly parse the command part: <az vm create> --name vm1
-            # Similar to knack.invocation.CommandInvoker._rudimentary_get_command, but we don't need to bother with
-            # positional args
-            nouns = []
-            for arg in args:
-                if arg and arg[0] != '-':
-                    nouns.append(arg)
-                else:
-                    break
-            return ' '.join(nouns).lower()
-
         # Clear the tables to make this method idempotent
         self.command_group_table.clear()
         self.command_table.clear()
@@ -404,8 +392,9 @@ def _roughly_parse_command(args):
                 _update_command_table_from_extensions([], index_extensions)
 
                 logger.debug("Loaded %d groups, %d commands.", len(self.command_group_table), len(self.command_table))
+                from azure.cli.core.util import roughly_parse_command
                 # The index may be outdated. Make sure the command appears in the loaded command table
-                command_str = _roughly_parse_command(args)
+                command_str = roughly_parse_command(args)
                 if command_str in self.command_table:
                     logger.debug("Found a match in the command table for '%s'", command_str)
                     return self.command_table

src/azure-cli-core/azure/cli/core/_session.py

@@ -113,3 +113,6 @@ def __len__(self):
 # it could be lagged behind and can be used to check whether
 # an upgrade of azure-cli happens
 VERSIONS = Session()
+
+# EXT_CMD_TREE provides command to extension name mapping
+EXT_CMD_TREE = Session()

src/azure-cli-core/azure/cli/core/extension/operations.py

@@ -85,8 +85,8 @@ def _validate_whl_extension(ext_file):
     check_version_compatibility(azext_metadata)
 
 
-def _add_whl_ext(cmd, source, ext_sha256=None, pip_extra_index_urls=None, pip_proxy=None, system=None):  # pylint: disable=too-many-statements
-    cmd.cli_ctx.get_progress_controller().add(message='Analyzing')
+def _add_whl_ext(cli_ctx, source, ext_sha256=None, pip_extra_index_urls=None, pip_proxy=None, system=None):  # pylint: disable=too-many-statements
+    cli_ctx.get_progress_controller().add(message='Analyzing')
     if not source.endswith('.whl'):
         raise ValueError('Unknown extension type. Only Python wheels are supported.')
     url_parse_result = urlparse(source)
@@ -108,7 +108,7 @@ def _add_whl_ext(cmd, source, ext_sha256=None, pip_extra_index_urls=None, pip_pr
         logger.debug('Downloading %s to %s', source, ext_file)
         import requests
         try:
-            cmd.cli_ctx.get_progress_controller().add(message='Downloading')
+            cli_ctx.get_progress_controller().add(message='Downloading')
             _whl_download_from_url(url_parse_result, ext_file)
         except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError) as err:
             raise CLIError('Please ensure you have network connection. Error detail: {}'.format(str(err)))
@@ -130,7 +130,7 @@ def _add_whl_ext(cmd, source, ext_sha256=None, pip_extra_index_urls=None, pip_pr
             raise CLIError("The checksum of the extension does not match the expected value. "
                            "Use --debug for more information.")
     try:
-        cmd.cli_ctx.get_progress_controller().add(message='Validating')
+        cli_ctx.get_progress_controller().add(message='Validating')
         _validate_whl_extension(ext_file)
     except AssertionError:
         logger.debug(traceback.format_exc())
@@ -140,7 +140,7 @@ def _add_whl_ext(cmd, source, ext_sha256=None, pip_extra_index_urls=None, pip_pr
     logger.debug('Validation successful on %s', ext_file)
     # Check for distro consistency
     check_distro_consistency()
-    cmd.cli_ctx.get_progress_controller().add(message='Installing')
+    cli_ctx.get_progress_controller().add(message='Installing')
     # Install with pip
     extension_path = build_extension_path(extension_name, system)
     pip_args = ['install', '--target', extension_path, ext_file]
@@ -206,15 +206,15 @@ def check_version_compatibility(azext_metadata):
         raise CLIError(min_max_msg_fmt)
 
 
-def add_extension(cmd, source=None, extension_name=None, index_url=None, yes=None,  # pylint: disable=unused-argument
+def add_extension(cmd=None, source=None, extension_name=None, index_url=None, yes=None,  # pylint: disable=unused-argument
                   pip_extra_index_urls=None, pip_proxy=None, system=None,
-                  version=None):
+                  version=None, cli_ctx=None):
     ext_sha256 = None
 
     version = None if version == 'latest' else version
-
+    cmd_cli_ctx = cli_ctx or cmd.cli_ctx
     if extension_name:
-        cmd.cli_ctx.get_progress_controller().add(message='Searching')
+        cmd_cli_ctx.get_progress_controller().add(message='Searching')
         ext = None
         try:
             ext = get_extension(extension_name)
@@ -236,7 +236,7 @@ def add_extension(cmd, source=None, extension_name=None, index_url=None, yes=Non
                 err = "No matching extensions for '{}'. Use --debug for more information.".format(extension_name)
             raise CLIError(err)
 
-    extension_name = _add_whl_ext(cmd=cmd, source=source, ext_sha256=ext_sha256,
+    extension_name = _add_whl_ext(cli_ctx=cmd_cli_ctx, source=source, ext_sha256=ext_sha256,
                                   pip_extra_index_urls=pip_extra_index_urls, pip_proxy=pip_proxy, system=system)
     try:
         ext = get_extension(extension_name)
@@ -289,8 +289,9 @@ def show_extension(extension_name):
         raise CLIError(e)
 
 
-def update_extension(cmd, extension_name, index_url=None, pip_extra_index_urls=None, pip_proxy=None):
+def update_extension(cmd=None, extension_name=None, index_url=None, pip_extra_index_urls=None, pip_proxy=None, cli_ctx=None):
     try:
+        cmd_cli_ctx = cli_ctx or cmd.cli_ctx
         ext = get_extension(extension_name, ext_type=WheelExtension)
         cur_version = ext.get_version()
         try:
@@ -307,7 +308,7 @@ def update_extension(cmd, extension_name, index_url=None, pip_extra_index_urls=N
         shutil.rmtree(extension_path)
         # Install newer version
         try:
-            _add_whl_ext(cmd=cmd, source=download_url, ext_sha256=ext_sha256,
+            _add_whl_ext(cli_ctx=cmd_cli_ctx, source=download_url, ext_sha256=ext_sha256,
                          pip_extra_index_urls=pip_extra_index_urls, pip_proxy=pip_proxy)
             logger.debug('Deleting backup of old extension at %s', backup_dir)
             shutil.rmtree(backup_dir)

src/azure-cli-core/azure/cli/core/parser.py

@@ -280,37 +280,167 @@ def parse_known_args(self, args=None, namespace=None):
         self._namespace, self._raw_arguments = super().parse_known_args(args=args, namespace=namespace)
         return self._namespace, self._raw_arguments
 
-    def _check_value(self, action, value):
+    def _get_extension_command_tree(self):
+        from azure.cli.core._session import EXT_CMD_TREE
+        import os
+        VALID_SECOND = 3600 * 24 * 10
+        # self.cli_ctx is None when self.prog is beyond 'az', such as 'az iot'.
+        # use cli_ctx from cli_help which is not lost.
+        cli_ctx = self.cli_ctx or (self.cli_help.cli_ctx if self.cli_help else None)
+        if not cli_ctx:
+            return None
+        EXT_CMD_TREE.load(os.path.join(cli_ctx.config.config_dir, 'extensionCommandTree.json'), VALID_SECOND)
+        if not EXT_CMD_TREE.data:
+            import requests
+            from azure.cli.core.util import should_disable_connection_verify
+            try:
+                response = requests.get(
+                    'https://azurecliextensionsync.blob.core.windows.net/cmd-index/extensionCommandTree.json',
+                    verify=(not should_disable_connection_verify()),
+                    timeout=300)
+            except Exception as ex:  # pylint: disable=broad-except
+                logger.info("Request failed for extension command tree: %s", str(ex))
+                return None
+            if response.status_code == 200:
+                EXT_CMD_TREE.data = response.json()
+                EXT_CMD_TREE.save_with_retry()
+            else:
+                logger.info("Error when retrieving extension command tree. Response code: %s", response.status_code)
+                return None
+        return EXT_CMD_TREE
+
+    def _search_in_extension_commands(self, command_str):
+        """Search the command in an extension commands dict which mimics a prefix tree.
+        If the value of the dict item is a string, then the key represents the end of a complete command
+        and the value is the name of the extension that the command belongs to.
+        An example of the dict read from extensionCommandTree.json:
+        {
+            "aks": {
+                "create": "aks-preview",
+                "update": "aks-preview",
+                "app": {
+                    "up": "deploy-to-azure"
+                },
+                "use-dev-spaces": "dev-spaces"
+            },
+            ...
+        }
+        """
+
+        cmd_chain = self._get_extension_command_tree()
+        for part in command_str.split():
+            try:
+                if isinstance(cmd_chain[part], str):
+                    return cmd_chain[part]
+                cmd_chain = cmd_chain[part]
+            except KeyError:
+                return None
+        return None
+
+    def _get_extension_use_dynamic_install_config(self):
+        cli_ctx = self.cli_ctx or (self.cli_help.cli_ctx if self.cli_help else None)
+        use_dynamic_install = cli_ctx.config.get(
+            'extension', 'use_dynamic_install', 'no').lower() if cli_ctx else 'no'
+        if use_dynamic_install not in ['no', 'yes_prompt', 'yes_without_prompt']:
+            use_dynamic_install = 'no'
+        return use_dynamic_install
+
+    def _check_value(self, action, value):  # pylint: disable=too-many-statements, too-many-locals
         # Override to customize the error message when a argument is not among the available choices
         # converted value must be one of the choices (if specified)
-        if action.choices is not None and value not in action.choices:
+        if action.choices is not None and value not in action.choices:  # pylint: disable=too-many-nested-blocks
+            caused_by_extension_not_installed = False
             if not self.command_source:
-                # parser has no `command_source`, value is part of command itself
-                extensions_link = 'https://docs.microsoft.com/en-us/cli/azure/azure-cli-extensions-overview'
-                error_msg = ("{prog}: '{value}' is not in the '{prog}' command group. See '{prog} --help'. "
-                             "If the command is from an extension, "
-                             "please make sure the corresponding extension is installed. "
-                             "To learn more about extensions, please visit "
-                             "{extensions_link}").format(prog=self.prog, value=value, extensions_link=extensions_link)
+                candidates = difflib.get_close_matches(value, action.choices, cutoff=0.7)
+                error_msg = None
+                # self.cli_ctx is None when self.prog is beyond 'az', such as 'az iot'.
+                # use cli_ctx from cli_help which is not lost.
+                cli_ctx = self.cli_ctx or (self.cli_help.cli_ctx if self.cli_help else None)
+                use_dynamic_install = self._get_extension_use_dynamic_install_config()
+                if use_dynamic_install != 'no' and not candidates:
+                    # Check if the command is from an extension
+                    from azure.cli.core.util import roughly_parse_command
+                    cmd_list = self.prog.split() + self._raw_arguments
+                    command_str = roughly_parse_command(cmd_list[1:])
+                    ext_name = self._search_in_extension_commands(command_str)
+                    if ext_name:
+                        caused_by_extension_not_installed = True
+                        telemetry.set_command_details(command_str,
+                                                      parameters=AzCliCommandInvoker._extract_parameter_names(cmd_list),  # pylint: disable=protected-access
+                                                      extension_name=ext_name)
+                        run_after_extension_installed = cli_ctx.config.getboolean('extension',
+                                                                                  'run_after_dynamic_install',
+                                                                                  False) if cli_ctx else False
+                        if use_dynamic_install == 'yes_without_prompt':
+                            logger.warning('The command requires the extension %s. '
+                                           'It will be installed first.', ext_name)
+                            go_on = True
+                        else:
+                            from knack.prompting import prompt_y_n, NoTTYException
+                            prompt_msg = 'The command requires the extension {}. ' \
+                                'Do you want to install it now?'.format(ext_name)
+                            if run_after_extension_installed:
+                                prompt_msg = '{} The command will continue to run after the extension is installed.' \
+                                    .format(prompt_msg)
+                            NO_PROMPT_CONFIG_MSG = "Run 'az config set extension.use_dynamic_install=" \
+                                "yes_without_prompt' to allow installing extensions without prompt."
+                            try:
+                                go_on = prompt_y_n(prompt_msg, default='y')
+                                if go_on:
+                                    logger.warning(NO_PROMPT_CONFIG_MSG)
+                            except NoTTYException:
+                                logger.warning("The command requires the extension %s.\n "
+                                               "Unable to prompt for extension install confirmation as no tty "
+                                               "available. %s", ext_name, NO_PROMPT_CONFIG_MSG)
+                                go_on = False
+                        if go_on:
+                            from azure.cli.core.extension.operations import add_extension
+                            add_extension(cli_ctx=cli_ctx, extension_name=ext_name)
+                            if run_after_extension_installed:
+                                import subprocess
+                                import platform
+                                exit_code = subprocess.call(cmd_list, shell=platform.system() == 'Windows')
+                                telemetry.set_user_fault("Extension {} dynamically installed and commands will be "
+                                                         "rerun automatically.".format(ext_name))
+                                self.exit(exit_code)
+                            else:
+                                error_msg = 'Extension {} installed. Please rerun your command.'.format(ext_name)
+                        else:
+                            error_msg = "The command requires the extension {ext_name}. " \
+                                "To install, run 'az extension add -n {ext_name}'.".format(ext_name=ext_name)
+                if not error_msg:
+                    # parser has no `command_source`, value is part of command itself
+                    error_msg = "{prog}: '{value}' is not in the '{prog}' command group. See '{prog} --help'." \
+                        .format(prog=self.prog, value=value)
+                    if use_dynamic_install.lower() == 'no':
+                        extensions_link = 'https://docs.microsoft.com/en-us/cli/azure/azure-cli-extensions-overview'
+                        error_msg = ("{msg} "
+                                     "If the command is from an extension, "
+                                     "please make sure the corresponding extension is installed. "
+                                     "To learn more about extensions, please visit "
+                                     "{extensions_link}").format(msg=error_msg, extensions_link=extensions_link)
             else:
                 # `command_source` indicates command values have been parsed, value is an argument
                 parameter = action.option_strings[0] if action.option_strings else action.dest
                 error_msg = "{prog}: '{value}' is not a valid value for '{param}'. See '{prog} --help'.".format(
                     prog=self.prog, value=value, param=parameter)
+                candidates = difflib.get_close_matches(value, action.choices, cutoff=0.7)
+
             telemetry.set_user_fault(error_msg)
             with CommandLoggerContext(logger):
                 logger.error(error_msg)
-            candidates = difflib.get_close_matches(value, action.choices, cutoff=0.7)
-            if candidates:
-                print_args = {
-                    's': 's' if len(candidates) > 1 else '',
-                    'verb': 'are' if len(candidates) > 1 else 'is',
-                    'value': value
-                }
-                self._suggestion_msg.append("\nThe most similar choice{s} to '{value}' {verb}:".format(**print_args))
-                self._suggestion_msg.append('\n'.join(['\t' + candidate for candidate in candidates]))
-
-            failure_recovery_recommendations = self._get_failure_recovery_recommendations(action)
-            self._suggestion_msg.extend(failure_recovery_recommendations)
-            self._print_suggestion_msg(sys.stderr)
+            if not caused_by_extension_not_installed:
+                if candidates:
+                    print_args = {
+                        's': 's' if len(candidates) > 1 else '',
+                        'verb': 'are' if len(candidates) > 1 else 'is',
+                        'value': value
+                    }
+                    self._suggestion_msg.append("\nThe most similar choice{s} to '{value}' {verb}:"
+                                                .format(**print_args))
+                    self._suggestion_msg.append('\n'.join(['\t' + candidate for candidate in candidates]))
+
+                failure_recovery_recommendations = self._get_failure_recovery_recommendations(action)
+                self._suggestion_msg.extend(failure_recovery_recommendations)
+                self._print_suggestion_msg(sys.stderr)
             self.exit(2)