SERVER-124059 fix nested bazel wrapper python installs (#51811)

GitOrigin-RevId: 6c8022df8625f61340036f7db6b4e1fbbb9ae7e1

bazel/wrapper_hook/BUILD.bazel

@@ -51,3 +51,14 @@ py_test(
         ":wrapper_hook",
     ],
 )
+
+py_test(
+    name = "install_modules_test",
+    srcs = [
+        "install_modules_test.py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":wrapper_hook",
+    ],
+)

bazel/wrapper_hook/install_modules.py

@@ -12,6 +12,8 @@ sys.path.append(str(REPO_ROOT))
 
 from bazel.wrapper_hook.wrapper_debug import wrapper_debug
 
+MODULES_READY_ENV = "MONGO_BAZEL_WRAPPER_MODULES_READY"
+
 
 def get_deps_dirs(deps):
     tmp_dir = pathlib.Path(os.environ["Temp"] if platform.system() == "Windows" else "/tmp")
@@ -100,6 +102,26 @@ def skip_cplusplus_toolchain(args):
     return False
 
 
+def _reexec_current_python(env_var: str = MODULES_READY_ENV) -> None:
+    wrapper_debug("python deps changed; restarting wrapper interpreter")
+    env = os.environ.copy()
+    env[env_var] = "1"
+    os.execve(sys.executable, [sys.executable, *sys.argv], env)
+
+
+def bootstrap_modules(bazel, args):
+    # Nested Bazel installs can refresh the repo-rule python tree under the
+    # running interpreter. Re-exec so later stdlib imports come from the
+    # refreshed tree instead of the potentially stale one this process started
+    # with.
+    if os.environ.get(MODULES_READY_ENV) == "1":
+        setup_python_path()
+        return
+
+    if install_modules(bazel, args):
+        _reexec_current_python()
+
+
 def install_modules(bazel, args):
     need_to_install = False
     pwd_hash = hashlib.md5(str(REPO_ROOT).encode()).hexdigest()
@@ -159,3 +181,4 @@ def install_modules(bazel, args):
         if deps_missing:
             raise Exception(f"Failed to install python deps {deps_missing}")
     setup_python_path()
+    return need_to_install

bazel/wrapper_hook/install_modules_test.py

@@ -0,0 +1,63 @@
+"""Unit tests for bazel.wrapper_hook.install_modules."""
+
+import os
+import unittest
+from unittest import mock
+
+from bazel.wrapper_hook import install_modules
+
+
+class BootstrapModulesTest(unittest.TestCase):
+    def test_second_phase_only_reapplies_python_path(self):
+        with (
+            mock.patch.dict(
+                os.environ,
+                {install_modules.MODULES_READY_ENV: "1"},
+                clear=True,
+            ),
+            mock.patch.object(install_modules, "setup_python_path") as mock_setup,
+            mock.patch.object(install_modules, "install_modules") as mock_install,
+            mock.patch.object(install_modules.os, "execve") as mock_execve,
+        ):
+            install_modules.bootstrap_modules("bazel", ["bazel", "build", "//:lint"])
+
+        mock_setup.assert_called_once_with()
+        mock_install.assert_not_called()
+        mock_execve.assert_not_called()
+
+    def test_first_phase_returns_without_reexec_when_no_install_needed(self):
+        with (
+            mock.patch.dict(os.environ, {}, clear=True),
+            mock.patch.object(
+                install_modules, "install_modules", return_value=False
+            ) as mock_install,
+            mock.patch.object(install_modules.os, "execve") as mock_execve,
+        ):
+            install_modules.bootstrap_modules("bazel", ["bazel", "build", "//:lint"])
+
+        mock_install.assert_called_once_with("bazel", ["bazel", "build", "//:lint"])
+        mock_execve.assert_not_called()
+
+    def test_first_phase_reexecs_after_nested_install(self):
+        argv = ["wrapper_hook.py", "bazel", "build", "//:lint"]
+        with (
+            mock.patch.dict(os.environ, {}, clear=True),
+            mock.patch.object(
+                install_modules, "install_modules", return_value=True
+            ) as mock_install,
+            mock.patch.object(install_modules.os, "execve") as mock_execve,
+            mock.patch.object(install_modules.sys, "executable", "/tmp/repo-python"),
+            mock.patch.object(install_modules.sys, "argv", argv),
+        ):
+            install_modules.bootstrap_modules("bazel", argv[1:])
+
+        mock_install.assert_called_once_with("bazel", argv[1:])
+        mock_execve.assert_called_once()
+        exec_path, exec_argv, exec_env = mock_execve.call_args.args
+        self.assertEqual(exec_path, "/tmp/repo-python")
+        self.assertEqual(exec_argv, ["/tmp/repo-python", *argv])
+        self.assertEqual(exec_env[install_modules.MODULES_READY_ENV], "1")
+
+
+if __name__ == "__main__":
+    unittest.main()

bazel/wrapper_hook/post_bazel_hook.py

@@ -6,14 +6,14 @@ import sys
 REPO_ROOT = pathlib.Path(__file__).parent.parent.parent
 sys.path.append(str(REPO_ROOT))
 
-from bazel.wrapper_hook.install_modules import install_modules
+from bazel.wrapper_hook.install_modules import bootstrap_modules
 
 BAZEL_USER_NAMESPACE = "user-prod"
 BAZEL_CI_NAMESPACE = "ci-prod"
 
 
 def main():
-    install_modules(sys.argv[1], sys.argv[1:])
+    bootstrap_modules(sys.argv[1], sys.argv[1:])
 
     from bazel.wrapper_hook.compiledb import finalize_compiledb_posthook
     from bazel.wrapper_hook.flag_sync import sync_flags

bazel/wrapper_hook/wrapper_hook.py

@@ -12,7 +12,7 @@ sys.path.append(str(REPO_ROOT))
 # may be expecting certain stdout, always print to stderr.
 sys.stdout = sys.stderr
 
-from bazel.wrapper_hook.install_modules import install_modules
+from bazel.wrapper_hook.install_modules import bootstrap_modules
 from bazel.wrapper_hook.wrapper_debug import wrapper_debug
 from bazel.wrapper_hook.wrapper_util import get_terminal_stream
 
@@ -81,7 +81,7 @@ def run_with_terminal_output(func, *args, **kwargs):
 
 
 def main():
-    install_modules(sys.argv[1], sys.argv[1:])
+    bootstrap_modules(sys.argv[1], sys.argv[1:])
 
     from bazel.auto_header.auto_header import gen_auto_headers
     from bazel.auto_header.gen_all_headers import spawn_all_headers_thread