Merge pull request #1215 from Nihlus/fix-connection-loader

ansible_mitogen: Fix usage of connection_loader__get.

ansible_mitogen/loaders.py

@@ -99,5 +99,5 @@ from ansible.plugins.loader import shell_loader
 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

ansible_mitogen/plugins/connection/mitogen_kubectl.py

@@ -46,14 +46,12 @@ import ansible_mitogen.connection
 import ansible_mitogen.loaders
 
 
-_get_result = ansible_mitogen.loaders.connection_loader__get(
-    'kubectl',
-    class_only=True,
-)
-
-
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'kubectl'
+    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context(
+        'kubectl',
+        class_only=True,
+    )
 
     not_supported_msg = (
         'The "mitogen_kubectl" plug-in requires a version of Ansible '
@@ -61,17 +59,12 @@ class Connection(ansible_mitogen.connection.Connection):
     )
 
     def __init__(self, *args, **kwargs):
-        if not _get_result:
+        if not Connection.vanilla_class:
             raise ansible.errors.AnsibleConnectionFailure(self.not_supported_msg)
         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
+        connection_options = Connection.vanilla_class.connection_options
         parameters = []
         for key in connection_options:
             task_var_name = 'ansible_%s' % key

ansible_mitogen/plugins/connection/mitogen_ssh.py

@@ -60,7 +60,7 @@ import ansible_mitogen.loaders
 
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'ssh'
-    vanilla_class = ansible_mitogen.loaders.connection_loader__get(
+    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context(
         'ssh',
         class_only=True,
     )

ansible_mitogen/strategy.py

@@ -47,6 +47,8 @@ import ansible_mitogen.process
 import ansible.executor.process.worker
 import ansible.template
 import ansible.utils.sentinel
+import ansible.playbook.play_context
+import ansible.plugins.loader
 
 
 def _patch_awx_callback():
@@ -76,12 +78,12 @@ def _patch_awx_callback():
 _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,26 @@ 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 ansible.plugins.loader.get_with_context_result(
+                adorned_klass,
+                context
+            )
+
+        return ansible.plugins.loader.get_with_context_result(
+            adorned_klass(*args, **kwargs),
+            context
+        )
+
+    return ansible.plugins.loader.get_with_context_result(None, context)
 
 
 REDIRECTED_CONNECTION_PLUGINS = (
@@ -115,15 +130,26 @@ REDIRECTED_CONNECTION_PLUGINS = (
 )
 
 
-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:
+    is_play_using_mitogen_connection = None
+    if len(args) > 0 and isinstance(args[0], ansible.playbook.play_context.PlayContext):
+        play_context = args[0]
+        is_play_using_mitogen_connection = play_context.connection in REDIRECTED_CONNECTION_PLUGINS
+
+    # assume true if we're not in a play context since we're using a Mitogen strategy
+    if is_play_using_mitogen_connection is None:
+        is_play_using_mitogen_connection = True
+
+    redirect_connection = name in REDIRECTED_CONNECTION_PLUGINS and is_play_using_mitogen_connection
+    if redirect_connection:
         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 +199,8 @@ class AnsibleWrappers(object):
         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 +210,11 @@ class AnsibleWrappers(object):
         """
         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
 

tests/ansible/regression/all.yml

@@ -13,6 +13,7 @@
 - import_playbook: issue_591__setuptools_cwd_crash.yml
 - import_playbook: issue_615__streaming_transfer.yml
 - import_playbook: issue_655__wait_for_connection_error.yml
+- import_playbook: issue_766__get_with_context.yml
 - import_playbook: issue_776__load_plugins_called_twice.yml
 - import_playbook: issue_952__ask_become_pass.yml
 - import_playbook: issue_1066__add_host__host_key_checking.yml

tests/ansible/regression/issue_655__wait_for_connection_error.yml

@@ -11,7 +11,10 @@
   tasks:
     - meta: end_play
       when:
-        # Podman versions available in Homebrew have dropped macOS 12 support.
+        # Podman versions available in Homebrew require macOS 13+ (Ventura).
+        # https://formulae.brew.sh/formula/podman
+        # See also
+        #   - issue_766__get_with_context.yml
         - ansible_facts.system == 'Darwin'
         - ansible_facts.distribution_version is version('13.0', '<', strict=True)
 

tests/ansible/regression/issue_766__get_with_context.yml

@@ -0,0 +1,64 @@
+# https://github.com/mitogen-hq/mitogen/issues/776
+---
+- name: regression/issue_766__get_with_context.yml
+  hosts: localhost
+  # Gather facts to use *and* to trigger any "could not recover task_vars" error
+  # https://github.com/mitogen-hq/mitogen/pull/1215#issuecomment-2596421111
+  gather_facts: true
+  vars:
+    netconf_container_image: ghcr.io/mitogen-hq/sysrepo-netopeer2:latest
+    netconf_container_name: sysprep
+    netconf_container_port: 8030
+
+  tasks:
+    - meta: end_play
+      when:
+        # Podman can be installed on macOS, but authenticating to gchr.io isn't
+        # worth the trouble right now.
+        # See also
+        #   - issue_655__wait_for_connection_error.yml
+        - ansible_facts.system == 'Darwin'
+
+    - meta: end_play
+      when:
+        # A failure during the ansible.netcommon.netconf_get task, when run
+        # with Ansible 4 (ansible-core 2.11) & associated collections.
+        # ansible.module_utils.connection.ConnectionError: Method not found
+        # https://github.com/mitogen-hq/mitogen/actions/runs/12854359099/job/35838635886
+        - ansible_version.full is version('2.11', '>=', strict=True)
+        - ansible_version.full is version('2.12', '<', strict=True)
+
+    - block:
+        - name: Start container
+          command:
+            cmd: >-
+              podman run
+              --name "{{ netconf_container_name }}"
+              --detach
+              --rm
+              --publish "{{ netconf_container_port }}:830"
+              "{{ netconf_container_image }}"
+          changed_when: true
+
+        - name: Wait for container
+          # TODO robust condition. wait_for + search_regex? wait_for_connection?
+          wait_for:
+            timeout: 5
+
+        - name: Get running configuration and state data
+          vars:
+            ansible_connection: netconf
+            ansible_user: netconf
+            ansible_password: netconf
+            ansible_port: "{{ netconf_container_port }}"
+            ansible_host_key_checking: false
+            ansible_python_interpreter: "{{ ansible_playbook_python }}"
+          ansible.netcommon.netconf_get:
+
+      always:
+        - name: Cleanup container
+          command:
+            cmd: podman stop "{{ netconf_container_name }}"
+          changed_when: true
+  tags:
+    - issue_766

tests/ansible/requirements.txt

@@ -1,7 +1,13 @@
-paramiko==2.3.2  # Last 2.6-compat version.
+paramiko==2.12.0;           python_version <= '2.7'
+paramiko==3.5.0;            python_version >= '3.6'
+
 # Incompatible with pip >= 72, due to removal of `setup.py test`:
 #   ModuleNotFoundError: No module named 'setuptools.command.test'
 # https://github.com/pypa/setuptools/issues/4519
 hdrhistogram==0.6.1
+
+ncclient==0.6.13; python_version <= '2.7'
+ncclient==0.6.16; python_version > '2.7'
+
 PyYAML==3.11; python_version < '2.7'
 PyYAML==5.3.1; python_version >= '2.7'  # Latest release (Jan 2021)