diff --git a/awscli/lazy_emitter.py b/awscli/lazy_emitter.py new file mode 100644 index 000000000000..2043902c77cd --- /dev/null +++ b/awscli/lazy_emitter.py @@ -0,0 +1,157 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Lazy-initializing event emitter for the AWS CLI plugin system. + +LazyInitEmitter wraps a HierarchicalEmitter and uses the plugin registry +to initialize plugins on demand. Before emitting event X, it finds all +initializer entries whose event patterns match X (using the same +prefix/wildcard semantics as HierarchicalEmitter), calls each initializer +at most once, then delegates to the underlying emitter for normal dispatch. +""" + +import copy +import importlib +import logging + +from botocore.hooks import HierarchicalEmitter, PrefixTrie + +from awscli.handlers_registry import ( + PLUGIN_REGISTRY, + CommandTableOp, +) +from awscli.lazy import LazyCommand + +logger = logging.getLogger(__name__) + + +class LazyInitEmitter(HierarchicalEmitter): + """HierarchicalEmitter that lazily initializes plugins from a registry. + + The registry maps event patterns to lists of (module, fn_name) tuples. + Before emitting any event, this emitter checks whether there are + uninitialised plugins whose event patterns match the event being emitted. + If so, it imports and calls them, then proceeds with normal dispatch. + """ + + def __init__(self, plugin_registry=None, main_command_table_ops=None): + super().__init__() + self._init_trie = PrefixTrie() + # set of (module, fn_name) + self._initialized = set() + # number of entries not yet initialized + self._pending_count = 0 + # event_name -> list of entries from init trie + self._init_cache: dict[str, list] = {} + self._main_ops = main_command_table_ops + self._main_ops_applied = False + if plugin_registry: + self.load_registry(plugin_registry) + + def load_registry(self, registry): + unique = set() + for event_pattern, entries in registry.items(): + for entry in entries: + self._init_trie.append_item(event_pattern, entry) + if entry not in unique: + unique.add(entry) + self._pending_count += 1 + self._init_cache = {} + + @property + def initialized_count(self): + return len(self._initialized) + + def _apply_main_command_table_ops(self, kwargs): + """Apply pre-computed renames and LazyCommand additions. + + This replaces the normal lazy-init path for + building-command-table.main entries, avoiding the import of + heavy plugin modules until the command is actually invoked. + """ + command_table = kwargs.get('command_table') + session = kwargs.get('session') + if command_table is None or session is None: + return + + for op in self._main_ops: + if op[0] == CommandTableOp.RENAME: + _, old_name, new_name = op + if old_name in command_table: + current = command_table[old_name] + command_table[new_name] = current + current.name = new_name + del command_table[old_name] + elif op[0] == CommandTableOp.ADD: + _, cmd_name, cmd_module, cmd_class = op + command_table[cmd_name] = LazyCommand( + cmd_name, + session, + cmd_module, + cmd_class, + ) + else: + raise RuntimeError(f'Unknown command table ops entry: {op}') + + # Mark all building-command-table.main entries as initialized so + # _ensure_initialized never imports them. The ops list fully + # accounts for all plugins registered against this event. + for entry in PLUGIN_REGISTRY.get('building-command-table.main', []): + entry = tuple(entry) + if entry not in self._initialized: + self._initialized.add(entry) + self._pending_count -= 1 + + def _ensure_initialized(self, event_name): + """Initialize any plugins whose event patterns match event_name.""" + if self._pending_count == 0: + return + candidates = self._init_cache.get(event_name) + if candidates is None: + candidates = self._init_trie.prefix_search(event_name) + self._init_cache[event_name] = candidates + for entry in candidates: + if entry not in self._initialized: + self._initialized.add(entry) + self._pending_count -= 1 + module_path, fn_name = entry + logger.debug( + 'Lazy-initializing plugin %s.%s for event %s', + module_path, + fn_name, + event_name, + ) + mod = importlib.import_module(module_path) + fn = getattr(mod, fn_name) + fn(self) + + def _emit(self, event_name, kwargs, stop_on_response=False): + if ( + self._main_ops + and not self._main_ops_applied + and event_name == 'building-command-table.main' + ): + self._main_ops_applied = True + self._apply_main_command_table_ops(kwargs) + self._ensure_initialized(event_name) + return super()._emit(event_name, kwargs, stop_on_response) + + def __copy__(self): + new_instance = self.__class__() + new_state = self.__dict__.copy() + new_state['_handlers'] = copy.copy(self._handlers) + new_state['_unique_id_handlers'] = copy.copy(self._unique_id_handlers) + new_state['_init_trie'] = copy.copy(self._init_trie) + new_state['_initialized'] = copy.copy(self._initialized) + new_state['_main_ops_applied'] = self._main_ops_applied + new_instance.__dict__ = new_state + return new_instance diff --git a/tests/unit/test_lazy_emitter.py b/tests/unit/test_lazy_emitter.py new file mode 100644 index 000000000000..281b6d38a6d7 --- /dev/null +++ b/tests/unit/test_lazy_emitter.py @@ -0,0 +1,236 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +from unittest.mock import MagicMock, patch + +import pytest + +from awscli.handlers_registry import CommandTableOp +from awscli.lazy import LazyCommand +from awscli.lazy_emitter import LazyInitEmitter + + +@pytest.fixture +def mock_module(): + """Create a mock module with a callable init function.""" + module = MagicMock() + module.my_init = MagicMock() + return module + + +class TestLazyInitEmitterPrefixMatching: + def test_bare_prefix_entry_initialized_on_dotted_emit(self, mock_module): + registry = { + 'building-command-table': [ + ('test.module', 'my_init'), + ], + } + emitter = LazyInitEmitter() + emitter.load_registry(registry) + + with patch('importlib.import_module', return_value=mock_module): + emitter.emit('building-command-table.main', command_table={}) + + assert emitter.initialized_count == 1 + mock_module.my_init.assert_called_once() + + def test_exact_match_entry_initialized(self, mock_module): + registry = { + 'building-command-table.main': [ + ('test.module', 'my_init'), + ], + } + emitter = LazyInitEmitter() + emitter.load_registry(registry) + + with patch('importlib.import_module', return_value=mock_module): + emitter.emit('building-command-table.main', command_table={}) + + assert emitter.initialized_count == 1 + mock_module.my_init.assert_called_once() + + def test_unrelated_entry_not_initialized(self, mock_module): + registry = { + 'building-command-table.ecs': [ + ('test.module', 'my_init'), + ], + } + emitter = LazyInitEmitter() + emitter.load_registry(registry) + + with patch('importlib.import_module', return_value=mock_module): + emitter.emit('building-command-table.main', command_table={}) + + assert emitter.initialized_count == 0 + mock_module.my_init.assert_not_called() + + def test_multiple_prefix_levels_all_initialized(self, mock_module): + mock_module.other_init = MagicMock() + registry = { + 'building-command-table': [ + ('test.module', 'my_init'), + ], + 'building-command-table.main': [ + ('test.module', 'other_init'), + ], + } + emitter = LazyInitEmitter() + emitter.load_registry(registry) + + with patch('importlib.import_module', return_value=mock_module): + emitter.emit('building-command-table.main', command_table={}) + + assert emitter.initialized_count == 2 + mock_module.my_init.assert_called_once() + mock_module.other_init.assert_called_once() + + def test_entry_initialized_only_once(self, mock_module): + registry = { + 'building-command-table': [ + ('test.module', 'my_init'), + ], + } + emitter = LazyInitEmitter() + emitter.load_registry(registry) + + with patch('importlib.import_module', return_value=mock_module): + emitter.emit('building-command-table.main', command_table={}) + emitter.emit('building-command-table.ecs', command_table={}) + + assert emitter.initialized_count == 1 + mock_module.my_init.assert_called_once() + + +class TestMainCommandTableOps: + def test_covered_plugins_not_imported(self, mock_module): + registry = { + 'building-command-table.main': [ + ('heavy.module', 'register_add_cmd'), + ], + } + main_ops = [ + (CommandTableOp.ADD, 'my-cmd', 'heavy.module', 'MyCommand'), + ] + emitter = LazyInitEmitter(main_command_table_ops=main_ops) + emitter.load_registry(registry) + + command_table = {'existing': MagicMock(name='existing')} + session = MagicMock() + + with ( + patch('awscli.lazy_emitter.PLUGIN_REGISTRY', registry), + patch('importlib.import_module', return_value=mock_module) as imp, + ): + emitter.emit( + 'building-command-table.main', + command_table=command_table, + session=session, + ) + + imp.assert_not_called() + assert emitter.initialized_count == 1 + + def test_rename_op_applied_to_command_table(self): + registry = { + 'building-command-table.main': [ + ('rename.module', 'register_rename'), + ], + } + main_ops = [ + (CommandTableOp.RENAME, 'old-name', 'new-name'), + ] + emitter = LazyInitEmitter(main_command_table_ops=main_ops) + emitter.load_registry(registry) + + cmd = MagicMock() + cmd.name = 'old-name' + command_table = {'old-name': cmd} + session = MagicMock() + + with patch('awscli.lazy_emitter.PLUGIN_REGISTRY', registry): + emitter.emit( + 'building-command-table.main', + command_table=command_table, + session=session, + ) + + assert 'old-name' not in command_table + assert 'new-name' in command_table + assert command_table['new-name'] is cmd + assert cmd.name == 'new-name' + + def test_add_op_creates_lazy_command(self): + registry = { + 'building-command-table.main': [ + ('add.module', 'register_add'), + ], + } + main_ops = [ + (CommandTableOp.ADD, 'my-cmd', 'add.module.impl', 'MyCommand'), + ] + emitter = LazyInitEmitter(main_command_table_ops=main_ops) + emitter.load_registry(registry) + + command_table = {} + session = MagicMock() + + with patch('awscli.lazy_emitter.PLUGIN_REGISTRY', registry): + emitter.emit( + 'building-command-table.main', + command_table=command_table, + session=session, + ) + + assert 'my-cmd' in command_table + assert isinstance(command_table['my-cmd'], LazyCommand) + assert command_table['my-cmd'].name == 'my-cmd' + + def test_main_ops_skips_covered_but_initializes_bare_prefix( + self, mock_module + ): + registry = { + 'building-command-table': [ + ('global.module', 'register_global'), + ], + 'building-command-table.main': [ + ('heavy.module', 'register_add_cmd'), + ], + } + main_ops = [ + (CommandTableOp.ADD, 'my-cmd', 'heavy.module', 'MyCommand'), + ] + emitter = LazyInitEmitter(main_command_table_ops=main_ops) + emitter.load_registry(registry) + + command_table = {} + session = MagicMock() + + mock_global = MagicMock() + mock_global.register_global = MagicMock() + + def import_side_effect(mod_path): + if mod_path == 'global.module': + return mock_global + raise AssertionError(f'Unexpected import of {mod_path!r}') + + with ( + patch('awscli.lazy_emitter.PLUGIN_REGISTRY', registry), + patch('importlib.import_module', side_effect=import_side_effect), + ): + emitter.emit( + 'building-command-table.main', + command_table=command_table, + session=session, + ) + + mock_global.register_global.assert_called_once() + assert emitter.initialized_count == 2