add optional module_utils import support (#73832) (#73918)

Treat core and collections module_utils imports nested within any Python block statement (eg, `try`, `if`) as optional. This allows Ansible modules to implement runtime fallback behavior for missing module_utils (eg from a newer version of ansible-core), where previously, the module payload builder would always fail when unable to locate a module_util (regardless of any runtime behavior the module may implement).

* sanity test fixes

(cherry picked from commit 3e1f6484d77f2d7546952cfa22a8534d74ed3dc6)

6 files changed, 162 insertions, 21 deletions

diff --git a/changelogs/fragments/optional_module_utils.yml b/changelogs/fragments/optional_module_utils.yml
new file mode 100644
index 0000000000..e9ff22c4c5
--- /dev/null
+++ b/changelogs/fragments/optional_module_utils.yml
@@ -0,0 +1,4 @@
+minor_changes:
+- module payload builder - module_utils imports in any nested block (eg, ``try``, ``if``) are treated as optional during
+  module payload builds; this allows modules to implement runtime fallback behavior for module_utils that do not exist
+  in older versions of Ansible.
diff --git a/lib/ansible/executor/module_common.py b/lib/ansible/executor/module_common.py
index 15f2506a75..623b5e29c1 100644
--- a/lib/ansible/executor/module_common.py
+++ b/lib/ansible/executor/module_common.py
@@ -29,6 +29,7 @@ import shlex
 import zipfile
 import re
 import pkgutil
+from ast import AST, Import, ImportFrom
 from io import BytesIO
 
 from ansible.release import __version__, __author__
@@ -65,7 +66,7 @@ except NameError:
 
 display = Display()
 
-ModuleUtilsProcessEntry = namedtuple('ModuleUtilsInfo', ['name_parts', 'is_ambiguous', 'has_redirected_child'])
+ModuleUtilsProcessEntry = namedtuple('ModuleUtilsInfo', ['name_parts', 'is_ambiguous', 'has_redirected_child', 'is_optional'])
 
 REPLACER = b"#<<INCLUDE_ANSIBLE_MODULE_COMMON>>"
 REPLACER_VERSION = b"\"<<ANSIBLE_VERSION>>\""
@@ -442,7 +443,7 @@ NEW_STYLE_PYTHON_MODULE_RE = re.compile(
 
 
 class ModuleDepFinder(ast.NodeVisitor):
-    def __init__(self, module_fqn, is_pkg_init=False, *args, **kwargs):
+    def __init__(self, module_fqn, tree, is_pkg_init=False, *args, **kwargs):
         """
         Walk the ast tree for the python module.
         :arg module_fqn: The fully qualified name to reach this module in dotted notation.
@@ -466,10 +467,37 @@ class ModuleDepFinder(ast.NodeVisitor):
         .. seealso:: :python3:class:`ast.NodeVisitor`
         """
         super(ModuleDepFinder, self).__init__(*args, **kwargs)
+        self._tree = tree  # squirrel this away so we can compare node parents to it
         self.submodules = set()
+        self.optional_imports = set()
         self.module_fqn = module_fqn
         self.is_pkg_init = is_pkg_init
 
+        self._visit_map = {
+            Import: self.visit_Import,
+            ImportFrom: self.visit_ImportFrom,
+        }
+
+        self.visit(tree)
+
+    def generic_visit(self, node):
+        """Overridden ``generic_visit`` that makes some assumptions about our
+        use case, and improves performance by calling visitors directly instead
+        of calling ``visit`` to offload calling visitors.
+        """
+        generic_visit = self.generic_visit
+        visit_map = self._visit_map
+        for field, value in ast.iter_fields(node):
+            if isinstance(value, list):
+                for item in value:
+                    if isinstance(item, (Import, ImportFrom)):
+                        item.parent = node
+                        visit_map[item.__class__](item)
+                    elif isinstance(item, AST):
+                        generic_visit(item)
+
+    visit = generic_visit
+
     def visit_Import(self, node):
         """
         Handle import ansible.module_utils.MODLIB[.MODLIBn] [as asname]
@@ -482,6 +510,9 @@ class ModuleDepFinder(ast.NodeVisitor):
                     alias.name.startswith('ansible_collections.')):
                 py_mod = tuple(alias.name.split('.'))
                 self.submodules.add(py_mod)
+                # if the import's parent is the root document, it's a required import, otherwise it's optional
+                if node.parent != self._tree:
+                    self.optional_imports.add(py_mod)
         self.generic_visit(node)
 
     def visit_ImportFrom(self, node):
@@ -543,6 +574,9 @@ class ModuleDepFinder(ast.NodeVisitor):
         if py_mod:
             for alias in node.names:
                 self.submodules.add(py_mod + (alias.name,))
+                # if the import's parent is the root document, it's a required import, otherwise it's optional
+                if node.parent != self._tree:
+                    self.optional_imports.add(py_mod + (alias.name,))
 
         self.generic_visit(node)
 
@@ -606,12 +640,13 @@ def _get_shebang(interpreter, task_vars, templar, args=tuple()):
 
 
 class ModuleUtilLocatorBase:
-    def __init__(self, fq_name_parts, is_ambiguous=False, child_is_redirected=False):
+    def __init__(self, fq_name_parts, is_ambiguous=False, child_is_redirected=False, is_optional=False):
         self._is_ambiguous = is_ambiguous
         # a child package redirection could cause intermediate package levels to be missing, eg
         # from ansible.module_utils.x.y.z import foo; if x.y.z.foo is redirected, we may not have packages on disk for
         # the intermediate packages x.y.z, so we'll need to supply empty packages for those
         self._child_is_redirected = child_is_redirected
+        self._is_optional = is_optional
         self.found = False
         self.redirected = False
         self.fq_name_parts = fq_name_parts
@@ -640,6 +675,8 @@ class ModuleUtilLocatorBase:
         try:
             collection_metadata = _get_collection_metadata(self._collection_name)
         except ValueError as ve:  # collection not found or some other error related to collection load
+            if self._is_optional:
+                return False
             raise AnsibleError('error processing module_util {0} loading redirected collection {1}: {2}'
                                .format('.'.join(name_parts), self._collection_name, to_native(ve)))
 
@@ -798,8 +835,8 @@ class LegacyModuleUtilLocator(ModuleUtilLocatorBase):
 
 
 class CollectionModuleUtilLocator(ModuleUtilLocatorBase):
-    def __init__(self, fq_name_parts, is_ambiguous=False, child_is_redirected=False):
-        super(CollectionModuleUtilLocator, self).__init__(fq_name_parts, is_ambiguous, child_is_redirected)
+    def __init__(self, fq_name_parts, is_ambiguous=False, child_is_redirected=False, is_optional=False):
+        super(CollectionModuleUtilLocator, self).__init__(fq_name_parts, is_ambiguous, child_is_redirected, is_optional)
 
         if fq_name_parts[0] != 'ansible_collections':
             raise Exception('CollectionModuleUtilLocator can only locate from ansible_collections, got {0}'.format(fq_name_parts))
@@ -891,20 +928,19 @@ def recursive_finder(name, module_fqn, module_data, zf):
     except (SyntaxError, IndentationError) as e:
         raise AnsibleError("Unable to import %s due to %s" % (name, e.msg))
 
-    finder = ModuleDepFinder(module_fqn)
-    finder.visit(tree)
+    finder = ModuleDepFinder(module_fqn, tree)
 
     # the format of this set is a tuple of the module name and whether or not the import is ambiguous as a module name
     # or an attribute of a module (eg from x.y import z <-- is z a module or an attribute of x.y?)
-    modules_to_process = [ModuleUtilsProcessEntry(m, True, False) for m in finder.submodules]
+    modules_to_process = [ModuleUtilsProcessEntry(m, True, False, is_optional=m in finder.optional_imports) for m in finder.submodules]
 
     # HACK: basic is currently always required since module global init is currently tied up with AnsiballZ arg input
-    modules_to_process.append(ModuleUtilsProcessEntry(('ansible', 'module_utils', 'basic'), False, False))
+    modules_to_process.append(ModuleUtilsProcessEntry(('ansible', 'module_utils', 'basic'), False, False, is_optional=False))
 
     # we'll be adding new modules inline as we discover them, so just keep going til we've processed them all
     while modules_to_process:
         modules_to_process.sort()  # not strictly necessary, but nice to process things in predictable and repeatable order
-        py_module_name, is_ambiguous, child_is_redirected = modules_to_process.pop(0)
+        py_module_name, is_ambiguous, child_is_redirected, is_optional = modules_to_process.pop(0)
 
         if py_module_name in py_module_cache:
             # this is normal; we'll often see the same module imported many times, but we only need to process it once
@@ -914,7 +950,8 @@ def recursive_finder(name, module_fqn, module_data, zf):
             module_info = LegacyModuleUtilLocator(py_module_name, is_ambiguous=is_ambiguous,
                                                   mu_paths=module_utils_paths, child_is_redirected=child_is_redirected)
         elif py_module_name[0] == 'ansible_collections':
-            module_info = CollectionModuleUtilLocator(py_module_name, is_ambiguous=is_ambiguous, child_is_redirected=child_is_redirected)
+            module_info = CollectionModuleUtilLocator(py_module_name, is_ambiguous=is_ambiguous,
+                                                      child_is_redirected=child_is_redirected, is_optional=is_optional)
         else:
             # FIXME: dot-joined result
             display.warning('ModuleDepFinder improperly found a non-module_utils import %s'
@@ -923,6 +960,9 @@ def recursive_finder(name, module_fqn, module_data, zf):
 
         # Could not find the module.  Construct a helpful error message.
         if not module_info.found:
+            if is_optional:
+                # this was a best-effort optional import that we couldn't find, oh well, move along...
+                continue
             # FIXME: use dot-joined candidate names
             msg = 'Could not find imported module support code for {0}.  Looked for ({1})'.format(module_fqn, module_info.candidate_names_joined)
             raise AnsibleError(msg)
@@ -938,9 +978,9 @@ def recursive_finder(name, module_fqn, module_data, zf):
         except (SyntaxError, IndentationError) as e:
             raise AnsibleError("Unable to import %s due to %s" % (module_info.fq_name_parts, e.msg))
 
-        finder = ModuleDepFinder('.'.join(module_info.fq_name_parts), module_info.is_package)
-        finder.visit(tree)
-        modules_to_process.extend(ModuleUtilsProcessEntry(m, True, False) for m in finder.submodules if m not in py_module_cache)
+        finder = ModuleDepFinder('.'.join(module_info.fq_name_parts), tree, module_info.is_package)
+        modules_to_process.extend(ModuleUtilsProcessEntry(m, True, False, is_optional=m in finder.optional_imports)
+                                  for m in finder.submodules if m not in py_module_cache)
 
         # we've processed this item, add it to the output list
         py_module_cache[module_info.fq_name_parts] = (module_info.source_code, module_info.output_path)
@@ -951,7 +991,7 @@ def recursive_finder(name, module_fqn, module_data, zf):
             accumulated_pkg_name.append(pkg)  # we're accumulating this across iterations
             normalized_name = tuple(accumulated_pkg_name)  # extra machinations to get a hashable type (list is not)
             if normalized_name not in py_module_cache:
-                modules_to_process.append((normalized_name, False, module_info.redirected))
+                modules_to_process.append(ModuleUtilsProcessEntry(normalized_name, False, module_info.redirected, is_optional=is_optional))
 
     for py_module_name in py_module_cache:
         py_module_file_name = py_module_cache[py_module_name][1]
diff --git a/test/integration/targets/module_utils/collections/ansible_collections/testns/testcoll/plugins/module_utils/legit.py b/test/integration/targets/module_utils/collections/ansible_collections/testns/testcoll/plugins/module_utils/legit.py
new file mode 100644
index 0000000000..b9d63482a1
--- /dev/null
+++ b/test/integration/targets/module_utils/collections/ansible_collections/testns/testcoll/plugins/module_utils/legit.py
@@ -0,0 +1,6 @@
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+
+def importme():
+    return "successfully imported from testns.testcoll"
diff --git a/test/integration/targets/module_utils/library/test_failure.py b/test/integration/targets/module_utils/library/test_failure.py
index e1a87c2e71..e5257aefc9 100644
--- a/test/integration/targets/module_utils/library/test_failure.py
+++ b/test/integration/targets/module_utils/library/test_failure.py
@@ -4,11 +4,9 @@ results = {}
 # Test that we are rooted correctly
 # Following files:
 #   module_utils/yak/zebra/foo.py
-try:
-    from ansible.module_utils.zebra import foo
-    results['zebra'] = foo.data
-except ImportError:
-    results['zebra'] = 'Failed in module as expected but incorrectly did not fail in AnsiballZ construction'
+from ansible.module_utils.zebra import foo
+
+results['zebra'] = foo.data
 
 from ansible.module_utils.basic import AnsibleModule
 AnsibleModule(argument_spec=dict()).exit_json(**results)
diff --git a/test/integration/targets/module_utils/library/test_optional.py b/test/integration/targets/module_utils/library/test_optional.py
new file mode 100644
index 0000000000..4d0225d96f
--- /dev/null
+++ b/test/integration/targets/module_utils/library/test_optional.py
@@ -0,0 +1,84 @@
+#!/usr/bin/python
+# Most of these names are only available via PluginLoader so pylint doesn't
+# know they exist
+# pylint: disable=no-name-in-module
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+from ansible.module_utils.basic import AnsibleModule
+
+# internal constants to keep pylint from griping about constant-valued conditionals
+_private_false = False
+_private_true = True
+
+# module_utils import statements nested below any block are considered optional "best-effort" for AnsiballZ to include.
+# test a number of different import shapes and nesting types to exercise this...
+
+# first, some nested imports that should succeed...
+try:
+    from ansible.module_utils.urls import fetch_url as yep1
+except ImportError:
+    yep1 = None
+
+try:
+    import ansible.module_utils.common.text.converters as yep2
+except ImportError:
+    yep2 = None
+
+try:
+    # optional import from a legit collection
+    from ansible_collections.testns.testcoll.plugins.module_utils.legit import importme as yep3
+except ImportError:
+    yep3 = None
+
+# and a bunch that should fail to be found, but not break the module_utils payload build in the process...
+try:
+    from ansible.module_utils.bogus import fromnope1
+except ImportError:
+    fromnope1 = None
+
+if _private_false:
+    from ansible.module_utils.alsobogus import fromnope2
+else:
+    fromnope2 = None
+
+try:
+    import ansible.module_utils.verybogus
+    nope1 = ansible.module_utils.verybogus
+except ImportError:
+    nope1 = None
+
+# deepish nested with multiple block types- make sure the AST walker made it all the way down
+try:
+    if _private_true:
+        if _private_true:
+            if _private_true:
+                if _private_true:
+                    try:
+                        import ansible.module_utils.stillbogus as nope2
+                    except ImportError:
+                        raise
+except ImportError:
+    nope2 = None
+
+try:
+    # optional import from a valid collection with an invalid package
+    from ansible_collections.testns.testcoll.plugins.module_utils.bogus import collnope1
+except ImportError:
+    collnope1 = None
+
+try:
+    # optional import from a bogus collection
+    from ansible_collections.bogusns.boguscoll.plugins.module_utils.bogus import collnope2
+except ImportError:
+    collnope2 = None
+
+module = AnsibleModule(argument_spec={})
+
+if not all([yep1, yep2, yep3]):
+    module.fail_json(msg='one or more existing optional imports did not resolve')
+
+if any([fromnope1, fromnope2, nope1, nope2, collnope1, collnope2]):
+    module.fail_json(msg='one or more missing optional imports resolved unexpectedly')
+
+module.exit_json(msg='all missing optional imports behaved as expected')
diff --git a/test/integration/targets/module_utils/module_utils_test.yml b/test/integration/targets/module_utils/module_utils_test.yml
index 943bf4ee25..0550b9f016 100644
--- a/test/integration/targets/module_utils/module_utils_test.yml
+++ b/test/integration/targets/module_utils/module_utils_test.yml
@@ -43,7 +43,6 @@
     ignore_errors: True
     register: result
 
-  - debug: var=result
   - name: Make sure we failed in AnsiBallZ
     assert:
       that:
@@ -60,3 +59,13 @@
       that:
         - result.deprecations[0].msg == "Alias 'baz' is deprecated. See the module docs for more information"
         - result.deprecations[0].version == '9.99'
+
+
+  - name: Test that optional imports behave properly
+    test_optional:
+    register: optionaltest
+
+  - assert:
+      that:
+      - optionaltest is success
+      - optionaltest.msg == 'all missing optional imports behaved as expected'
\ No newline at end of file