feature: add rebuild-index command

Author: DanielMicrosoft

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

@@ -517,7 +517,7 @@ def _get_extension_suppressions(mod_loaders):
             else:
                 logger.debug("No module found from index for '%s'", args)
                 if command_index_is_valid and command_index_has_entries and args and not args[0].startswith('-') and \
-                        not self.cli_ctx.data['completer_active']:
+                    not self.cli_ctx.data['completer_active'] and args[0] != 'rebuild-index':
                     logger.debug("Valid command index has no entry for '%s'. Skip full reload and fail fast.",
                                  args[0])
                     return self.command_table
@@ -895,6 +895,28 @@ def invalidate(self):
         self.INDEX[self._HELP_INDEX] = {}
         logger.debug("Command index has been invalidated.")
 
+    def rebuild(self):
+        """Invalidate command index and rebuild it if command index is enabled."""
+        self.invalidate()
+
+        if not self.cli_ctx:
+            return
+
+        if not self.cli_ctx.config.getboolean('core', 'use_command_index', fallback=True):
+            return
+
+        invocation = getattr(self.cli_ctx, 'invocation', None)
+        commands_loader = getattr(invocation, 'commands_loader', None) if invocation else None
+        if not commands_loader:
+            logger.debug('Skipping command index rebuild because commands_loader is unavailable.')
+            return
+
+        try:
+            logger.debug('Rebuilding command index after extension operation.')
+            commands_loader.load_command_table(None)
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning('Failed to rebuild command index after extension operation: %s', ex)
+
 
 class ModExtensionSuppress:  # pylint: disable=too-few-public-methods
 

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

@@ -276,29 +276,6 @@ def _augment_telemetry_with_ext_info(extension_name, ext=None):
         pass
 
 
-def _rebuild_command_index(cli_ctx):
-    command_index = CommandIndex(cli_ctx)
-    command_index.invalidate()
-
-    if not cli_ctx:
-        return
-
-    if not cli_ctx.config.getboolean('core', 'use_command_index', fallback=True):
-        return
-
-    invocation = getattr(cli_ctx, 'invocation', None)
-    commands_loader = getattr(invocation, 'commands_loader', None) if invocation else None
-    if not commands_loader:
-        logger.debug('Skipping command index rebuild because commands_loader is unavailable.')
-        return
-
-    try:
-        logger.debug('Rebuilding command index after extension operation.')
-        commands_loader.load_command_table(None)
-    except Exception as ex:  # pylint: disable=broad-except
-        logger.warning('Failed to rebuild command index after extension operation: %s', ex)
-
-
 def check_version_compatibility(azext_metadata):
     is_compatible, cli_core_version, min_required, max_required, min_ext_required = ext_compat_with_cli(azext_metadata)
     # logger.debug("Extension compatibility result: is_compatible=%s cli_core_version=%s min_required=%s "
