ClickHouse/docker/packager/packager

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import logging
import os
import subprocess
import sys
from pathlib import Path
from typing import List, Optional

SCRIPT_PATH = Path(__file__).absolute()
IMAGE_TYPE = "binary-builder"
IMAGE_NAME = f"clickhouse/{IMAGE_TYPE}"


class BuildException(Exception):
    pass


def check_image_exists_locally(image_name: str) -> bool:
    try:
        output = subprocess.check_output(
            f"docker images -q {image_name} 2> /dev/null", shell=True
        )
        return output != b""
    except subprocess.CalledProcessError:
        return False


def pull_image(image_name: str) -> bool:
    try:
        subprocess.check_call(f"docker pull {image_name}", shell=True)
        return True
    except subprocess.CalledProcessError:
        logging.info("Cannot pull image %s", image_name)
        return False


def build_image(image_name: str, filepath: Path) -> None:
    context = filepath.parent
    build_cmd = f"docker build --network=host -t {image_name} -f {filepath} {context}"
    logging.info("Will build image with cmd: '%s'", build_cmd)
    subprocess.check_call(
        build_cmd,
        shell=True,
    )


def pre_build(repo_path: Path, env_variables: List[str]) -> None:
    if "WITH_PERFORMANCE=1" in env_variables:
        current_branch = subprocess.check_output(
            "git branch --show-current", shell=True, encoding="utf-8"
        ).strip()
        is_shallow = (
            subprocess.check_output(
                "git rev-parse --is-shallow-repository", shell=True, encoding="utf-8"
            )
            == "true\n"
        )
        if is_shallow:
            # I've spent quite some time on looking around the problem, and my
            # conclusion is: in the current state the easiest way to go is to force
            # unshallow repository for performance artifacts.
            # To change it we need to rework our performance tests docker image
            raise BuildException(
                "shallow repository is not suitable for performance builds"
            )
        if current_branch != "master":
            cmd = (
                f"git -C {repo_path} fetch --no-recurse-submodules "
                "--no-tags origin master:master"
            )
            logging.info("Getting master branch for performance artifact: '%s'", cmd)
            subprocess.check_call(cmd, shell=True)


