Merge pull request #6183: Add support for loading .pth files in st2 venv

cognifloyd · web-flow · commit a56f527f4679 · 2024-04-23T10:40:55.000-05:00
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -253,13 +253,21 @@ jobs:
         run: |
           script -e -c "make .ci-prepare-integration" && exit 0
       - name: Extend the path for upcoming tasks
-        run: |
-          echo ${HOME}/work/st2/st2/virtualenv/bin
-          echo ${HOME}/work/st2/st2/virtualenv/bin >> $GITHUB_PATH
+        # pants uses PEP 660 editable wheels to add our code to the virtualenv.
+        # But PEP 660 editable wheels do not include 'scripts'.
+        # https://peps.python.org/pep-0660/#limitations
+        # So, we need to include each bin dir in PATH instead of virtualenv/bin.
+        run: |
+          for component_bin in ${GITHUB_WORKSPACE}/st2*/bin; do
+          echo ${component_bin} | tee -a $GITHUB_PATH
+          done
+          echo ${GITHUB_WORKSPACE}/virtualenv/bin | tee -a $GITHUB_PATH
       - name: Create symlinks to find the binaries when running st2 actions
+        # st2 is actually a console_script entry point, not just a 'script'
+        # so it IS included in the virtualenv. But, st2-run-pack-tests might not be included.
         run: |
-          ln -s ${HOME}/work/st2/st2/virtualenv/bin/st2 /usr/local/bin/st2
-          ln -s ${HOME}/work/st2/st2/virtualenv/bin/st2-run-pack-tests /usr/local/bin/st2-run-pack-tests
+          ln -s ${GITHUB_WORKSPACE}/virtualenv/bin/st2 /usr/local/bin/st2
+          ln -s ${GITHUB_WORKSPACE}/st2common/bin/st2-run-pack-tests /usr/local/bin/st2-run-pack-tests
       - name: Install st2client
         timeout-minutes: 5
         run: |
@@ -270,7 +278,7 @@ jobs:
         env:
           ST2_CONF: /home/runner/work/st2/st2/conf/st2.ci.conf
         run: |
-          sudo -E ST2_AUTH_TOKEN=$(st2 auth testu -p 'testp' -t) PATH=${PATH} virtualenv/bin/st2-self-check
+          sudo -E ST2_AUTH_TOKEN=$(st2 auth testu -p 'testp' -t) PATH=${PATH} st2common/bin/st2-self-check
       - name: Compress Service Logs Before upload
         if: ${{ failure() }}
         run: |
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -24,10 +24,12 @@ Added
 * Continue introducing `pants <https://www.pantsbuild.org/docs>`_ to improve DX (Developer Experience)
   working on StackStorm, improve our security posture, and improve CI reliability thanks in part
   to pants' use of PEX lockfiles. This is not a user-facing addition.
-  #6118 #6141 #6133 #6120 #6181
+  #6118 #6141 #6133 #6120 #6181 #6183
   Contributed by @cognifloyd
 * Build of ST2 EL9 packages #6153
   Contributed by @amanda11
+* Ensure `.pth` files in the st2 virtualenv get loaded by pack virtualenvs. #6183
+  Contributed by @cognifloyd
 
 3.8.1 - December 13, 2023
 -------------------------
diff --git a/st2common/bin/st2-run-pack-tests b/st2common/bin/st2-run-pack-tests
@@ -63,7 +63,8 @@ ENABLE_TIMING=false
 
 VIRTUALENV_ACTIVATED=false
 
-STACKSTORM_VIRTUALENV_BIN="/opt/stackstorm/st2/bin"
+STACKSTORM_VIRTUALENV="/opt/stackstorm/st2"
+STACKSTORM_VIRTUALENV_BIN="${STACKSTORM_VIRTUALENV}/bin"
 STACKSTORM_VIRTUALENV_PYTHON_BINARY="${STACKSTORM_VIRTUALENV_BIN}/python"
 
 ####################
@@ -194,6 +195,13 @@ if [ "${CREATE_VIRTUALENV}" = true ]; then
         mkdir -p ${VIRTUALENVS_DIR}
         virtualenv --no-download --system-site-packages ${VIRTUALENV_DIR}
 