@@ -375,7 +352,7 @@ def add_extension(cmd=None, source=None, extension_name=None, index_url=None, ye
     except ExtensionNotInstalledException:
         pass
 
-    _rebuild_command_index(cmd_cli_ctx)
+    CommandIndex(cmd_cli_ctx).rebuild()
 
 
 def is_cloud_shell_system_extension(ext_path):
@@ -399,7 +376,7 @@ def remove_extension(extension_name, cli_ctx=None):
         # We call this just before we remove the extension so we can get the metadata before it is gone
         _augment_telemetry_with_ext_info(extension_name, ext)
         rmtree_with_retry(ext.path)
-        _rebuild_command_index(cli_ctx)
+        CommandIndex(cli_ctx).rebuild()
     except ExtensionNotInstalledException as e:
         raise CLIError(e)
 
@@ -459,7 +436,7 @@ def update_extension(cmd=None, extension_name=None, index_url=None, pip_extra_in
             logger.debug('Copying %s to %s', backup_dir, extension_path)
             shutil.copytree(backup_dir, extension_path)
             raise CLIError('Failed to update. Rolled {} back to {}.'.format(extension_name, cur_version))
-        _rebuild_command_index(cmd_cli_ctx)
+        CommandIndex(cmd_cli_ctx).rebuild()
     except ExtensionNotInstalledException as e:
         raise CLIError(e)
 

src/azure-cli-core/azure/cli/core/tests/test_extension.py

@@ -10,14 +10,14 @@
 
 from unittest import mock
 from azure.cli.core.mock import DummyCli
+from azure.cli.core import CommandIndex
 
 from azure.cli.core.extension import (get_extensions, build_extension_path, extension_exists,
                                       get_extension, get_extension_names, get_extension_modname, ext_compat_with_cli,
                                       ExtensionNotInstalledException, WheelExtension,
                                       EXTENSIONS_MOD_PREFIX, EXT_METADATA_MINCLICOREVERSION, EXT_METADATA_MAXCLICOREVERSION)
 
 from azure.cli.core.extension.operations import list_available_extensions, add_extension, show_extension, remove_extension, list_extensions, update_extension
-from azure.cli.core.extension import operations as extension_operations
 
 # The test extension name
 EXT_NAME = 'myfirstcliextension'
@@ -108,35 +108,32 @@ def test_rebuild_command_index_when_enabled(self):
         cli_ctx.config.getboolean.return_value = True
         cli_ctx.invocation.commands_loader = mock.MagicMock()
 
-        with mock.patch('azure.cli.core.extension.operations.CommandIndex') as command_index:
-            extension_operations._rebuild_command_index(cli_ctx)  # pylint: disable=protected-access
+        with mock.patch.object(CommandIndex, 'invalidate') as invalidate:
+            CommandIndex(cli_ctx).rebuild()
 
-        command_index.assert_called_once_with(cli_ctx)
-        command_index.return_value.invalidate.assert_called_once()
+        invalidate.assert_called_once()
         cli_ctx.invocation.commands_loader.load_command_table.assert_called_once_with(None)
 
     def test_rebuild_command_index_skips_when_disabled(self):
         cli_ctx = mock.MagicMock()
         cli_ctx.config.getboolean.return_value = False
         cli_ctx.invocation.commands_loader = mock.MagicMock()
 
-        with mock.patch('azure.cli.core.extension.operations.CommandIndex') as command_index:
-            extension_operations._rebuild_command_index(cli_ctx)  # pylint: disable=protected-access
+        with mock.patch.object(CommandIndex, 'invalidate') as invalidate:
+            CommandIndex(cli_ctx).rebuild()
 
-        command_index.assert_called_once_with(cli_ctx)
-        command_index.return_value.invalidate.assert_called_once()
+        invalidate.assert_called_once()
         cli_ctx.invocation.commands_loader.load_command_table.assert_not_called()
 
     def test_rebuild_command_index_skips_without_loader(self):
         cli_ctx = mock.MagicMock()
         cli_ctx.config.getboolean.return_value = True
         cli_ctx.invocation = None
 
-        with mock.patch('azure.cli.core.extension.operations.CommandIndex') as command_index:
-            extension_operations._rebuild_command_index(cli_ctx)  # pylint: disable=protected-access
+        with mock.patch.object(CommandIndex, 'invalidate') as invalidate:
+            CommandIndex(cli_ctx).rebuild()
 
-        command_index.assert_called_once_with(cli_ctx)
-        command_index.return_value.invalidate.assert_called_once()
+        invalidate.assert_called_once()
 
     def test_no_extensions_in_dir(self):
         """ Directory exists but there are no extensions """

src/azure-cli/azure/cli/command_modules/util/_help.py

@@ -57,6 +57,11 @@
 short-summary: Upgrade Azure CLI and extensions
 """
 
+helps['rebuild-index'] = """
+type: command
+short-summary: Rebuild the Azure CLI command index cache.
+"""
+
 helps['demo'] = """
 type: group
 short-summary: Demos for designing, developing and demonstrating Azure CLI.

src/azure-cli/azure/cli/command_modules/util/commands.py

@@ -15,6 +15,9 @@ def load_command_table(self, _):
     with self.command_group('') as g:
         g.custom_command('upgrade', 'upgrade_version', is_preview=True)
 
+    with self.command_group('') as g:
+        g.custom_command('rebuild-index', 'rebuild_index')
+
     with self.command_group('demo', deprecate_info=g.deprecate(hide=True)) as g:
         g.custom_command('style', 'demo_style')
         g.custom_command('byo-access-token', 'byo_access_token')

src/azure-cli/azure/cli/command_modules/util/custom.py

@@ -39,6 +39,12 @@ def show_version(cmd):  # pylint: disable=unused-argument
     return versions
 
 
+def rebuild_index(cmd):
+    from azure.cli.core import CommandIndex
+    CommandIndex(cmd.cli_ctx).rebuild()
+    return {'status': 'rebuilt'}
+
+
 def upgrade_version(cmd, update_all=None, yes=None, allow_preview=None):  # pylint: disable=too-many-locals, too-many-statements, too-many-branches, no-member, unused-argument
     import platform
     import sys