scripts/portage_cmd_wrapper.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.

"""Wrapper for board specific portage commands.

This script is meant to be used in generated wrapper scripts, not used directly.
"""

import logging
import os
from pathlib import Path
from typing import Iterable, List, Optional

from chromite.third_party.opentelemetry import trace

from chromite.lib import build_query
from chromite.lib import chromite_config
from chromite.lib import commandline
from chromite.lib import constants
from chromite.lib import cros_build_lib
from chromite.lib import osutils
from chromite.lib import portage_util
from chromite.lib import sysroot_lib
from chromite.lib.parser import package_info
from chromite.utils import telemetry


tracer = trace.get_tracer(__name__)


def get_parser() -> commandline.ArgumentParser:
    """Build the argument parser."""
    parser = commandline.ArgumentParser(description=__doc__)

    parser.add_argument(
        "--build-target",
        required=True,
        help="The build target name.",
    )
    parser.add_argument(
        "--sysroot",
        type="path",
        required=True,
        help="The path to the sysroot for which the command will be created.",
    )
    parser.add_argument(
        "--chost",
        required=True,
        help="The CHOST value for the sysroot.",
    )
    parser.add_argument(
        "command",
        nargs="+",
        help="The command to run.",
    )

    return parser


def parse_arguments(argv: List[str]) -> commandline.ArgumentNamespace:
    """Parse and validate arguments."""
    parser = get_parser()
    opts = parser.parse_args(argv)

    opts.Freeze()
    return opts


@tracer.start_as_current_span("portage_cmd_wrapper.parse_pkgs")
def parse_pkgs(command: List[str], build_target_name: str) -> Iterable[str]:
    """Parse packages from a command."""
    span = trace.get_current_span()
    span.update_name(f"portage_cmd_wrapper.{command[0]}.parse_pkgs")

    pkg_fragments = set()
    for arg in command[1:]:
        if arg.startswith("-"):
            # Skip --arguments.
            continue

        try:
            pkg = package_info.parse(arg)
        except ValueError:
            # e.g. /some/path.
            continue

        if pkg.cpvr or pkg.atom:
            # We have at least an atom, that's good enough.
            yield pkg.cpvr or pkg.atom
        else:
            # The fragment gets parsed as the package name.
            pkg_fragments.add(pkg.package)

    if pkg_fragments:
        ebuilds = build_query.Query(build_query.Ebuild, board=build_target_name)
        for ebuild in ebuilds:
            if ebuild.package_info.package in pkg_fragments:
                yield ebuild.package_info.cpvr


# TODO: Find a better name and a reusable location for this.
@tracer.start_as_current_span("portage_cmd_wrapper.sudo_run_cmd")
def sudo_run_cmd_with_failed_pkg_parsing(command, extra_env):
    """Wrapper for sudo_run that adds CROS_METRICS_DIR usage."""
    span = trace.get_current_span()
    span.update_name(f"portage_cmd_wrapper.{command[0]}.sudo_run_cmd")

    extra_env = extra_env.copy()
    with osutils.TempDir() as tempdir:
        extra_env[constants.CROS_METRICS_DIR_ENVVAR] = tempdir
        try:
            return cros_build_lib.sudo_run(
                command,
                preserve_env=True,
                extra_env=extra_env,
            )
        except cros_build_lib.RunCommandError as e:
            raise sysroot_lib.PackageInstallError(
                "Merging board packages failed",
                e.result,
                exception=e,
                packages=portage_util.ParseDieHookStatusFile(tempdir),
            )


@tracer.start_as_current_span("portage_cmd_wrapper.execute_cmd")
def execute_cmd(opts: commandline.ArgumentNamespace) -> int:
    extra_env = {
        "CHOST": opts.chost,
        "PORTAGE_CONFIGROOT": opts.sysroot,
        "SYSROOT": opts.sysroot,
        "ROOT": opts.sysroot,
        "PORTAGE_USERNAME": (
            os.environ.get("PORTAGE_USERNAME") or Path("~").expanduser().name
        ),
    }

    # If we try to use sudo when the sandbox is active, we get ugly warnings
    # that just confuse developers.
    if os.environ.get("SANDBOX_ON") == "1":
        os.environ["SANDBOX_ON"] = "0"
    os.environ.pop("LD_PRELOAD", None)

    pkgs = []
    if opts.command[0] == "emerge":
        pkgs = list(parse_pkgs(opts.command, opts.build_target))

    span = trace.get_current_span()
    span.update_name(f"portage_cmd_wrapper.{opts.command[0]}.execute_cmd")
    span.set_attributes(
        {
            "executable": opts.command[0],
            "command": opts.command,
            "build_target": opts.build_target,
            "extra_env": [f"{k}={v}" for k, v in extra_env.items()],
            "packages": pkgs,
        }
    )

    return sudo_run_cmd_with_failed_pkg_parsing(
        opts.command, extra_env
    ).returncode


@tracer.start_as_current_span("portage_cmd_wrapper.main")
def main(argv: Optional[List[str]]) -> Optional[int]:
    """Main."""
    commandline.RunInsideChroot()

    opts = parse_arguments(argv)

    if opts.command[0] != "equery":
        # There's a *lot* more equery calls than any other command. Specifically
        # lots of parallel executions in cros clean-outdated-packages.
        # Nothing wrong with those usages, but it's pretty noisy for our data
        # given it's not one we're currently concerned about. So for now, just
        # skip all equery invocations.
        chromite_config.initialize()
        telemetry.initialize(chromite_config.TELEMETRY_CONFIG, debug=opts.debug)

    span = trace.get_current_span()
    span.update_name(f"portage_cmd_wrapper.{opts.command[0]}.main")

    try:
        return execute_cmd(opts)
    except cros_build_lib.RunCommandError as e:
        logging.error(e)
        logging.error("Error running %s.", opts.command[0])
        logging.error(
            "Full command: %s", cros_build_lib.CmdToStr(e.result.args)
        )
        # sysroot_lib.PackageInstallError is a subclass of RunCommandError.
        if hasattr(e, "failed_packages"):
            logging.error("Failed Packages: %s", " ".join(e.failed_packages))
        return e.result.returncode