+        if [ -f "${STACKSTORM_VIRTUALENV_PYTHON_BINARY}" ]; then
+            # ensure any .pth files in st2 venv get loaded with the pack venv too.
+            ST2_SITE_PACKAGES=$(${STACKSTORM_VIRTUALENV_PYTHON_BINARY} -c "import sysconfig;print(sysconfig.get_path('platlib'))")
+            PACK_SITE_PACKAGES=$(${VIRTUALENV_DIR}/bin/python3 -c "import sysconfig;print(sysconfig.get_path('platlib'))")
+            echo "import sys; addsitedir('${ST2_SITE_PACKAGES}', known_paths)" > "${PACK_SITE_PACKAGES}/zzzzzzzzzz__st2__.pth"
+        fi
+
         # Activate the virtualenv
         activate_virtualenv
 
@@ -346,16 +354,26 @@ if [ "${ENABLE_TIMING}" = true ]; then
   NOSE_OPTS+=(--with-timer)
 fi
 
+NOSE=(nosetests)
+if head -n 1 $(command -v nosetests) | grep -q ' -sE$'; then
+    # workaround pants+pex default of hermetic scripts so we can run nosetests with PYTHONPATH
+    if [ -f "${STACKSTORM_VIRTUALENV_PYTHON_BINARY}" ]; then
+        NOSE=(${STACKSTORM_VIRTUALENV_PYTHON_BINARY} -m "nose")
+    else
+        NOSE=(python3 -m "nose")
+    fi
+fi
+
 # Change to the pack's directory (required for test coverage reporting)
 pushd ${PACK_PATH} > /dev/null
 
 # Execute the tests
 if [ "${TEST_LOCATION}" ]; then
     # Run a specific test file, class or method
-    nosetests ${NOSE_OPTS[@]} ${TEST_LOCATION}
+    ${NOSE[@]} ${NOSE_OPTS[@]} ${TEST_LOCATION}
 else
     # Run all tests inside the pack
-    nosetests ${NOSE_OPTS[@]} ${PACK_TESTS_PATH}
+    ${NOSE[@]} ${NOSE_OPTS[@]} ${PACK_TESTS_PATH}
 fi
 TESTS_EXIT_CODE=$?
 
diff --git a/st2common/st2common/util/sandboxing.py b/st2common/st2common/util/sandboxing.py
@@ -24,6 +24,7 @@
 import os
 import sys
 from sysconfig import get_path
+from typing import Optional
 
 from oslo_config import cfg
 
@@ -32,20 +33,51 @@
 from st2common.content.utils import get_pack_base_path
 
 
-def get_python_lib():
-    """Replacement for distutil.sysconfig.get_python_lib, returns a string with the python platform lib path (to site-packages)"""
-    return get_path("platlib")
-
-
 __all__ = [
+    "get_site_packages_dir",
+    "get_virtualenv_lib_path",
     "get_sandbox_python_binary_path",
     "get_sandbox_python_path",
     "get_sandbox_python_path_for_python_action",
     "get_sandbox_path",
     "get_sandbox_virtualenv_path",
+    "is_in_virtualenv",
 ]
 
 
