{Core} Build and use command index

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

@@ -105,6 +105,9 @@ def __len__(self):
 # SESSION provides read-write session variables
 SESSION = Session()
 
+# INDEX contains {top-level command: [command_modules and extensions]} mapping index
+INDEX = Session()
+
 # VERSIONS provides local versions and pypi versions.
 # DO NOT USE it to get the current version of azure-cli,
 # it could be lagged behind and can be used to check whether

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

@@ -537,6 +537,8 @@ def update_cloud(cli_ctx, cloud):
     if not _get_cloud(cli_ctx, cloud.name):
         raise CloudNotRegisteredException(cloud.name)
     _save_cloud(cloud, overwrite=True)
+    from azure.cli.core import invalidate_command_index
+    invalidate_command_index()
 
 
 def remove_cloud(cli_ctx, cloud_name):

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

@@ -21,7 +21,7 @@
 
 # pylint: disable=unused-import
 from azure.cli.core.commands.constants import (
-    BLACKLISTED_MODS, DEFAULT_QUERY_TIME_RANGE, CLI_COMMON_KWARGS, CLI_COMMAND_KWARGS, CLI_PARAM_KWARGS,
+    BLOCKED_MODS, DEFAULT_QUERY_TIME_RANGE, CLI_COMMON_KWARGS, CLI_COMMAND_KWARGS, CLI_PARAM_KWARGS,
     CLI_POSITIONAL_PARAM_KWARGS, CONFIRM_PARAM_NAME)
 from azure.cli.core.commands.parameters import (
     AzArgumentContext, patch_arg_make_required, patch_arg_make_optional)

src/azure-cli-core/azure/cli/core/commands/constants.py

@@ -30,7 +30,7 @@
 # 1 hour in milliseconds
 DEFAULT_QUERY_TIME_RANGE = 3600000
 
-BLACKLISTED_MODS = ['context', 'shell', 'documentdb', 'component']
+BLOCKED_MODS = ['context', 'shell', 'documentdb', 'component']
 
 SURVEY_PROMPT = 'Please let us know how we are doing: https://aka.ms/clihats'
 SURVEY_PROMPT_COLOR = Fore.YELLOW + Style.BRIGHT + 'Please let us know how we are doing: ' + Fore.BLUE + \

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

@@ -18,6 +18,7 @@
 import requests
 from pkg_resources import parse_version
 
+from azure.cli.core import invalidate_command_index
 from azure.cli.core.util import CLIError, reload_module
 from azure.cli.core.extension import (extension_exists, build_extension_path, get_extensions, get_extension_modname,
                                       get_extension, ext_compat_with_cli,
@@ -188,8 +189,8 @@ def _augment_telemetry_with_ext_info(extension_name, ext=None):
 
 def check_version_compatibility(azext_metadata):
     is_compatible, cli_core_version, min_required, max_required = ext_compat_with_cli(azext_metadata)
-    logger.debug("Extension compatibility result: is_compatible=%s cli_core_version=%s min_required=%s "
-                 "max_required=%s", is_compatible, cli_core_version, min_required, max_required)
+    # logger.debug("Extension compatibility result: is_compatible=%s cli_core_version=%s min_required=%s "
+    #              "max_required=%s", is_compatible, cli_core_version, min_required, max_required)
     if not is_compatible:
         min_max_msg_fmt = "The '{}' extension is not compatible with this version of the CLI.\n" \
                           "You have CLI core version {} and this extension " \
@@ -234,6 +235,7 @@ def add_extension(cmd, source=None, extension_name=None, index_url=None, yes=Non
                            "Please use with discretion.", extension_name)
         elif extension_name and ext.preview:
             logger.warning("The installed extension '%s' is in preview.", extension_name)
+        invalidate_command_index()
     except ExtensionNotInstalledException:
         pass
 
@@ -253,6 +255,7 @@ def log_err(func, path, exc_info):
         # 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)
         shutil.rmtree(ext.path, onerror=log_err)
+        invalidate_command_index()
     except ExtensionNotInstalledException as e:
         raise CLIError(e)
 
@@ -305,6 +308,7 @@ def update_extension(cmd, extension_name, index_url=None, pip_extra_index_urls=N
             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))
