Feat: add index_url paramter for setup env

CHANGELOG.rst

@@ -99,6 +99,9 @@ Added
   each parallel test run would be a maintenance burden. #6282
   Contributed by @cognifloyd
 
+* Add new option `index_url` to allow user manually configure python package download url.
+  Contributed by @yhkl-dev
+
 3.8.1 - December 13, 2023
 -------------------------
 Fixed

contrib/packs/actions/pack_mgmt/setup_virtualenv.py

@@ -80,13 +80,16 @@ def __init__(self, config=None, action_service=None):
         ):
             os.environ["no_proxy"] = self.no_proxy
 
-    def run(self, packs, update=False, no_download=True):
+    def run(self, packs, index_url, update=False, no_download=True):
         """
         :param packs: A list of packs to create the environment for.
         :type: packs: ``list``
 
         :param update: True to update dependencies inside the virtual environment.
         :type update: ``bool``
+
+        :param index_url: Package index options.
+        :type index_url: ``str``
         """
 
         for pack_name in packs:
@@ -96,6 +99,7 @@ def run(self, packs, update=False, no_download=True):
                 logger=self.logger,
                 proxy_config=self.proxy_config,
                 no_download=no_download,
+                index_url=index_url,
             )
 
         message = "Successfully set up virtualenv for the following packs: %s" % (

contrib/packs/actions/setup_virtualenv.yaml

@@ -29,3 +29,9 @@
       required: false
       description: Action timeout in seconds. Action will get killed if it doesn't finish in timeout
       type: integer
+    index_url:
+      default: "https://pypi.org/simple"
+      required: false
+      description: Package Index options
+      type: string
+

st2common/st2common/cmd/setup_pack_virtualenv.py

@@ -47,6 +47,11 @@ def _register_cli_opts():
                 "exist, it will create it.."
             ),
         ),
+        cfg.StrOpt(
+            "index_url",
+            default=False,
+            help="Package Index options",
+        ),
     ]
     do_register_cli_opts(cli_opts)
 
@@ -64,6 +69,7 @@ def main(argv):
 
     packs = cfg.CONF.pack
     update = cfg.CONF.update
+    index_url = cfg.CONF.index_url
 
     proxy_config = get_and_set_proxy_config()
 
@@ -76,6 +82,7 @@ def main(argv):
             logger=LOG,
             proxy_config=proxy_config,
             no_download=True,
+            index_url=index_url,
         )
         LOG.info('Successfully set up virtualenv for pack "%s"' % (pack))
 

st2common/st2common/util/virtualenvs.py

@@ -61,6 +61,7 @@ def setup_pack_virtualenv(
     no_download=True,
     force_owner_group=True,
     inject_parent_virtualenv_sites=True,
+    index_url=None,
 ):
 
     """
@@ -78,6 +79,9 @@ def setup_pack_virtualenv(
     :param no_download: Do not download and install latest version of pre-installed packages such
                         as pip and setuptools.
     :type no_download: ``bool``
+
+    :param index_url: Package index options.
+    :type index_url: ``str``
     """
     logger = logger or LOG
 
@@ -137,6 +141,7 @@ def setup_pack_virtualenv(
             requirement=requirement,
             proxy_config=proxy_config,
             logger=logger,
+            index_url=index_url,
         )
 
     # 4. Install pack-specific requirements
@@ -152,6 +157,7 @@ def setup_pack_virtualenv(
             requirements_file_path=requirements_file_path,
             proxy_config=proxy_config,
             logger=logger,
+            index_url=index_url,
         )
     else:
         logger.debug("No pack specific requirements found")
@@ -319,7 +325,11 @@ def inject_st2_pth_into_virtualenv(virtualenv_path: str, logger: Logger = None)
 
 
 def install_requirements(
-    virtualenv_path, requirements_file_path, proxy_config=None, logger=None
+    virtualenv_path,
+    requirements_file_path,
+    proxy_config=None,
+    logger=None,
+    index_url=None,
 ):
     """
     Install requirements from a file.
@@ -346,6 +356,8 @@ def install_requirements(
     cmd.append("install")
     cmd.extend(pip_opts)
     cmd.extend(["-U", "-r", requirements_file_path])
+    if index_url:
+        cmd.extend(["-i", index_url])
 
     env = get_env_for_subprocess_command()
 
@@ -372,7 +384,13 @@ def install_requirements(
     return True
 
 
-def install_requirement(virtualenv_path, requirement, proxy_config=None, logger=None):
+def install_requirement(
+    virtualenv_path,
+    requirement,
+    proxy_config=None,
+    logger=None,
+    index_url=None,
+):
     """
     Install a single requirement.
 
@@ -400,6 +418,9 @@ def install_requirement(virtualenv_path, requirement, proxy_config=None, logger=
     cmd.append("install")
     cmd.extend(pip_opts)
     cmd.extend([requirement])
+    if index_url:
+        cmd.extend(["-i", index_url])
+
     env = get_env_for_subprocess_command()
     logger.debug(
         "Installing requirement %s with command %s.", requirement, " ".join(cmd)

st2common/tests/unit/test_virtualenvs.py

@@ -180,6 +180,31 @@ def test_install_requirement_without_proxy(self):
         }
         virtualenvs.run_command.assert_called_once_with(**expected_args)
 
+    @mock.patch.object(
+        virtualenvs, "run_command", mock.MagicMock(return_value=(0, "", ""))
+    )
+    @mock.patch.object(
+        virtualenvs, "get_env_for_subprocess_command", mock.MagicMock(return_value={})
+    )
+    def test_install_requirement_with_index_url(self):
+        pack_virtualenv_dir = "/opt/stackstorm/virtualenvs/dummy_pack_tests/"
+        requirement = "six>=1.9.0"
+        index_url = "https://test.com/sample"
+        install_requirement(
+            pack_virtualenv_dir, requirement, proxy_config=None, index_url=index_url
+        )
+        expected_args = {
+            "cmd": [
+                "/opt/stackstorm/virtualenvs/dummy_pack_tests/bin/pip",
+                "install",
+                "six>=1.9.0",
+                "-i",
+                index_url,
+            ],
+            "env": {},
+        }
+        virtualenvs.run_command.assert_called_once_with(**expected_args)
+
     @mock.patch.object(
         virtualenvs, "run_command", mock.MagicMock(return_value=(0, "", ""))
     )