+def get_site_packages_dir() -> str:
+    """Returns a string with the python platform lib path (to site-packages)."""
+    # This assumes we are running in the primary st2 virtualenv (typically /opt/stackstorm/st2)
+    site_packages_dir = get_path("platlib")
+
+    sys_prefix = os.path.abspath(sys.prefix)
+    if sys_prefix not in site_packages_dir:
+        raise ValueError(
+            f'The file with "{sys_prefix}" is not found in "{site_packages_dir}".'
+        )
+
+    return site_packages_dir
+
+
+def get_virtualenv_lib_path(virtualenv_path: str) -> str:
+    """Returns the path to a virtualenv's lib/python3.* directory."""
+    if not (virtualenv_path and os.path.isdir(virtualenv_path)):
+        raise OSError(
+            f"virtualenv_path must be an existing directory. virtualenv_path={virtualenv_path}"
+        )
+
+    pack_virtualenv_lib_path = os.path.join(virtualenv_path, "lib")
+
+    virtualenv_directories = os.listdir(pack_virtualenv_lib_path)
+    virtualenv_directories = [
+        dir_name
+        for dir_name in virtualenv_directories
+        if fnmatch.fnmatch(dir_name, "python*")
+    ]
+
+    return os.path.join(pack_virtualenv_lib_path, virtualenv_directories[0])
+
+
 def get_sandbox_python_binary_path(pack=None):
     """
     Return path to the Python binary for the provided pack.
@@ -114,14 +146,7 @@ def get_sandbox_python_path(inherit_from_parent=True, inherit_parent_virtualenv=
 
     if inherit_parent_virtualenv and is_in_virtualenv():
         # We are running inside virtualenv
-        site_packages_dir = get_python_lib()
-
-        sys_prefix = os.path.abspath(sys.prefix)
-        if sys_prefix not in site_packages_dir:
-            raise ValueError(
-                f'The file with "{sys_prefix}" is not found in "{site_packages_dir}".'
-            )
-
+        site_packages_dir = get_site_packages_dir()
         sandbox_python_path.append(site_packages_dir)
 
     sandbox_python_path = ":".join(sandbox_python_path)
@@ -146,30 +171,20 @@ def get_sandbox_python_path_for_python_action(
     virtualenv_path = get_sandbox_virtualenv_path(pack=pack)
 
     if virtualenv_path and os.path.isdir(virtualenv_path):
-        pack_virtualenv_lib_path = os.path.join(virtualenv_path, "lib")
-
-        virtualenv_directories = os.listdir(pack_virtualenv_lib_path)
-        virtualenv_directories = [
-            dir_name
-            for dir_name in virtualenv_directories
-            if fnmatch.fnmatch(dir_name, "python*")
-        ]
-
         # Add the pack's lib directory (lib/python3.x) in front of the PYTHONPATH.
-        pack_actions_lib_paths = os.path.join(
-            pack_base_path, "actions", ACTION_LIBS_DIR
-        )
-        pack_virtualenv_lib_path = os.path.join(virtualenv_path, "lib")
-        pack_venv_lib_directory = os.path.join(
-            pack_virtualenv_lib_path, virtualenv_directories[0]
-        )
+        pack_venv_lib_directory = get_virtualenv_lib_path(virtualenv_path)
 
         # Add the pack's site-packages directory (lib/python3.x/site-packages)
         # in front of the Python system site-packages This is important because
         # we want Python 3 compatible libraries to be used from the pack virtual
         # environment and not system ones.
         pack_venv_site_packages_directory = os.path.join(
-            pack_virtualenv_lib_path, virtualenv_directories[0], "site-packages"
+            pack_venv_lib_directory, "site-packages"
+        )
+
+        # Then add the actions/lib directory in the pack.
+        pack_actions_lib_paths = os.path.join(
+            pack_base_path, "actions", ACTION_LIBS_DIR
         )
 
         full_sandbox_python_path = [
@@ -185,7 +200,7 @@ def get_sandbox_python_path_for_python_action(
     return sandbox_python_path
 
 
-def get_sandbox_virtualenv_path(pack):
+def get_sandbox_virtualenv_path(pack: str) -> Optional[str]:
     """
     Return a path to the virtual environment for the provided pack.
     """
@@ -194,7 +209,7 @@ def get_sandbox_virtualenv_path(pack):
         virtualenv_path = None
     else:
         system_base_path = cfg.CONF.system.base_path
-        virtualenv_path = os.path.join(system_base_path, "virtualenvs", pack)
+        virtualenv_path = str(os.path.join(system_base_path, "virtualenvs", pack))
 
     return virtualenv_path
 
diff --git a/st2common/st2common/util/virtualenvs.py b/st2common/st2common/util/virtualenvs.py
@@ -22,6 +22,9 @@
 import os
 import re
 import shutil
+from logging import Logger
+from pathlib import Path
+from textwrap import dedent
 
 import six
 from oslo_config import cfg
@@ -30,14 +33,19 @@
 from st2common.constants.pack import PACK_REF_WHITELIST_REGEX
 from st2common.constants.pack import BASE_PACK_REQUIREMENTS
 from st2common.constants.pack import RESERVED_PACK_LIST
+from st2common.util.sandboxing import (
+    get_site_packages_dir,
+    get_virtualenv_lib_path,
+    is_in_virtualenv,
+)
 from st2common.util.shell import run_command
 from st2common.util.shell import quote_unix
 from st2common.util.compat import to_ascii
 from st2common.util.pack_management import apply_pack_owner_group
 from st2common.content.utils import get_packs_base_paths
 from st2common.content.utils import get_pack_directory
 
-__all__ = ["setup_pack_virtualenv"]
+__all__ = ["setup_pack_virtualenv", "inject_st2_pth_into_virtualenv"]
 
 LOG = logging.getLogger(__name__)
 
@@ -52,6 +60,7 @@ def setup_pack_virtualenv(
     proxy_config=None,
     no_download=True,
     force_owner_group=True,
+    inject_parent_virtualenv_sites=True,
 ):
 
     """