def run_docker_image_with_env(
    image_name: str,
    as_root: bool,
    output_dir: Path,
    env_variables: List[str],
    ch_root: Path,
    cargo_cache_dir: Path,
    ccache_dir: Optional[Path],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    cargo_cache_dir.mkdir(parents=True, exist_ok=True)

    env_part = " -e ".join(env_variables)
    if env_part:
        env_part = " -e " + env_part

    if sys.stdout.isatty():
        interactive = "-it"
    else:
        interactive = ""

    if as_root:
        user = "0:0"
    else:
        user = f"{os.geteuid()}:{os.getegid()}"

    ccache_mount = f"--volume={ccache_dir}:/ccache"
    if ccache_dir is None:
        ccache_mount = ""

    cmd = (
        f"docker run --network=host --user={user} --rm {ccache_mount} "
        f"--volume={output_dir}:/output --volume={ch_root}:/build {env_part} "
        f"--volume={cargo_cache_dir}:/rust/cargo/registry {interactive} {image_name}"
    )

    logging.info("Will build ClickHouse pkg with cmd: '%s'", cmd)

    subprocess.check_call(cmd, shell=True)


def is_release_build(
    debug_build: bool, package_type: str, sanitizer: str, coverage: bool
) -> bool:
    return (
        not debug_build and package_type == "deb" and sanitizer == "" and not coverage
    )


def parse_env_variables(
    debug_build: bool,
    coverage: bool,
    compiler: str,
    sanitizer: str,
    package_type: str,
    cache: str,
    s3_bucket: str,
    s3_directory: str,
    s3_rw_access: bool,
    clang_tidy: bool,
    version: str,
    official: bool,
    additional_pkgs: bool,
    with_profiler: bool,
    with_coverage: bool,
    with_binaries: str,
) -> List[str]:
    DARWIN_SUFFIX = "-darwin"
    DARWIN_ARM_SUFFIX = "-darwin-aarch64"
    ARM_SUFFIX = "-aarch64"
    ARM_V80COMPAT_SUFFIX = "-aarch64-v80compat"
    FREEBSD_SUFFIX = "-freebsd"
    PPC_SUFFIX = "-ppc64le"
    RISCV_SUFFIX = "-riscv64"
    S390X_SUFFIX = "-s390x"
    LOONGARCH_SUFFIX = "-loongarch64"
    AMD64_COMPAT_SUFFIX = "-amd64-compat"
    AMD64_MUSL_SUFFIX = "-amd64-musl"

    result = []
    result.append("OUTPUT_DIR=/output")
    cmake_flags = ["$CMAKE_FLAGS"]
    if package_type == "fuzzers":
        build_target = "fuzzers"
    else:
        build_target = "clickhouse-bundle"

    is_cross_darwin = compiler.endswith(DARWIN_SUFFIX)
    is_cross_darwin_arm = compiler.endswith(DARWIN_ARM_SUFFIX)
    is_cross_arm = compiler.endswith(ARM_SUFFIX)
    is_cross_arm_v80compat = compiler.endswith(ARM_V80COMPAT_SUFFIX)
    is_cross_ppc = compiler.endswith(PPC_SUFFIX)
    is_cross_riscv = compiler.endswith(RISCV_SUFFIX)
    is_cross_s390x = compiler.endswith(S390X_SUFFIX)
    is_cross_loongarch = compiler.endswith(LOONGARCH_SUFFIX)
    is_cross_freebsd = compiler.endswith(FREEBSD_SUFFIX)
    is_amd64_compat = compiler.endswith(AMD64_COMPAT_SUFFIX)
    is_amd64_musl = compiler.endswith(AMD64_MUSL_SUFFIX)

    if is_cross_darwin:
        cc = compiler[: -len(DARWIN_SUFFIX)]
        cmake_flags.append("-DCMAKE_AR:FILEPATH=/cctools/bin/x86_64-apple-darwin-ar")
        cmake_flags.append(
            "-DCMAKE_INSTALL_NAME_TOOL=/cctools/bin/"
            "x86_64-apple-darwin-install_name_tool"
        )
        cmake_flags.append(
            "-DCMAKE_RANLIB:FILEPATH=/cctools/bin/x86_64-apple-darwin-ranlib"
        )
        cmake_flags.append("-DLINKER_NAME=/cctools/bin/x86_64-apple-darwin-ld")
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/darwin/toolchain-x86_64.cmake"
        )
        result.append("EXTRACT_TOOLCHAIN_DARWIN=1")
        result.append("EXPORT_SOURCES_WITH_SUBMODULES=1")
    elif is_cross_darwin_arm:
        cc = compiler[: -len(DARWIN_ARM_SUFFIX)]
        cmake_flags.append("-DCMAKE_AR:FILEPATH=/cctools/bin/aarch64-apple-darwin-ar")
        cmake_flags.append(
            "-DCMAKE_INSTALL_NAME_TOOL=/cctools/bin/"
            "aarch64-apple-darwin-install_name_tool"
        )
        cmake_flags.append(
            "-DCMAKE_RANLIB:FILEPATH=/cctools/bin/aarch64-apple-darwin-ranlib"
        )
        cmake_flags.append("-DLINKER_NAME=/cctools/bin/aarch64-apple-darwin-ld")
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/darwin/toolchain-aarch64.cmake"
        )
        result.append("EXTRACT_TOOLCHAIN_DARWIN=1")
    elif is_cross_arm:
        cc = compiler[: -len(ARM_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-aarch64.cmake"
        )
        result.append("DEB_ARCH=arm64")
    elif is_cross_arm_v80compat:
        cc = compiler[: -len(ARM_V80COMPAT_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-aarch64.cmake"
        )
        cmake_flags.append("-DNO_ARMV81_OR_HIGHER=1")
        result.append("DEB_ARCH=arm64")
    elif is_cross_freebsd:
        cc = compiler[: -len(FREEBSD_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/freebsd/toolchain-x86_64.cmake"
        )
    elif is_cross_ppc:
        cc = compiler[: -len(PPC_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-ppc64le.cmake"
        )
    elif is_cross_riscv:
        cc = compiler[: -len(RISCV_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-riscv64.cmake"
        )
    elif is_cross_s390x:
        cc = compiler[: -len(S390X_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-s390x.cmake"
        )
    elif is_cross_loongarch:
        cc = compiler[: -len(LOONGARCH_SUFFIX)]
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-loongarch64.cmake"
        )
    elif is_amd64_compat:
        cc = compiler[: -len(AMD64_COMPAT_SUFFIX)]
        result.append("DEB_ARCH=amd64")
        cmake_flags.append("-DNO_SSE3_OR_HIGHER=1")
    elif is_amd64_musl:
        cc = compiler[: -len(AMD64_MUSL_SUFFIX)]
        result.append("DEB_ARCH=amd64")
        cmake_flags.append(
            "-DCMAKE_TOOLCHAIN_FILE=/build/cmake/linux/toolchain-x86_64-musl.cmake"
        )
    else:
        cc = compiler
        result.append("DEB_ARCH=amd64")

    cxx = cc.replace("clang", "clang++")

    if package_type == "deb":
        # NOTE: This is the env for packages/build script
        result.append("MAKE_DEB=true")
        cmake_flags.append("-DENABLE_TESTS=0")
        cmake_flags.append("-DENABLE_UTILS=0")
        cmake_flags.append("-DCMAKE_FIND_PACKAGE_NO_PACKAGE_REGISTRY=ON")
        cmake_flags.append("-DCMAKE_INSTALL_PREFIX=/usr")
        cmake_flags.append("-DCMAKE_INSTALL_SYSCONFDIR=/etc")
        cmake_flags.append("-DCMAKE_INSTALL_LOCALSTATEDIR=/var")
        # Reduce linking and building time by avoid *install/all dependencies
        cmake_flags.append("-DCMAKE_SKIP_INSTALL_ALL_DEPENDENCY=ON")
        # Add bridges to the build target
        build_target = (
            f"{build_target} clickhouse-odbc-bridge clickhouse-library-bridge"
        )
        if is_release_build(debug_build, package_type, sanitizer, coverage):
            cmake_flags.append("-DSPLIT_DEBUG_SYMBOLS=ON")
            result.append("WITH_PERFORMANCE=1")
            cmake_flags.append("-DBUILD_STANDALONE_KEEPER=1")
    elif package_type == "fuzzers":
        cmake_flags.append("-DENABLE_FUZZING=1")
        cmake_flags.append("-DENABLE_PROTOBUF=1")
        cmake_flags.append("-DWITH_COVERAGE=1")
        # Reduce linking and building time by avoid *install/all dependencies
        cmake_flags.append("-DCMAKE_SKIP_INSTALL_ALL_DEPENDENCY=ON")

    result.append(f"CC={cc}")
    result.append(f"CXX={cxx}")
    cmake_flags.append(f"-DCMAKE_C_COMPILER={cc}")
    cmake_flags.append(f"-DCMAKE_CXX_COMPILER={cxx}")

    if sanitizer:
        result.append(f"SANITIZER={sanitizer}")
    if debug_build:
        result.append("BUILD_TYPE=Debug")
    else:
        result.append("BUILD_TYPE=None")

    if coverage:
        cmake_flags.append("-DSANITIZE_COVERAGE=1 -DBUILD_STANDALONE_KEEPER=0")

    if not cache:
        cmake_flags.append("-DCOMPILER_CACHE=disabled")

    if cache == "ccache":
        cmake_flags.append("-DCOMPILER_CACHE=ccache")
        result.append("CCACHE_DIR=/ccache")
        result.append("CCACHE_COMPRESSLEVEL=5")
        result.append("CCACHE_BASEDIR=/build")
        result.append("CCACHE_NOHASHDIR=true")
        result.append("CCACHE_COMPILERCHECK=content")
        result.append("CCACHE_MAXSIZE=15G")

    if cache == "sccache":
        cmake_flags.append("-DCOMPILER_CACHE=sccache")
        # see https://github.com/mozilla/sccache/blob/main/docs/S3.md
        result.append(f"SCCACHE_BUCKET={s3_bucket}")
        sccache_dir = "sccache"
        if s3_directory:
            sccache_dir = f"{s3_directory}/{sccache_dir}"
        result.append(f"SCCACHE_S3_KEY_PREFIX={sccache_dir}")
        if not s3_rw_access:
            result.append("SCCACHE_S3_NO_CREDENTIALS=true")

    if clang_tidy:
        # `CTCACHE_DIR` has the same purpose as the `CCACHE_DIR` above.
        # It's there to have the clang-tidy cache embedded into our standard `CCACHE_DIR`
        if cache == "ccache":
            result.append("CTCACHE_DIR=/ccache/clang-tidy-cache")
        if s3_bucket:
            # see https://github.com/matus-chochlik/ctcache#environment-variables
            ctcache_dir = "clang-tidy-cache"
            if s3_directory:
                ctcache_dir = f"{s3_directory}/{ctcache_dir}"
            result.append(f"CTCACHE_S3_BUCKET={s3_bucket}")
            result.append(f"CTCACHE_S3_FOLDER={ctcache_dir}")
            if not s3_rw_access:
                result.append("CTCACHE_S3_NO_CREDENTIALS=true")

    if additional_pkgs:
        # NOTE: This are the env for packages/build script
        result.append("MAKE_RPM=true")
        result.append("MAKE_TGZ=true")

    if with_binaries == "programs":
        result.append("BINARY_OUTPUT=programs")
    elif with_binaries == "tests":
        result.append("ENABLE_TESTS=1")
        result.append("BINARY_OUTPUT=tests")
        cmake_flags.append("-DENABLE_TESTS=1")

    if clang_tidy:
        cmake_flags.append("-DENABLE_CLANG_TIDY=1")
        cmake_flags.append("-DENABLE_TESTS=1")
        cmake_flags.append("-DENABLE_EXAMPLES=1")
        # Don't stop on first error to find more clang-tidy errors in one run.
        result.append("NINJA_FLAGS=-k0")

        cmake_flags.append("-DENABLE_UTILS=1")
        # utils are not included into clickhouse-bundle, so build everything
        build_target = "all"

    if with_profiler:
        cmake_flags.append("-DENABLE_BUILD_PROFILING=1")

    if with_coverage:
        cmake_flags.append("-DWITH_COVERAGE=1")

    if version:
        result.append(f"VERSION_STRING='{version}'")

    if official:
        cmake_flags.append("-DCLICKHOUSE_OFFICIAL_BUILD=1")

    result.append('CMAKE_FLAGS="' + " ".join(cmake_flags) + '"')
    result.append(f"BUILD_TARGET='{build_target}'")

    return result


def dir_name(name: str) -> Path:
    path = Path(name)
    if not path.is_absolute():
        path = Path.cwd() / name
    return path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="ClickHouse building script using prebuilt Docker image",
    )
    parser.add_argument(
        "--package-type",
        choices=["deb", "binary", "fuzzers"],
        required=True,
    )
    parser.add_argument(
        "--clickhouse-repo-path",
        default=SCRIPT_PATH.parents[2],
        type=dir_name,
        help="ClickHouse git repository",
    )
    parser.add_argument("--output-dir", type=dir_name, required=True)
    parser.add_argument("--debug-build", action="store_true")

    parser.add_argument(
        "--compiler",
        choices=(
            "clang-18",
            "clang-18-darwin",
            "clang-18-darwin-aarch64",
            "clang-18-aarch64",
            "clang-18-aarch64-v80compat",
            "clang-18-ppc64le",
            "clang-18-riscv64",
            "clang-18-s390x",
            "clang-18-loongarch64",
            "clang-18-amd64-compat",
            "clang-18-amd64-musl",
            "clang-18-freebsd",
        ),
        default="clang-18",
        help="a compiler to use",
    )
    parser.add_argument(
        "--sanitizer",
        choices=("address", "thread", "memory", "undefined", ""),
        default="",
    )
    parser.add_argument(
        "--coverage",
        action="store_true",
        help="enable granular coverage with introspection",
    )

    parser.add_argument("--clang-tidy", action="store_true")
    parser.add_argument(
        "--cache",
        choices=("ccache", "sccache", ""),
        default="",
        help="ccache or sccache for objects caching; sccache uses only S3 buckets",
    )
    parser.add_argument(
        "--ccache-dir",
        default=Path.home() / ".ccache",
        type=dir_name,
        help="a directory with ccache",
    )
    parser.add_argument(
        "--s3-bucket",
        help="an S3 bucket used for sscache and clang-tidy-cache",
    )
    parser.add_argument(
        "--s3-directory",
        default="ccache",
        help="an S3 directory prefix used for sscache and clang-tidy-cache",
    )
    parser.add_argument(
        "--s3-rw-access",
        action="store_true",
        help="if set, the build fails on errors writing cache to S3",
    )
    parser.add_argument(
        "--cargo-cache-dir",
        default=Path(os.getenv("CARGO_HOME", "") or Path.home() / ".cargo")
        / "registry",
        type=dir_name,
        help="a directory to preserve the rust cargo crates",
    )
    parser.add_argument("--force-build-image", action="store_true")
    parser.add_argument("--version")
    parser.add_argument("--official", action="store_true")
    parser.add_argument("--additional-pkgs", action="store_true")
    parser.add_argument("--with-profiler", action="store_true")
    parser.add_argument("--with-coverage", action="store_true")
    parser.add_argument(
        "--with-binaries", choices=("programs", "tests", ""), default=""
    )
    parser.add_argument(
        "--docker-image-version", default="latest", help="docker image tag to use"
    )
    parser.add_argument(
        "--as-root", action="store_true", help="if the container should run as root"
    )

    args = parser.parse_args()

    if args.additional_pkgs and args.package_type != "deb":
        raise argparse.ArgumentTypeError(
            "Can build additional packages only in deb build"
        )

    if args.cache != "ccache":
        args.ccache_dir = None

    if args.with_binaries != "":
        if args.package_type != "deb":
            raise argparse.ArgumentTypeError(
                "Can add additional binaries only in deb build"
            )
        logging.info("Should place %s to output", args.with_binaries)

    if args.cache == "sccache":
        if not args.s3_bucket:
            raise argparse.ArgumentTypeError("sccache must have --s3-bucket set")

    return args


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    args = parse_args()

    ch_root = args.clickhouse_repo_path

    dockerfile = ch_root / "docker/packager" / IMAGE_TYPE / "Dockerfile"
    image_with_version = IMAGE_NAME + ":" + args.docker_image_version
    if args.force_build_image:
        build_image(image_with_version, dockerfile)
    elif not (
        check_image_exists_locally(image_with_version) or pull_image(image_with_version)
    ):
        build_image(image_with_version, dockerfile)

    env_prepared = parse_env_variables(
        args.debug_build,
        args.coverage,
        args.compiler,
        args.sanitizer,
        args.package_type,
        args.cache,
        args.s3_bucket,
        args.s3_directory,
        args.s3_rw_access,
        args.clang_tidy,
        args.version,
        args.official,
        args.additional_pkgs,
        args.with_profiler,
        args.with_coverage,
        args.with_binaries,
    )

    pre_build(args.clickhouse_repo_path, env_prepared)
    run_docker_image_with_env(
        image_with_version,
        args.as_root,
        args.output_dir,
        env_prepared,
        ch_root,
        args.cargo_cache_dir,
        args.ccache_dir,
    )
    logging.info("Output placed into %s", args.output_dir)


if __name__ == "__main__":
    main()