diff --git a/ansible_mitogen/loaders.py b/ansible_mitogen/loaders.py index 632b11b14..123dd4ac6 100644 --- a/ansible_mitogen/loaders.py +++ b/ansible_mitogen/loaders.py @@ -99,5 +99,5 @@ def assert_supported_release(): from ansible.plugins.loader import strategy_loader # These are original, unwrapped implementations -action_loader__get = action_loader.get -connection_loader__get = connection_loader.get_with_context +action_loader__get_with_context = action_loader.get_with_context +connection_loader__get_with_context = connection_loader.get_with_context diff --git a/ansible_mitogen/plugins/connection/mitogen_kubectl.py b/ansible_mitogen/plugins/connection/mitogen_kubectl.py index bd9626aff..33ae49e6a 100644 --- a/ansible_mitogen/plugins/connection/mitogen_kubectl.py +++ b/ansible_mitogen/plugins/connection/mitogen_kubectl.py @@ -48,7 +48,7 @@ class Connection(ansible_mitogen.connection.Connection): transport = 'kubectl' - (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get( + (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context( 'kubectl', class_only=True, ) diff --git a/ansible_mitogen/plugins/connection/mitogen_ssh.py b/ansible_mitogen/plugins/connection/mitogen_ssh.py index 74dd616c9..b953edba3 100644 --- a/ansible_mitogen/plugins/connection/mitogen_ssh.py +++ b/ansible_mitogen/plugins/connection/mitogen_ssh.py @@ -60,7 +60,7 @@ class Connection(ansible_mitogen.connection.Connection): transport = 'ssh' - (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get( + (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context( 'ssh', class_only=True, ) diff --git a/ansible_mitogen/strategy.py b/ansible_mitogen/strategy.py index 440e58112..8851a673f 100644 --- a/ansible_mitogen/strategy.py +++ b/ansible_mitogen/strategy.py @@ -48,6 +48,8 @@ import ansible.template import ansible.utils.sentinel +from ansible.plugins.loader import get_with_context_result + def _patch_awx_callback(): """ @@ -76,12 +78,12 @@ def patch_add_local(self, **kwargs): _patch_awx_callback() -def wrap_action_loader__get(name, *args, **kwargs): +def wrap_action_loader__get_with_context(name, *args, **kwargs): """ - While the mitogen strategy is active, trap action_loader.get() calls, - augmenting any fetched class with ActionModuleMixin, which replaces various - helper methods inherited from ActionBase with implementations that avoid - the use of shell fragments wherever possible. + While the mitogen strategy is active, trap action_loader.get_with_context() + calls, augmenting any fetched class with ActionModuleMixin, which replaces + various helper methods inherited from ActionBase with implementations that + avoid the use of shell fragments wherever possible. This is used instead of static subclassing as it generalizes to third party action plugins outside the Ansible tree. @@ -91,13 +93,19 @@ def wrap_action_loader__get(name, *args, **kwargs): name = 'mitogen_' + name get_kwargs['collection_list'] = kwargs.pop('collection_list', None) - klass = ansible_mitogen.loaders.action_loader__get(name, **get_kwargs) + (klass, context) = ansible_mitogen.loaders.action_loader__get_with_context( + name, + **get_kwargs + ) + if klass: bases = (ansible_mitogen.mixins.ActionModuleMixin, klass) adorned_klass = type(str(name), bases, {}) if kwargs.get('class_only'): - return adorned_klass - return adorned_klass(*args, **kwargs) + return get_with_context_result(adorned_klass, context) + return get_with_context_result(adorned_klass(*args, **kwargs), context) + + return get_with_context_result(None, context) REDIRECTED_CONNECTION_PLUGINS = ( @@ -115,15 +123,16 @@ def wrap_action_loader__get(name, *args, **kwargs): ) -def wrap_connection_loader__get(name, *args, **kwargs): +def wrap_connection_loader__get_with_context(name, *args, **kwargs): """ - While a Mitogen strategy is active, rewrite connection_loader.get() calls - for some transports into requests for a compatible Mitogen transport. + While a Mitogen strategy is active, rewrite + connection_loader.get_with_context() calls for some transports into + requests for a compatible Mitogen transport. """ if name in REDIRECTED_CONNECTION_PLUGINS: name = 'mitogen_' + name - return ansible_mitogen.loaders.connection_loader__get(name, *args, **kwargs) + return ansible_mitogen.loaders.connection_loader__get_with_context(name, *args, **kwargs) def wrap_worker__run(self): @@ -173,8 +182,8 @@ def _install_wrappers(self): Install our PluginLoader monkey patches and update global variables with references to the real functions. """ - ansible_mitogen.loaders.action_loader.get = wrap_action_loader__get - ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get + ansible_mitogen.loaders.action_loader.get_with_context = wrap_action_loader__get_with_context + ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get_with_context global worker__run worker__run = ansible.executor.process.worker.WorkerProcess.run @@ -184,11 +193,11 @@ def _remove_wrappers(self): """ Uninstall the PluginLoader monkey patches. """ - ansible_mitogen.loaders.action_loader.get = ( - ansible_mitogen.loaders.action_loader__get + ansible_mitogen.loaders.action_loader.get_with_context = ( + ansible_mitogen.loaders.action_loader__get_with_context ) ansible_mitogen.loaders.connection_loader.get_with_context = ( - ansible_mitogen.loaders.connection_loader__get + ansible_mitogen.loaders.connection_loader__get_with_context ) ansible.executor.process.worker.WorkerProcess.run = worker__run