loader: support extra modules in sys.path

salt/loader.py

@@ -260,6 +260,7 @@ def minion_mods(
         whitelist=whitelist,
         loaded_base_name=loaded_base_name,
         static_modules=static_modules,
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
 
     ret.pack["__salt__"] = ret
@@ -344,7 +345,13 @@ def engines(opts, functions, runners, utils, proxy=None):
         "__proxy__": proxy,
         "__utils__": utils,
     }
-    return LazyLoader(_module_dirs(opts, "engines"), opts, tag="engines", pack=pack,)
+    return LazyLoader(
+        _module_dirs(opts, "engines"),
+        opts,
+        tag="engines",
+        pack=pack,
+        extra_module_dirs=utils.module_dirs if utils else None,
+    )
 
 
 def proxy(opts, functions=None, returners=None, whitelist=None, utils=None):
@@ -356,6 +363,7 @@ def proxy(opts, functions=None, returners=None, whitelist=None, utils=None):
         opts,
         tag="proxy",
         pack={"__salt__": functions, "__ret__": returners, "__utils__": utils},
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
 
     ret.pack["__proxy__"] = ret
@@ -393,11 +401,13 @@ def pillars(opts, functions, context=None):
     """
     Returns the pillars modules
     """
+    _utils = utils(opts)
     ret = LazyLoader(
         _module_dirs(opts, "pillar"),
         opts,
         tag="pillar",
-        pack={"__salt__": functions, "__context__": context, "__utils__": utils(opts)},
+        pack={"__salt__": functions, "__context__": context, "__utils__": _utils},
+        extra_module_dirs=_utils.module_dirs,
     )
     ret.pack["__ext_pillar__"] = ret
     return FilterDictWrapper(ret, ".ext_pillar")
@@ -487,12 +497,14 @@ def fileserver(opts, backends):
     """
     Returns the file server modules
     """
+    _utils = utils(opts)
     return LazyLoader(
         _module_dirs(opts, "fileserver"),
         opts,
         tag="fileserver",
         whitelist=backends,
-        pack={"__utils__": utils(opts)},
+        pack={"__utils__": _utils},
+        extra_module_dirs=_utils.module_dirs,
     )
 
 
