mongo/bazel/utils.bzl

# General starlark utility functions
load("//bazel/platforms:normalize.bzl", "ARCH_NORMALIZE_MAP")

def write_target_impl(ctx):
    out = ctx.actions.declare_file(ctx.label.name + ".gen_source_list")
    ctx.actions.write(
        out,
        "//" + ctx.label.package + ":" + ctx.attr.target_name,
    )
    return [
        DefaultInfo(
            files = depset([out]),
        ),
    ]

write_target = rule(
    write_target_impl,
    attrs = {
        "target_name": attr.string(
            doc = "the name of the target to record",
        ),
    },
)

def retry_download_and_extract(ctx, tries, **kwargs):
    sleep_time = 1
    for attempt in range(tries):
        is_retriable = attempt + 1 < tries
        result = ctx.download_and_extract(allow_fail = is_retriable, **kwargs)
        if result.success:
            return result
        else:
            print("Download failed (Attempt #%s), sleeping for %s seconds then retrying..." % (attempt + 1, sleep_time))
            ctx.execute(["sleep", str(sleep_time)])
            sleep_time *= 2

def retry_download(ctx, tries, **kwargs):
    sleep_time = 1
    for attempt in range(tries):
        is_retriable = attempt + 1 < tries
        result = ctx.download(allow_fail = is_retriable, **kwargs)
        if result.success:
            return result
        else:
            print("Download failed (Attempt #%s), sleeping for %s seconds then retrying..." % (attempt + 1, sleep_time))
            ctx.execute(["sleep", str(sleep_time)])
            sleep_time *= 2

def write_python_pyc_cache_prefix_customization(ctx, customization_file, pycache_dirname = "bazel_pycache"):
    """Write a site/usercustomize module to redirect .pyc writes to /tmp.

    Toolchain and runtime Python distributions often live under Bazel-managed paths
    (external repositories, runfiles, etc.). If Python writes `__pycache__/*.pyc`
    there, it can cause non-hermetic filesystem changes and Bazel cache churn.

    This helper writes either `sitecustomize.py` or `usercustomize.py` (callers
    choose the filename) to set `sys.pycache_prefix` to a temp directory.
    """
    ctx.file(
        customization_file,
        """
import os
import sys
import tempfile

# Prevent bytecode cache writes under Bazel-managed paths (external/, runfiles/).
sys.pycache_prefix = os.path.join(tempfile.gettempdir(), "{pycache_dirname}")
""".format(pycache_dirname = pycache_dirname),
    )

def generate_noop_toolchain(ctx, substitutions):
    # BUILD file is required for a no-op.
    # Keep a stub mongo_toolchain target so unconditional register_toolchains()
    # calls don't fail when the toolchain is intentionally skipped/unsupported.
    # Create a stub clang-format script so that targets referencing
    # //:clang_format (e.g. the format_multirun rule) can still build on
    # platforms where the mongo toolchain is unavailable (macOS, etc.).
    ctx.file(
        "clang_format_noop.sh",
        "#!/usr/bin/env bash\n# Stub: mongo toolchain clang-format is not available on this platform.\nexit 0\n",
        executable = True,
    )

    ctx.file(
        "BUILD.bazel",
        """
# {} not supported on this platform

package(default_visibility = ["//visibility:public"])

filegroup(
    name = "all_files",
    srcs = [],
)

filegroup(
    name = "clang_tidy",
    srcs = [],
)

sh_binary(
    name = "clang_format",
    srcs = ["clang_format_noop.sh"],
)

filegroup(
    name = "llvm_symbolizer",
    srcs = [],
)

filegroup(
    name = "llvm_symbolizer_libs",
    srcs = [],
)

toolchain(
    name = "mongo_toolchain",
    exec_compatible_with = ["@platforms//:incompatible"],
    target_compatible_with = ["@platforms//:incompatible"],
    toolchain = "@bazel_tools//tools/cpp:current_cc_toolchain",
    toolchain_type = "@bazel_tools//tools/cpp:toolchain_type",
)
""".format(ctx.attr.version),
    )

def get_toolchain_subs(ctx):
    if ctx.attr.os:
        os = ctx.attr.os
    else:
        os = ctx.os.name

    if ctx.attr.arch:
        arch = ctx.attr.arch
    else:
        arch = ctx.os.arch

    arch = ARCH_NORMALIZE_MAP[arch]

    version = ctx.attr.version

    distro = get_host_distro_major_version(ctx)

    if os != "linux":
        substitutions = {
            "{platforms_arch}": "arm64",
            "{bazel_toolchain_cpu}": arch,
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
        generate_noop_toolchain(ctx, substitutions)
        ctx.report_progress("mongo toolchain not supported on " + os + " and " + arch)

    if arch == "aarch64":
        substitutions = {
            "{platforms_arch}": "arm64",
            "{bazel_toolchain_cpu}": arch,
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
    elif arch == "x86_64":
        substitutions = {
            "{platforms_arch}": "x86_64",
            "{bazel_toolchain_cpu}": "x86_64",
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
    elif arch == "ppc64le":
        substitutions = {
            "{platforms_arch}": "ppc64le",
            "{bazel_toolchain_cpu}": "ppc64le",
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
    elif arch == "s390x":
        substitutions = {
            "{platforms_arch}": "s390x",
            "{bazel_toolchain_cpu}": arch,
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
    else:
        substitutions = {
            "{platforms_arch}": "none",
            "{bazel_toolchain_cpu}": arch,
            "{arch}": arch,
            "{version}": version,
            "{distro}": distro,
        }
        generate_noop_toolchain(ctx, substitutions)
        ctx.report_progress("mongo toolchain not supported on " + os + " and " + arch)

    return distro, arch, substitutions

def get_host_distro_major_version(repository_ctx):
    _DISTRO_PATTERN_MAP = {
        "Ubuntu 18*": "ubuntu18",
        "Ubuntu 20*": "ubuntu20",
        "Ubuntu 22*": "ubuntu22",
        "Pop!_OS 22*": "ubuntu22",
        "Ubuntu 24*": "ubuntu24",
        "Amazon Linux 2": "amazon_linux_2",
        "Amazon Linux 2023": "amazon_linux_2023",
        "Debian GNU/Linux 10": "debian10",
        "Debian GNU/Linux 12": "debian12",
        "Red Hat Enterprise Linux 8*": "rhel8",
        "Red Hat Enterprise Linux 9*": "rhel9",
        "Red Hat Enterprise Linux 10*": "rhel10",
        "Fedora*": "rhel10",
        "SLES 15*": "suse15",
    }

    if repository_ctx.os.name != "linux":
        return None

    result = repository_ctx.execute([
        "sed",
        "-n",
        "/^\\(NAME\\|VERSION_ID\\)=/{s/[^=]*=//;s/\"//g;p}",
        "/etc/os-release",
    ])

    if result.return_code != 0:
        print("Failed to determine system distro, parsing os-release failed with the error: " + result.stderr)
        return None

    distro_seq = result.stdout.splitlines()
    if len(distro_seq) != 2:
        print("Failed to determine system distro, parsing os-release returned: " + result.stdout)
        return None

    distro_str = "{distro_name} {distro_version}".format(
        distro_name = distro_seq[0],
        distro_version = distro_seq[1],
    )

    for distro_pattern, simplified_name in _DISTRO_PATTERN_MAP.items():
        if "*" in distro_pattern:
            prefix_suffix = distro_pattern.split("*")
            if distro_str.startswith(prefix_suffix[0]) and distro_str.endswith(prefix_suffix[1]):
                return simplified_name
        elif distro_str == distro_pattern:
            return simplified_name
    return None