@@ -112,7 +121,15 @@ def setup_pack_virtualenv(
             no_download=no_download,
         )
 
-    # 2. Install base requirements which are common to all the packs
+    # 2. Inject the st2 site-packages dir to enable .pth file loading
+    if inject_parent_virtualenv_sites:
+        logger.debug("Injecting st2 venv site-packages via .pth file")
+        inject_st2_pth_into_virtualenv(
+            virtualenv_path=virtualenv_path,
+            logger=logger,
+        )
+
+    # 3. Install base requirements which are common to all the packs
     logger.debug("Installing base requirements")
     for requirement in BASE_PACK_REQUIREMENTS:
         install_requirement(
@@ -122,7 +139,7 @@ def setup_pack_virtualenv(
             logger=logger,
         )
 
-    # 3. Install pack-specific requirements
+    # 4. Install pack-specific requirements
     requirements_file_path = os.path.join(pack_path, "requirements.txt")
     has_requirements = os.path.isfile(requirements_file_path)
 
@@ -139,7 +156,7 @@ def setup_pack_virtualenv(
     else:
         logger.debug("No pack specific requirements found")
 
-    # 4. Set the owner group
+    # 5. Set the owner group
     if force_owner_group:
         apply_pack_owner_group(pack_path=virtualenv_path)
 
@@ -252,6 +269,55 @@ def remove_virtualenv(virtualenv_path, logger=None):
     return True
 
 
+def inject_st2_pth_into_virtualenv(virtualenv_path: str, logger: Logger = None) -> None:
+    """
+    Add a .pth file to the pack virtualenv that loads any .pth files from the st2 virtualenv.
+
+    If the primary st2 venv (typically /opt/stackstorm/st2) contains any .pth files,
+    the pack's virtualenv would not load them because that directory is on the PYTHONPATH,
+    but it is not a "sites" directory that the "site" module will try to load from.
+    To work around this, we add an .pth file that registers the st2 venv's
+    site-packages directory as a "sites" directory.
+
+    Sites dirs get loaded from (in order): venv, [user site-packages,] system site-packages.
+    After the sites dirs are loaded (including loading their .pth files), sitecustomize gets
+    imported. We do not use sitecustomize because we need the st2 venv's .pth files to load
+    after the pack venv and before system site-packages. So, we use a .pth file that loads
+    as late as possible.
+    """
+    logger = logger or LOG
+    if not is_in_virtualenv():
+        # If this is installed in the system-packages directory, then
+        # its site-packages directory should already be included.
+        logger.debug("Not in a virtualenv; Skipping st2 .pth injection.")
+        return
+
+    parent_site_packages_dir = get_site_packages_dir()
+
+    contents = dedent(
+        # The .pth format requires that any code be on one line.
+        # https://docs.python.org/3/library/site.html
+        # The line gets passed through exec() which uses the scope of site.addpackage()
+        f"""\
+        # This file is auto-generated by StackStorm.
+        import sys; addsitedir('{parent_site_packages_dir}', known_paths)
+        """
+        # TODO: Maybe use this to adjust PATH and PYTHONPATH as well.
+        #       This could replace manipulations in python_runner and sensor process_container:
+        #       env["PATH"] = st2common.util.sandboxing.get_sandbox_path(...)
+        #       env["PYTHONPATH"] = st2common.util.sandboxing.get_sandbox_python_path*(...)
+    )
+
+    # .pth filenames are sorted() before loading, so use many "z"s to ensure
+    # any st2 virtualenv .pth files get loaded after .pth files in the pack's virtualenv.
+    pth_file = (
+        Path(get_virtualenv_lib_path(virtualenv_path))
+        / "site-packages"
+        / "zzzzzzzzzz__st2__.pth"
+    )
+    pth_file.write_text(contents)
+
+
 def install_requirements(
     virtualenv_path, requirements_file_path, proxy_config=None, logger=None
 ):
diff --git a/st2common/tests/unit/test_util_sandboxing.py b/st2common/tests/unit/test_util_sandboxing.py
diff --git a/st2reactor/tests/integration/test_sensor_container.py b/st2reactor/tests/integration/test_sensor_container.py
