mongo_toolchain.py

import os
import subprocess
import sys
from pathlib import Path
from typing import Optional

import typer
from typing_extensions import Annotated

# Get relative imports to work when the package is not installed on the PYTHONPATH.
if __name__ == "__main__" and __package__ is None:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(os.path.realpath(__file__)))))


SUPPORTED_VERSIONS = ("v4", "v5")


class MongoToolchainError(RuntimeError):
    pass


class MongoToolchainNotFoundError(MongoToolchainError):
    def __init__(self, extra_msg: str):
        super().__init__(extra_msg)


class ToolNotFoundError(MongoToolchainError):
    def __init__(self, tool: str, extra_msg: str | None):
        msg = f"Couldn't find {tool} in mongo toolchain"
        msg = msg + f": {extra_msg}" if extra_msg is not None else msg
        super().__init__(msg)


class MongoToolchain:
    def __init__(self, root_dir_path: Path):
        self._root_dir = root_dir_path

    def get_tool_path(self, tool: str) -> str:
        path = self._get_bin_dir_path() / tool
        if not path.exists():
            raise ToolNotFoundError(tool, f"{path} does not exist")
        if not path.is_file() or not os.access(path, os.X_OK):
            raise ToolNotFoundError(tool, f"{path} is not an executable")
        return str(path)

    def check_exists(self) -> None:
        for directory in (
            self._root_dir,
            self._get_bin_dir_path(),
            self._get_include_dir_path(),
            self._get_lib_dir_path(),
        ):
            if not directory.is_dir():
                raise MongoToolchainNotFoundError(f"{directory} is not a directory")

    def get_root_dir(self) -> str:
        return str(self._root_dir)

    def get_bin_dir(self) -> str:
        return str(self._get_bin_dir_path())

    def get_include_dir(self) -> str:
        return str(self._get_include_dir_path())

    def get_lib_dir(self) -> str:
        return str(self._get_include_dir_path())

    def _get_bin_dir_path(self) -> Path:
        return self._root_dir / "bin"

    def _get_include_dir_path(self) -> Path:
        return self._root_dir / "include"

    def _get_lib_dir_path(self) -> Path:
        return self._root_dir / "lib"


def _run_command(cmd: str) -> str:
    return subprocess.check_output(cmd, shell=True, text=True).strip()


def _fetch_bazel_toolchain(version: str) -> None:
    try:
        _run_command(f"bazel build @mongo_toolchain_{version}//:all")
    except subprocess.CalledProcessError as e:
        raise MongoToolchainNotFoundError(
            f"Failed to fetch bazel toolchain: `{e.cmd}` exited with code {e.returncode}"
        )


def _get_bazel_execroot() -> Path:
    try:
        execroot_str = _run_command("bazel info execution_root")
    except subprocess.CalledProcessError as e:
        raise MongoToolchainNotFoundError(
            f"Couldn't find bazel execroot: `{e.cmd}` exited with code {e.returncode}"
        )
    execroot = Path(execroot_str)
    return execroot


def _get_bazel_toolchain_path(version: str) -> Path:
    return _get_bazel_execroot() / "external" / f"mongo_toolchain_{version}" / version


def _get_toolchain_from_path(path: str) -> MongoToolchain:
    toolchain = MongoToolchain(Path(path).resolve())
    toolchain.check_exists()
    return toolchain


def _get_bazel_toolchain(version: str) -> MongoToolchain:
    path = _get_bazel_toolchain_path(version)
    if not path.exists():
        _fetch_bazel_toolchain(version)
    if not path.is_dir():
        raise MongoToolchainNotFoundError(
            f"Couldn't find bazel toolchain: {path} is not a directory"
        )
    return _get_toolchain_from_path(path)


def _get_installed_toolchain_path(version: str) -> Path:
    return Path("/opt/mongodbtoolchain") / version


def _get_installed_toolchain(version: str):
    return _get_toolchain_from_path(_get_installed_toolchain_path(version))


def get_mongo_toolchain(
    *, version: str | None = None, from_bazel: bool | None = None
) -> MongoToolchain:
    # When running under bazel this environment variable will be set and will point to the
    # toolchain the target was configured to use. It can also be set manually to override
    # a script's selection of toolchain.
    toolchain_path = os.environ.get("MONGO_TOOLCHAIN_PATH", None)
    if toolchain_path is not None:
        return _get_toolchain_from_path(toolchain_path)

    # If no version given, look in the environment or default to v4.
    if version is None:
        version = os.environ.get("MONGO_TOOLCHAIN_VERSION", "v4")
    assert version is not None
    if version not in SUPPORTED_VERSIONS:
        raise MongoToolchainNotFoundError(f"Unknown toolchain version {version}")

    # If from_bazel is unspecified. Look in the environment or fall back to a default based on
    # the version (v4 -> not from bazel, v5 -> from bazel).
    def _parse_from_bazel_envvar(value: str) -> bool:
        v = value.lower()
        if v in ("true", "1"):
            return True
        elif v in ("false", "0"):
            return False
        else:
            raise ValueError(f"Invalid value {value} for MONGO_TOOLCHAIN_FROM_BAZEL")

    if from_bazel is None:
        from_bazel_default = "false" if version == "v4" else "true"
        from_bazel_value = os.environ.get("MONGO_TOOLCHAIN_FROM_BAZEL", from_bazel_default)
        from_bazel = _parse_from_bazel_envvar(from_bazel_value)

    if from_bazel:
        return _get_bazel_toolchain(version)
    return _get_installed_toolchain(version)


def try_get_mongo_toolchain(*args, **kwargs) -> MongoToolchain | None:
    try:
        return get_mongo_toolchain(*args, **kwargs)
    except MongoToolchainError:
        return None


_app = typer.Typer(add_completion=False)


@_app.command()
def main(
    tool: Annotated[Optional[str], typer.Argument()] = None,
    version: Annotated[Optional[str], typer.Option("--version")] = None,
    from_bazel: Annotated[Optional[bool], typer.Option("--bazel/--no-bazel")] = None,
):
    """
    Prints the path to tools in the mongo toolchain or the toolchain's root directory (which
    should contain bin/, include/, and so on).
    If MONGO_TOOLCHAIN_PATH is set in the environment, it overrides all options to this script.
    Otherwise, MONGO_TOOLCHAIN_VERSION is a lower-precedence alternative to --version and
    MONGO_TOOLCHAIN_FROM_BAZEL is a lower-precedence alternative to --bazel/--no-bazel.
    None of these are required, and if none are given, the default configuration will be used.
    """
    toolchain = get_mongo_toolchain(version=version, from_bazel=from_bazel)
    if tool is not None:
        print(toolchain.get_tool_path(tool))
    else:
        print(toolchain.get_root_dir())


if __name__ == "__main__":
    _app()