SERVER-103495 skip wrapper script for certain commands (#34898)

GitOrigin-RevId: 01352faa03ee15baaf72ae3783bce3c29fb06a80

bazel/wrapper_hook/bazel_commands.commands

@@ -0,0 +1,21 @@
+aquery
+build
+canonicalize-flags
+clean
+coverage
+cquery
+dump
+fetch
+help
+info
+license
+mobile-install
+mod
+print_action
+query
+run
+shutdown
+sync
+test
+vendor
+version

bazel/wrapper_hook/plus_interface.py

@@ -51,6 +51,15 @@ def get_buildozer_output(autocomplete_query):
     return p.stdout
 
 
+def check_bazel_command_type(args):
+    with open(pathlib.Path(__file__).parent / "bazel_commands.commands") as f:
+        bazel_commands = [line.strip() for line in f.readlines()]
+
+    for arg in args:
+        if arg in bazel_commands:
+            return arg
+
+
 def test_runner_interface(args, autocomplete_query, get_buildozer_output=get_buildozer_output):
     start = time.time()
 
@@ -72,30 +81,7 @@ def test_runner_interface(args, autocomplete_query, get_buildozer_output=get_bui
     bin_targets = []
     source_targets = {}
 
-    bazel_commands = [
-        "aquery",
-        "build",
-        "canonicalize-flags",
-        "clean",
-        "coverage",
-        "cquery",
-        "dump",
-        "fetch",
-        "help",
-        "info",
-        "license",
-        "mobile-install",
-        "mod",
-        "print_action",
-        "query",
-        "run",
-        "shutdown",
-        "sync",
-        "test",
-        "vendor",
-        "version",
-    ]
-    current_bazel_command = None
+    current_bazel_command = check_bazel_command_type(args)
 
     if autocomplete_query:
         str_args = " ".join(args)
@@ -106,8 +92,6 @@ def test_runner_interface(args, autocomplete_query, get_buildozer_output=get_bui
         persistent_compdb = False
 
     for arg in args:
-        if not current_bazel_command and arg in bazel_commands:
-            current_bazel_command = arg
         if arg in compiledb_targets:
             compiledb_target = True
         if arg in lint_targets:

bazel/wrapper_hook/wrapper_hook.py

@@ -21,19 +21,32 @@ def main():
     from bazel.wrapper_hook.autogenerated_targets import autogenerate_targets
     from bazel.wrapper_hook.developer_bes_keywords import write_workstation_bazelrc
     from bazel.wrapper_hook.engflow_check import engflow_auth
-    from bazel.wrapper_hook.plus_interface import test_runner_interface
-
-    engflow_auth(sys.argv)
-
-    write_workstation_bazelrc(sys.argv)
+    from bazel.wrapper_hook.plus_interface import check_bazel_command_type, test_runner_interface
 
     # This is used to autogenerate a BUILD.bazel that creates
     # Filegroups for select tags - used to group targets for installing
     autogenerate_targets(sys.argv, sys.argv[1])
 
+    if check_bazel_command_type(sys.argv[1:]) not in ["clean", "shutdown", "version", None]:
+        args = sys.argv
+        enterprise_mod = REPO_ROOT / "src" / "mongo" / "db" / "modules" / "enterprise"
+        if not enterprise_mod.exists():
+            print(
+                f"{enterprise_mod.relative_to(REPO_ROOT).as_posix()} missing, defaulting to local enterprise build. Create an empty directory to bypass this."
+            )
+            args += ["--build_enterprise=False", "--config=local"]
+
+        engflow_auth(args)
+
+        write_workstation_bazelrc(args)
+
         args = test_runner_interface(
             sys.argv[1:], autocomplete_query=os.environ.get("MONGO_AUTOCOMPLETE_QUERY") == "1"
         )
+
+    else:
+        args = sys.argv[2:]
+
     os.chmod(os.environ.get("MONGO_BAZEL_WRAPPER_ARGS"), 0o644)
     with open(os.environ.get("MONGO_BAZEL_WRAPPER_ARGS"), "w") as f:
         f.write("\n".join(args))

buildscripts/bazel_rules_mongo/engflow_auth/engflow_auth.py

@@ -8,6 +8,7 @@ import os
 import platform
 import stat
 import subprocess
+import sys
 import time
 import urllib.request
 from datetime import datetime
@@ -115,6 +116,22 @@ def authenticate(binary_path: str, verbose: bool) -> bool:
             print("Already authenticated. Skipping authentication.")
         return True
 
+    try:
+        print(
+            "Attempting to authenticate with MongoDB remote build service. On any device, login via the following link to complete EngFlow authentication:\nGenerating link url and opening browser in 10 seconds:\n(use CTRL+C or COMMAND+C to skip if not an internal mongodb user)"
+        )
+        countdown = 10
+        for i in reversed(range(countdown + 1)):
+            if i == 9:
+                sys.stdout.write("\b\b  \b\b")
+            elif i < 9:
+                sys.stdout.write("\b \b")
+            if i == 0:
+                break
+            sys.stdout.write(str(i))
+            sys.stdout.flush()
+            time.sleep(1)
+
         p = subprocess.Popen(
             f"{binary_path} login -store=file {CLUSTER}",
             shell=True,
@@ -135,14 +152,13 @@ def authenticate(binary_path: str, verbose: bool) -> bool:
             print("CLI had unexpected output.")
             p.kill()
             return False
-
-    print(
-        f"On any device with a browser, login via the following link to complete EngFlow authentication:\n{login_url}"
-    )
+        else:
+            print(login_url)
 
         try:
             p.wait(timeout=LOGIN_TIMEOUT)
             print("Sucessfully authenticated with EngFlow!")
+
         except subprocess.TimeoutExpired:
             print(
                 "Timed out waiting for login attempt. Failed to authenticate with EngFlow. Builds will be run locally..."
@@ -150,6 +166,13 @@ def authenticate(binary_path: str, verbose: bool) -> bool:
             p.kill()
             return False
 
+    except KeyboardInterrupt:
+        print(
+            "Skipping authentication, use '--config=local' to skip trying to authenticate in the future."
+        )
+        time.sleep(3)
+        return False
+
     return True
 
 

tools/bazel

@@ -109,6 +109,46 @@ for cert in ${cert_locs[@]}; do
   fi
 done
 
+skip_python=0
+bazel_commands=()
+while IFS= read -r line; do
+  bazel_commands+=("$line")
+done < $REPO_ROOT/bazel/wrapper_hook/bazel_commands.commands
+
+current_bazel_command=""
+for arg do
+  for command in "${bazel_commands[@]}"
+  do
+    if [ "$command" == "$arg" ] ; then
+      current_bazel_command="$arg"
+      break
+    fi
+  done
+  if [ ! -z $current_bazel_command ]; then
+    break
+  fi
+done
+
+if [ -z $current_bazel_command ]; then
+  skip_python=1
+fi
+
+if [ "$skip_python" == "0" ]; then
+  skip_commands=("clean" "version" "shutdown")
+
+  for command in "${skip_commands[@]}"
+  do
+    if [ "$current_bazel_command" == "$command" ] ; then
+      skip_python=1
+      break
+    fi
+  done
+fi
+
+if [ "$skip_python" == "1" ]; then
+  exec "$bazel_real" $@
+fi
+
 cur_dir=$(basename $REPO_ROOT)
 bazel_python="$REPO_ROOT/bazel-$cur_dir/external/_main~setup_mongo_python_toolchains~py_${os}_${ARCH}/dist/bin/python3"
 compdb_python="$REPO_ROOT/.compiledb/compiledb-$cur_dir/external/py_${os}_${ARCH}/dist/bin/python3"

tools/bazel.bat

@@ -7,6 +7,40 @@ set REPO_ROOT=%~dp0..
 
 for %%I in (%REPO_ROOT%) do set cur_dir=%%~nxI
 
+set skip_python="0"
+set bazelCommandCount=0
+for /F "delims=" %%a in (%REPO_ROOT%\bazel\wrapper_hook\bazel_commands.commands) do (
+    set /A bazelCommandCount+=1
+    set "bazel_commands[!bazelCommandCount!]=%%~a"
+)
+
+set argCount=0
+for %%x in (%*) do (
+   set /A argCount+=1
+   set "argVec[!argCount!]=%%~x"
+)
+
+set current_bazel_command=""
+for /L %%i in (1,1,%argCount%) do (
+  for /L %%j in (1,1,%bazelCommandCount%) do (
+    if "!bazel_commands[%%j]!"=="!argVec[%%i]!" (
+      set current_bazel_command="!argVec[%%i]!"
+      goto :found_command
+    )
+  )
+)
+:found_command
+if !current_bazel_command!=="" set skip_python="1"
+
+if !skip_python!=="0" if !current_bazel_command!=="clean" set skip_python="1"
+if !skip_python!=="0" if !current_bazel_command!=="version" set skip_python="1"
+if !skip_python!=="0" if !current_bazel_command!=="shutdown" set skip_python="1"
+
+if !skip_python!=="1" (
+    "%BAZEL_REAL%" %*
+    exit %ERRORLEVEL%
+)
+
 set bazel_python="%REPO_ROOT%\bazel-%cur_dir%\external\_main~setup_mongo_python_toolchains~py_windows_x86_64\dist\python.exe"
 set compdb_python="%REPO_ROOT%\.compiledb\compiledb-%cur_dir%\external\py_windows_x86_64\dist\python.exe"
 set python=%bazel_python%