+        invalidate_command_index()
     except ExtensionNotInstalledException as e:
         raise CLIError(e)
 

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

@@ -146,10 +146,14 @@ def _mock_import_lib(_):
         return mock_obj
 
     def _mock_iter_modules(_):
-        return [(None, __name__, None)]
+        return [(None, "hello", None),
+                (None, "extra", None)]
 
     def _mock_extension_modname(ext_name, ext_dir):
-        return ext_name
+        if ext_name.endswith('.ExtCommandsLoader'):
+            return "azext_hello1"
+        if ext_name.endswith('.Ext2CommandsLoader'):
+            return "azext_hello2"
 
     def _mock_get_extensions():
         MockExtension = namedtuple('Extension', ['name', 'preview', 'experimental', 'path', 'get_metadata'])
@@ -163,7 +167,18 @@ class TestCommandsLoader(AzCommandsLoader):
             def load_command_table(self, args):
                 super(TestCommandsLoader, self).load_command_table(args)
                 with self.command_group('hello', operations_tmpl='{}#TestCommandRegistration.{{}}'.format(__name__)) as g:
-                    g.command('world', 'sample_vm_get')
+                    g.command('mod-only', 'sample_vm_get')
+                    g.command('overridden', 'sample_vm_get')
+                self.__module__ = "azure.cli.command_modules.hello"
+                return self.command_table
+
+        class Test2CommandsLoader(AzCommandsLoader):
+            # An extra group that is not loaded if a module is found from the index
+            def load_command_table(self, args):
+                super(Test2CommandsLoader, self).load_command_table(args)
+                with self.command_group('extra', operations_tmpl='{}#TestCommandRegistration.{{}}'.format(__name__)) as g:
+                    g.command('unused', 'sample_vm_get')
+                self.__module__ = "azure.cli.command_modules.extra"
                 return self.command_table
 
         # A command from an extension
@@ -172,7 +187,8 @@ class ExtCommandsLoader(AzCommandsLoader):
             def load_command_table(self, args):
                 super(ExtCommandsLoader, self).load_command_table(args)
                 with self.command_group('hello', operations_tmpl='{}#TestCommandRegistration.{{}}'.format(__name__)) as g:
-                    g.command('noodle', 'sample_vm_get')
+                    g.command('ext-only', 'sample_vm_get')
+                self.__module__ = "azext_hello1"
                 return self.command_table
 
         # A command from an extension that overrides the original command
@@ -181,22 +197,26 @@ class Ext2CommandsLoader(AzCommandsLoader):
             def load_command_table(self, args):
                 super(Ext2CommandsLoader, self).load_command_table(args)
                 with self.command_group('hello', operations_tmpl='{}#TestCommandRegistration.{{}}'.format(__name__)) as g:
-                    g.command('world', 'sample_vm_get')
+                    g.command('overridden', 'sample_vm_get')
+                self.__module__ = "azext_hello2"
                 return self.command_table
 
         if prefix == 'azure.cli.command_modules.':
-            command_loaders = {'TestCommandsLoader': TestCommandsLoader}
+            command_loaders = {'hello': TestCommandsLoader, 'extra': Test2CommandsLoader}
         else:
-            command_loaders = {'ExtCommandsLoader': ExtCommandsLoader, 'Ext2CommandsLoader': Ext2CommandsLoader}
+            command_loaders = {'azext_hello1': ExtCommandsLoader, 'azext_hello2': Ext2CommandsLoader}
 
         module_command_table = {}
