mongo/bazel/wrapper_hook/wrapper_hook.py

#!/usr/bin/env python3
import os
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).parent.parent.parent
sys.path.append(str(REPO_ROOT))

# This script should be careful not to disrupt automatic mechanism which
# 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.wrapper_debug import wrapper_debug
from bazel.wrapper_hook.wrapper_util import get_terminal_stream

wrapper_debug(f"wrapper hook script is using {sys.executable}")


def _supports_color(stream):
    if os.name == "nt":
        return False
    if os.environ.get("NO_COLOR"):
        return False
    try:
        return stream.isatty()
    except Exception:
        return False


def _info_prefix(stream):
    if _supports_color(stream):
        GREEN = "\x1b[0;32m"
        RESET = "\x1b[0m"
        return f"{GREEN}INFO{RESET}:"
    return "INFO:"


def _fmt_duration(seconds: float) -> str:
    return f"{seconds*1000:.1f} ms" if seconds < 1 else f"{seconds:.3f} s"


def _info(msg: str, printer=print, stream=None):
    term_err = get_terminal_stream("MONGO_WRAPPER_STDERR_FD")
    # Save current stdout/stderr
    old_stdout = sys.stdout
    old_stderr = sys.stderr

    try:
        sys.stdout = term_err
        sys.stderr = term_err

        stream = stream or sys.stdout
        prefix = _info_prefix(stream)
        printer(f"{prefix} {msg}")

    finally:
        # Restore original stdout/stderr to whatever wrapper has
        sys.stdout = old_stdout
        sys.stderr = old_stderr


def run_with_terminal_output(func, *args, **kwargs):
    term_err = get_terminal_stream("MONGO_WRAPPER_STDERR_FD")
    sys.stdout = None
    # Save current stdout/stderr
    old_stdout = sys.stdout
    old_stderr = sys.stderr

    try:
        sys.stdout = term_err
        sys.stderr = term_err
        return func(*args, **kwargs)

    finally:
        # Restore original stdout/stderr to whatever wrapper has
        sys.stdout = old_stdout
        sys.stderr = old_stderr


def main():
    install_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
    from bazel.wrapper_hook.autogenerated_targets import autogenerate_targets
    from bazel.wrapper_hook.check_resources import check_resource
    from bazel.wrapper_hook.engflow_check import engflow_auth
    from bazel.wrapper_hook.generate_common_bes_bazelrc import write_workstation_bazelrc

    # from bazel.wrapper_hook.git_age_check import check as git_age_check
    from bazel.wrapper_hook.lint import LinterFail
    from bazel.wrapper_hook.plus_interface import check_bazel_command_type, test_runner_interface
    from bazel.wrapper_hook.write_wrapper_hook_bazelrc import write_wrapper_hook_bazelrc

    th_all_header, hdr_state_all_header = spawn_all_headers_thread(REPO_ROOT)

    # Join the header generator before finalizing args.
    start = time.perf_counter()
    auto_hdr_state = gen_auto_headers(REPO_ROOT)
    th_all_header.join()
    if hdr_state_all_header["ok"]:
        wrapper_debug(f'({"wrote" if hdr_state_all_header["wrote"] else "nochange"})')
    else:
        print(f'[all_headers] failed: {hdr_state_all_header["err"]!r}')

    if auto_hdr_state["ok"]:
        wrapper_debug(f'({"wrote" if auto_hdr_state["wrote"] else "nochange"})')
    else:
        print(f'[auto_headers] failed: {auto_hdr_state["err"]!r}')
    t_total_s = time.perf_counter() - start

    _info(f"auto_header build generated: {_fmt_duration(t_total_s)}")

    # 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])

    enterprise = True
    atlas = True
    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():
            enterprise = False
            print(
                f"{enterprise_mod.relative_to(REPO_ROOT).as_posix()} missing, defaulting to local non-enterprise build (--config=local --//bazel/config:build_enterprise=False). Add the directory to not automatically add these options."
            )
            args += ["--//bazel/config:build_enterprise=False", "--config=local"]

        atlas_mod = REPO_ROOT / "src" / "mongo" / "db" / "modules" / "atlas"
        if not atlas_mod.exists():
            atlas = False
            args += ["--//bazel/config:build_atlas=False"]

        if any(arg.startswith("--include_mongot") for arg in args):
            os.makedirs("mongot-localdev", exist_ok=True)

        engflow_auth(args)
        write_workstation_bazelrc(args)
        write_wrapper_hook_bazelrc(args)
        # Disable git age check for now, to avoid issues wth merge commits
        # git_age_check()
        check_resource()

        try:
            args = run_with_terminal_output(
                test_runner_interface,
                args[1:],
                autocomplete_query=os.environ.get("MONGO_AUTOCOMPLETE_QUERY") == "1",
                enterprise=enterprise,
                atlas=atlas,
            )
        except LinterFail:
            # Linter fails preempt bazel run.
            sys.exit(3)
    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))
        f.write("\n")


if __name__ == "__main__":
    main()