scripts/bazel.py

# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Bazel command wrapper.

This wrapper sets up necessary symlinks for the workspace, ensures Bazelisk via
CIPD, and executes it.  It's also the right home for gathering any telemetry on
users' Bazel commands.
"""

import argparse
import logging
import os
from pathlib import Path
from typing import List, Optional, Tuple

from chromite.lib import cipd
from chromite.lib import constants
from chromite.lib import osutils


# TODO(jrosenth): We likely want to publish our own Bazelisk at some point
# instead of relying upon Skia's.
_BAZELISK_PACKAGE = "skia/bots/bazelisk_${os}_${arch}"
_BAZELISK_VERSION = "version:0"

# Symlinks which may exist in the workspace root without an underlying file in
# src/bazel/workspace_root.  These are symlinks generated by bazel itself.
_KNOWN_SYMLINKS = [
    "bazel-bin",
    "bazel-out",
    "bazel-src",
    "bazel-testlogs",
    "@portage",
]

# Workspaces for each project are defined here.
_PROJECTS = ["alchemy", "metallurgy", "fwsdk"]
_WORKSPACES_DIR = constants.BAZEL_WORKSPACE_ROOT / "bazel" / "workspace_root"


def _setup_workspace(project: str) -> None:
    """Setup the Bazel workspace root.

    Args:
        project: The temporary project type (e.g., metallurgy, alchemy, or
            fwsdk).  This argument will eventually be removed when all Bazel
            projects share a unified workspace.
    """
    known_symlinks = set(_KNOWN_SYMLINKS)
    for workspace in (
        _WORKSPACES_DIR / "general",
        _WORKSPACES_DIR / project,
    ):
        for path in workspace.iterdir():
            osutils.SafeSymlink(
                path.relative_to(constants.BAZEL_WORKSPACE_ROOT),
                constants.BAZEL_WORKSPACE_ROOT / path.name,
            )
            known_symlinks.add(path.name)

    # Remove any stale symlinks from the workspace root.
    for path in constants.BAZEL_WORKSPACE_ROOT.iterdir():
        if path.is_symlink() and not path.name in known_symlinks:
            osutils.SafeUnlink(path)


def _get_default_project() -> str:
    """Get the default value for --project.

    It's inconvenient to pass --project for each Bazel invocation.  We assume if
    the user has run with --project before, we can use the value from their last
    invocation.

    If no other default project can be found, the assumed project is "alchemy".

    This function will be removed once all projects unify into a single Bazel
    workspace.
    """
    workspace_file = constants.BAZEL_WORKSPACE_ROOT / "WORKSPACE.bazel"
    if workspace_file.is_symlink():
        project = Path(os.readlink(workspace_file)).parent.name
        if project in _PROJECTS:
            return project
        else:
            logging.warning(
                "Your checkout contains a WORKSPACE.bazel symlink which points "
                "to an unknown project (%s).",
                project,
            )

    logging.notice(
        "Assuming a default project of alchemy.  Pass --project if you want a "
        "different one."
    )
    return "alchemy"


def _get_bazelisk() -> Path:
    """Ensure Bazelisk from CIPD.

    Returns:
        The path to the Bazel executable.
    """
    cipd_path = cipd.GetCIPDFromCache()
    package_path = cipd.InstallPackage(
        cipd_path,
        _BAZELISK_PACKAGE,
        _BAZELISK_VERSION,
        print_cmd=False,
    )
    return package_path / "bazelisk"


def _get_parser() -> argparse.ArgumentParser:
    """Build the argument parser."""

    # Do not use commandline.ArgumentParser.
    # It adds extra flags with the same names as bazel's flags.
    # For example, bazel info --color=yes will actually throw an error:
    # bazel: error: argument --color: ignored explicit argument 'yes'
    parser = argparse.ArgumentParser(add_help=False)

    parser.add_argument(
        "--project",
        choices=_PROJECTS,
        default=_get_default_project(),
        help=(
            "The temporary project type.  This argument will be removed once "
            "all projects unify into a single Bazel workspace."
        ),
    )

    return parser


def parse_arguments(
    argv: Optional[List[str]],
) -> Tuple[argparse.Namespace, List[str]]:
    """Parse and validate arguments.

    Args:
        argv: The command line to parse.

    Returns:
        A two tuple, the parsed arguments, and the remaining arguments that
        should be passed to Bazel.
    """
    parser = _get_parser()
    opts, bazel_args = parser.parse_known_args(argv)
    return opts, bazel_args


def main(argv: Optional[List[str]]) -> Optional[int]:
    """Main."""
    opts, bazel_args = parse_arguments(argv)

    _setup_workspace(opts.project)

    bazelisk = _get_bazelisk()
    os.environ["CHROMITE_BAZEL_WRAPPER"] = "1"
    os.execv(bazelisk, [bazelisk, *bazel_args])