-        for _, loader_cls in command_loaders.items():
+        for mod_name, loader_cls in command_loaders.items():
+            # If name is provided, only load the named module
+            if name and name != mod_name:
+                continue
             command_loader = loader_cls(cli_ctx=loader.cli_ctx)
             command_table = command_loader.load_command_table(args)
             if command_table:
                 module_command_table.update(command_table)
                 loader.loaders.append(command_loader)  # this will be used later by the load_arguments method
-        return module_command_table, {}
+        return module_command_table, command_loader.command_group_table
 
     @mock.patch('importlib.import_module', _mock_import_lib)
     @mock.patch('pkgutil.iter_modules', _mock_iter_modules)
@@ -206,19 +226,131 @@ def load_command_table(self, args):
     def test_register_command_from_extension(self):
 
         from azure.cli.core.commands import _load_command_loader
+        from azure.cli.core._session import INDEX
+        from azure.cli.core import _COMMAND_INDEX, _COMMAND_INDEX_VERSION, invalidate_command_index
+
         cli = DummyCli()
+        # Clear the command index
+        INDEX[_COMMAND_INDEX] = {}
+        self.assertFalse(INDEX[_COMMAND_INDEX])
         main_loader = MainCommandsLoader(cli)
         cli.loader = main_loader
 
         cmd_tbl = cli.loader.load_command_table(None)
