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 bae41609c..8d043bdfb 100644 --- a/ansible_mitogen/plugins/connection/mitogen_kubectl.py +++ b/ansible_mitogen/plugins/connection/mitogen_kubectl.py @@ -46,12 +46,6 @@ import ansible_mitogen.loaders -_get_result = ansible_mitogen.loaders.connection_loader__get( - 'kubectl', - class_only=True, -) - - class Connection(ansible_mitogen.connection.Connection): transport = 'kubectl' @@ -61,17 +55,22 @@ class Connection(ansible_mitogen.connection.Connection): ) def __init__(self, *args, **kwargs): - if not _get_result: + result = ansible_mitogen.loaders.connection_loader__get_with_context( + 'kubectl', class_only=True, + ) + klass = result.object + if not klass: raise ansible.errors.AnsibleConnectionFailure(self.not_supported_msg) + + # FIXME Better name + # TODO Generalize? + self._klass = klass super(Connection, self).__init__(*args, **kwargs) def get_extra_args(self): - try: - # Ansible < 2.10, _get_result is the connection class - connection_options = _get_result.connection_options - except AttributeError: - # Ansible >= 2.10, _get_result is a get_with_context_result - connection_options = _get_result.object.connection_options + # TODO Can this overridden method be replaced by similar to used in mitogen_ssh? + # TODO Can self._klass.connection_options be a basis for something better than used in mitogen_ssh? + connection_options = self._klass.connection_options parameters = [] for key in connection_options: task_var_name = 'ansible_%s' % key diff --git a/ansible_mitogen/plugins/connection/mitogen_ssh.py b/ansible_mitogen/plugins/connection/mitogen_ssh.py index f6a27a6e7..da7e99023 100644 --- a/ansible_mitogen/plugins/connection/mitogen_ssh.py +++ b/ansible_mitogen/plugins/connection/mitogen_ssh.py @@ -33,6 +33,7 @@ import sys from ansible.plugins.connection.ssh import ( + Connection as _ansible_ssh_Connection, DOCUMENTATION as _ansible_ssh_DOCUMENTATION, ) @@ -55,18 +56,13 @@ del base_dir import ansible_mitogen.connection -import ansible_mitogen.loaders class Connection(ansible_mitogen.connection.Connection): transport = 'ssh' - vanilla_class = ansible_mitogen.loaders.connection_loader__get( - 'ssh', - class_only=True, - ) @staticmethod def _create_control_path(*args, **kwargs): """Forward _create_control_path() to the implementation in ssh.py.""" # https://github.com/dw/mitogen/issues/342 - return Connection.vanilla_class._create_control_path(*args, **kwargs) + return _ansible_ssh_Connection._create_control_path(*args, **kwargs) diff --git a/ansible_mitogen/strategy.py b/ansible_mitogen/strategy.py index 440e58112..5fdab355a 100644 --- a/ansible_mitogen/strategy.py +++ b/ansible_mitogen/strategy.py @@ -76,7 +76,7 @@ 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, class_only=False, collection_list=None, *args, **kwargs): """ While the mitogen strategy is active, trap action_loader.get() calls, augmenting any fetched class with ActionModuleMixin, which replaces various @@ -86,18 +86,23 @@ def wrap_action_loader__get(name, *args, **kwargs): This is used instead of static subclassing as it generalizes to third party action plugins outside the Ansible tree. """ - get_kwargs = {'class_only': True} if name in ('fetch',): name = 'mitogen_' + name - get_kwargs['collection_list'] = kwargs.pop('collection_list', None) - klass = ansible_mitogen.loaders.action_loader__get(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) + result = ansible_mitogen.loaders.action_loader__get_with_context( + name, class_only=True, collection_list=collection_list, + ) + + if not result.object: + return result + + adorned_bases = (ansible_mitogen.mixins.ActionModuleMixin, result.object) + adorned_class = type(str(name), adorned_bases, {}) + + if class_only: + return result._replace(object=adorned_class) + + return result._replace(object=adorned_class(*args, **kwargs)) REDIRECTED_CONNECTION_PLUGINS = ( @@ -115,7 +120,7 @@ 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. @@ -123,7 +128,9 @@ def wrap_connection_loader__get(name, *args, **kwargs): 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 +180,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,12 +191,8 @@ def _remove_wrappers(self): """ Uninstall the PluginLoader monkey patches. """ - ansible_mitogen.loaders.action_loader.get = ( - ansible_mitogen.loaders.action_loader__get - ) - ansible_mitogen.loaders.connection_loader.get_with_context = ( - ansible_mitogen.loaders.connection_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_with_context ansible.executor.process.worker.WorkerProcess.run = worker__run def install(self):