@@ -506,6 +518,7 @@ def roster(opts, runner=None, utils=None, whitelist=None):
         tag="roster",
         whitelist=whitelist,
         pack={"__runner__": runner, "__utils__": utils},
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
 
 
@@ -546,6 +559,7 @@ def states(
         tag="states",
         pack={"__salt__": functions, "__proxy__": proxy or {}},
         whitelist=whitelist,
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
     ret.pack["__states__"] = ret
     ret.pack["__utils__"] = utils
@@ -659,12 +673,14 @@ def grain_funcs(opts, proxy=None):
           __opts__ = salt.config.minion_config('/etc/salt/minion')
           grainfuncs = salt.loader.grain_funcs(__opts__)
     """
+    _utils = utils(opts, proxy=proxy)
     ret = LazyLoader(
         _module_dirs(opts, "grains", "grain", ext_type_dirs="grains_dirs",),
         opts,
         tag="grains",
+        extra_module_dirs=_utils.module_dirs,
     )
-    ret.pack["__utils__"] = utils(opts, proxy=proxy)
+    ret.pack["__utils__"] = _utils
     return ret
 
 
@@ -939,6 +955,7 @@ def runner(opts, utils=None, context=None, whitelist=None):
         tag="runners",
         pack={"__utils__": utils, "__context__": context},
         whitelist=whitelist,
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
     # TODO: change from __salt__ to something else, we overload __salt__ too much
     ret.pack["__salt__"] = ret
@@ -974,6 +991,7 @@ def sdb(opts, functions=None, whitelist=None, utils=None):
             "__salt__": minion_mods(opts, utils=utils),
         },
         whitelist=whitelist,
+        extra_module_dirs=utils.module_dirs if utils else None,
     )
 
 
@@ -1007,6 +1025,7 @@ def clouds(opts):
     """
     Return the cloud functions
     """
+    _utils = salt.loader.utils(opts)
     # Let's bring __active_provider_name__, defaulting to None, to all cloud
     # drivers. This will get temporarily updated/overridden with a context
     # manager when needed.
@@ -1020,7 +1039,8 @@ def clouds(opts):
         ),
         opts,
         tag="clouds",
-        pack={"__utils__": salt.loader.utils(opts), "__active_provider_name__": None},
+        pack={"__utils__": _utils, "__active_provider_name__": None},
+        extra_module_dirs=_utils.module_dirs,
     )
     for funcname in LIBCLOUD_FUNCS_NOT_SUPPORTED:
         log.trace(
@@ -1136,6 +1156,7 @@ class LazyLoader(salt.utils.lazy.LazyDict):
     :param bool virtual_enable: Whether or not to respect the __virtual__ function when loading modules.
     :param str virtual_funcs: The name of additional functions in the module to call to verify its functionality.
                                 If not true, the module will not load.
+    :param list extra_module_dirs: A list of directories that will be able to import from
     :returns: A LazyLoader object which functions as a dictionary. Keys are 'module.function' and values
     are function references themselves which are loaded on-demand.
     # TODO:
@@ -1158,6 +1179,7 @@ def __init__(
         static_modules=None,
         proxy=None,
         virtual_funcs=None,
+        extra_module_dirs=None,
     ):  # pylint: disable=W0231
         """
         In pack, if any of the values are None they will be replaced with an
@@ -1201,6 +1223,9 @@ def __init__(
             virtual_funcs = []
         self.virtual_funcs = virtual_funcs
 
+        self.extra_module_dirs = extra_module_dirs if extra_module_dirs else []
+        self._clean_module_dirs = []
+
         self.disabled = set(
             self.opts.get(
                 "disable_{0}{1}".format(self.tag, "" if self.tag[-1] == "s" else "s"),
@@ -1515,12 +1540,44 @@ def _reload_submodules(self, mod):
                 reload_module(submodule)
                 self._reload_submodules(submodule)
 
+    def __populate_sys_path(self):
+        for directory in self.extra_module_dirs:
+            if directory not in sys.path:
+                sys.path.append(directory)
+                self._clean_module_dirs.append(directory)
+
+    def __clean_sys_path(self):
+        invalidate_path_importer_cache = False
+        for directory in self._clean_module_dirs:
+            if directory in sys.path:
+                sys.path.remove(directory)
+                invalidate_path_importer_cache = True
+        self._clean_module_dirs = []
+
+        # Be sure that sys.path_importer_cache do not contains any
+        # invalid FileFinder references
+        if USE_IMPORTLIB:
+            importlib.invalidate_caches()
+
+        # Because we are mangling with importlib, we can find from
+        # time to time an invalidation issue with
+        # sys.path_importer_cache, that requires the removal of
+        # FileFinder that remain None for the extra_module_dirs
+        if invalidate_path_importer_cache:
+            for directory in self.extra_module_dirs:
+                if (
+                    directory in sys.path_importer_cache
+                    and sys.path_importer_cache[directory] is None
+                ):
+                    del sys.path_importer_cache[directory]
+
     def _load_module(self, name):
         mod = None
         fpath, suffix = self.file_mapping[name][:2]
         self.loaded_files.add(name)
         fpath_dirname = os.path.dirname(fpath)
         try:
+            self.__populate_sys_path()
             sys.path.append(fpath_dirname)
             if suffix == ".pyx":
                 mod = pyximport.load_module(name, fpath, tempfile.gettempdir())
@@ -1659,6 +1716,7 @@ def _load_module(self, name):
             return False
         finally:
             sys.path.remove(fpath_dirname)
+            self.__clean_sys_path()
 
         if hasattr(mod, "__opts__"):
             mod.__opts__.update(self.opts)

tests/unit/test_loader.py

@@ -133,6 +133,96 @@ def test_depends(self):
         self.assertTrue(self.module_name + ".not_loaded" not in self.loader)
 
 
+loader_template_module = """
+import my_utils
+
+def run():
+    return my_utils.run()
+"""
+
+loader_template_utils = """
+def run():
+    return True
+"""
+
+
+class LazyLoaderUtilsTest(TestCase):
+    """
+    Test the loader
+    """
+
+    module_name = "lazyloaderutilstest"
+    utils_name = "my_utils"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.opts = salt.config.minion_config(None)
+        cls.opts["grains"] = salt.loader.grains(cls.opts)
+        if not os.path.isdir(RUNTIME_VARS.TMP):
+            os.makedirs(RUNTIME_VARS.TMP)
+
+    def setUp(self):
+        # Setup the module
+        self.module_dir = tempfile.mkdtemp(dir=RUNTIME_VARS.TMP)
+        self.module_file = os.path.join(
+            self.module_dir, "{}.py".format(self.module_name)
+        )
+        with salt.utils.files.fopen(self.module_file, "w") as fh:
+            fh.write(salt.utils.stringutils.to_str(loader_template_module))
+            fh.flush()
+            os.fsync(fh.fileno())
+
+        self.utils_dir = tempfile.mkdtemp(dir=RUNTIME_VARS.TMP)
+        self.utils_file = os.path.join(self.utils_dir, "{}.py".format(self.utils_name))
+        with salt.utils.files.fopen(self.utils_file, "w") as fh:
+            fh.write(salt.utils.stringutils.to_str(loader_template_utils))
+            fh.flush()
+            os.fsync(fh.fileno())
+
+    def tearDown(self):
+        shutil.rmtree(self.module_dir)
+        if os.path.isdir(self.module_dir):
+            shutil.rmtree(self.module_dir)
+        shutil.rmtree(self.utils_dir)
+        if os.path.isdir(self.utils_dir):
+            shutil.rmtree(self.utils_dir)
+        del self.module_dir
+        del self.module_file
+        del self.utils_dir
+        del self.utils_file
+
+        if self.module_name in sys.modules:
+            del sys.modules[self.module_name]
+        if self.utils_name in sys.modules:
+            del sys.modules[self.utils_name]
+
+    @classmethod
+    def tearDownClass(cls):
+        del cls.opts
+
+    def test_utils_found(self):
+        """
+        Test that the extra module directory is available for imports
+        """
+        loader = salt.loader.LazyLoader(
+            [self.module_dir],
+            copy.deepcopy(self.opts),
+            tag="module",
+            extra_module_dirs=[self.utils_dir],
+        )
+        self.assertTrue(inspect.isfunction(loader[self.module_name + ".run"]))
+        self.assertTrue(loader[self.module_name + ".run"]())
+
+    def test_utils_not_found(self):
+        """
+        Test that the extra module directory is not available for imports
+        """
+        loader = salt.loader.LazyLoader(
+            [self.module_dir], copy.deepcopy(self.opts), tag="module"
+        )
+        self.assertTrue(self.module_name + ".run" not in loader)
+
+
 class LazyLoaderVirtualEnabledTest(TestCase):
     """
     Test the base loader of salt.