-        ext1 = cmd_tbl['hello noodle']
-        ext2 = cmd_tbl['hello world']
-
-        self.assertTrue(isinstance(ext1.command_source, ExtensionCommandSource))
-        self.assertFalse(ext1.command_source.overrides_command)
-
-        self.assertTrue(isinstance(ext2.command_source, ExtensionCommandSource))
-        self.assertTrue(ext2.command_source.overrides_command)
+        hello_mod_only_cmd = cmd_tbl['hello mod-only']
+        hello_ext_only_cmd = cmd_tbl['hello ext-only']
+        hello_overridden_cmd = cmd_tbl['hello overridden']
+
+        self.assertEqual(hello_mod_only_cmd.command_source, 'hello')
+        self.assertEqual(hello_mod_only_cmd.loader.__module__, 'azure.cli.command_modules.hello')
+
+        self.assertTrue(isinstance(hello_ext_only_cmd.command_source, ExtensionCommandSource))
+        self.assertFalse(hello_ext_only_cmd.command_source.overrides_command)
+
+        self.assertTrue(isinstance(hello_overridden_cmd.command_source, ExtensionCommandSource))
+        self.assertTrue(hello_overridden_cmd.command_source.overrides_command)
+
+        expected_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello2', 'azext_hello1'],
+                                  'extra': ['azure.cli.command_modules.extra']}
+        expected_command_table = ['hello mod-only', 'hello overridden', 'extra unused', 'hello ext-only']
+
+        # Test command index is built for None args
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+        # Test command index is built when `args` is provided
+        INDEX[_COMMAND_INDEX] = {}
+        cmd_tbl = cli.loader.load_command_table(["hello", "mod-only"])
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+        # Test rebuild command index if no module found
+        INDEX[_COMMAND_INDEX] = {"network": ["azure.cli.command_modules.network"]}
+        cli.loader.load_command_table(["hello", "mod-only"])
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+        with mock.patch("azure.cli.core.__version__", "2.7.0"):
+            # Test rebuild command index if version is not present
+            del INDEX[_COMMAND_INDEX_VERSION]
+            del INDEX[_COMMAND_INDEX]
+            cli.loader.load_command_table(["hello", "mod-only"])
+            self.assertEqual(INDEX[_COMMAND_INDEX_VERSION], "2.7.0")
+            self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+            # Test rebuild command index if version is not valid
+            INDEX[_COMMAND_INDEX_VERSION] = ""
+            INDEX[_COMMAND_INDEX] = {}
+            cli.loader.load_command_table(["hello", "mod-only"])
+            self.assertEqual(INDEX[_COMMAND_INDEX_VERSION], "2.7.0")
+            self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+            # Test rebuild command index if version is outdated
+            INDEX[_COMMAND_INDEX_VERSION] = "2.6.0"
+            INDEX[_COMMAND_INDEX] = {}
+            cli.loader.load_command_table(["hello", "mod-only"])
+            self.assertEqual(INDEX[_COMMAND_INDEX_VERSION], "2.7.0")
+            self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+        # Test rebuild command index if modules are found but outdated
+        # This only happens in dev environment. For users, the version check logic prevents it
+        INDEX[_COMMAND_INDEX] = {"hello": ["azure.cli.command_modules.extra"]}
+        cli.loader.load_command_table(["hello", "mod-only"])
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+
+        # Test irrelevant commands are not loaded
+        INDEX[_COMMAND_INDEX] = expected_command_index
+        cmd_tbl = cli.loader.load_command_table(["hello", "mod-only"])
+        self.assertEqual(['hello mod-only', 'hello overridden', 'hello ext-only'], list(cmd_tbl.keys()))
+
+        # Full scenario test 1: Installing an extension 'azext_hello1' that extends 'hello' group
+        outdated_command_index = {'hello': ['azure.cli.command_modules.hello'],
+                                  'extra': ['azure.cli.command_modules.extra']}
+        INDEX[_COMMAND_INDEX] = outdated_command_index
+
+        # Command for an outdated group
+        cmd_tbl = cli.loader.load_command_table(["hello", "-h"])
+        # Index is not updated, and only built-in commands are loaded
+        self.assertEqual(INDEX[_COMMAND_INDEX], outdated_command_index)
+        self.assertListEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden'])
+
+        # Command index is explicitly invalidated by azure.cli.core.extension.operations.add_extension
+        invalidate_command_index()
+
+        cmd_tbl = cli.loader.load_command_table(["hello", "-h"])
+        # Index is updated, and new commands are loaded
+        self.assertEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+        self.assertListEqual(list(cmd_tbl), expected_command_table)
+
+        # Full scenario test 2: Installing extension 'azext_hello2' that overrides existing command 'hello overridden'
+        outdated_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello1'],
+                                  'extra': ['azure.cli.command_modules.extra']}
+        INDEX[_COMMAND_INDEX] = outdated_command_index
+        # Command for an overridden command
+        cmd_tbl = cli.loader.load_command_table(["hello", "overridden"])
+        # Index is not updated
+        self.assertEqual(INDEX[_COMMAND_INDEX], outdated_command_index)
+        # With the old command index, 'hello overridden' is loaded from the build-in module
+        hello_overridden_cmd = cmd_tbl['hello overridden']
+        self.assertEqual(hello_overridden_cmd.command_source, 'hello')
+        self.assertListEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'hello ext-only'])
+
+        # Command index is explicitly invalidated by azure.cli.core.extension.operations.add_extension
+        invalidate_command_index()
+
+        # Command index is updated, and 'hello overridden' is loaded from the new extension
+        cmd_tbl = cli.loader.load_command_table(["hello", "overridden"])
+        hello_overridden_cmd = cmd_tbl['hello overridden']
+        self.assertTrue(isinstance(hello_overridden_cmd.command_source, ExtensionCommandSource))
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+        self.assertListEqual(list(cmd_tbl), expected_command_table)
+
+        # Call again with the new command index. Irrelevant commands are not loaded
+        cmd_tbl = cli.loader.load_command_table(["hello", "overridden"])
+        hello_overridden_cmd = cmd_tbl['hello overridden']
+        self.assertTrue(isinstance(hello_overridden_cmd.command_source, ExtensionCommandSource))
+        self.assertDictEqual(INDEX[_COMMAND_INDEX], expected_command_index)
+        self.assertListEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'hello ext-only'])
+
+        del INDEX[_COMMAND_INDEX_VERSION]
+        del INDEX[_COMMAND_INDEX]
 
     def test_argument_with_overrides(self):