spark/tools/py/deploy.py

import argparse
import dotenv
import subprocess

import publish_data_stream
import publish_sparkd

from core.argparse_utils import add_toggle_opt
from core.file_utils import Path, SparkPaths


def _get_arg_parser():
    parser = argparse.ArgumentParser(
        description=
        "Deploy sparkd, pubsub, md-publisher, webapp, and api server")
    parser.add_argument(
        "--cmake-jobs",
        type=int,
        help="The number of cmake parallel jobs",
    )
    parser.add_argument(
        "--profile",
        help="The configuration profile dirs to use, relative to each service")
    parser.add_argument("--compose-file",
                        help="The docker compose (yaml) file to use")
    parser.add_argument(
        "--clean",
        action="store_true",
        help="Perform a clean build (remove artifacts from previous build)",
    )
    parser.add_argument("--env-files",
                        action="append",
                        help=".env files to pass to docker compose")
    parser.add_argument(
        "--verbosity",
        metavar="{0-5}",
        choices=range(1, 6),
        type=int,
        help="The output verbosity level",
    )
    add_toggle_opt(parser,
                   "build",
                   help="Whether or not to build the services first")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview the build commands without running them",
    )
    # TODO: add more services (api-server, pubsub, dsqa web server, etc.)
    all_services = ["sparkd"]
    parser.add_argument("services",
                        help="this selects services defined in compose.yaml",
                        nargs="*")
    parser.set_defaults(
        configuration="release",
        services=all_services,
        build=True,
        clean=False,
        profile="dev",
        verbosity=2,
    )
    return parser


def _resolve_profile_dir(base_path: Path, profile: str):
    def resolve_as(profile: str):
        profile_dir = base_path.joinpath(profile)
        if profile_dir.exists():
            return profile_dir
        for p in profile_dir.parent.iterdir():
            if p.is_dir() and p.name.lower() == profile_dir.name.lower():
                return p
        return None

    # Use aliases to unify different naming conventions in c++ and c# projects
    known_aliases = [("dev", "Development"), ("prod", "Production")]
    res = resolve_as(profile)
    if not res:
        for alias_grp in known_aliases:
            if any(filter(lambda a: a.lower() == profile, alias_grp)):
                res = next(filter(lambda a: a, map(resolve_as, alias_grp)), None)
                break
    if res:
        return res
    raise ValueError(f"Profile '{profile}' not found")


def _get_env_files(services, profile):
    root = SparkPaths.root()
    files = []

    def add(file):
        if file.exists():
            files.append(file)

    add(root.joinpath(".env"))
    add(root.joinpath(".env.local"))
    if profile is not None:
        if "sparkd" in services or "sparkd-symbols" in services:
            profile_dir=_resolve_profile_dir(root.joinpath("src/app_sparkd/configurations"), profile)
            add(profile_dir.joinpath(".env"))
            add(profile_dir.joinpath(".env.local"))
        if "pubsub" in services:
            profile_dir = _resolve_profile_dir(root.joinpath("src/DataStreamSuite/PubSubServer/Configurations"), profile)
            add(profile_dir.joinpath(".env"))
            add(profile_dir.joinpath(".env.local"))
        if "md-publisher" in services:
            profile_dir = _resolve_profile_dir(root.joinpath("src/DataStreamSuite/DataStreamPublisher/Configurations"), profile)
            add(profile_dir.joinpath(".env"))
            add(profile_dir.joinpath(".env.local"))
    return files


def deploy(args):
    def run(cmd):
        print(cmd)
        if args.dry_run:
            return
        subprocess.run(cmd, shell=True).check_returncode()

    services: list[str] = args.services or []
    if args.build:
        should_build_sparkd = len(services) == 0 or any(
            filter(lambda s: s.startswith("sparkd"), services))
        if should_build_sparkd:
            publish_sparkd.publish(
                publish_sparkd.PublishArgs(
                    configuration="release",
                    cmake_jobs=args.cmake_jobs,
                    clean=args.clean,
                    verbosity=args.verbosity,
                    dry_run=args.dry_run,
                ))
        should_build_pubsub = any(filter(lambda s: s.startswith("pubsub") or s.startswith("md-publisher"), services))
        if should_build_pubsub:
            publish_data_stream.publish(
                publish_data_stream.PublishArgs(
                    configuration="release",
                    cmake_jobs=args.cmake_jobs,
                    clean=args.clean,
                    verbosity=args.verbosity,
                    dry_run=args.dry_run,
                )
            )
    service_args = " ".join(services)
    compose_cmd = "sudo docker compose"
    if args.compose_file:
        compose_cmd += f" -f {args.compose_file}"
    env_files = _get_env_files(args.services, args.profile)
    if len(env_files) > 0:
        compose_cmd += " " + " ".join([f"--env-file {f}" for f in env_files])
        for ef in env_files:
            if args.verbosity > 2:
                print(f"Loading env file {ef}")
                if args.verbosity > 3:
                    env_vars = dotenv.dotenv_values(ef)
                    for k, v in env_vars.items():
                        print(f"    {k}={v}")
            dotenv.load_dotenv(ef, override=True)
    run(f"{compose_cmd} down {service_args}")
    run(f"{compose_cmd} build {service_args}")
    run(f"{compose_cmd} up {service_args} -d")
    run(f"{compose_cmd} logs {service_args} --follow")


def _main():
    parser = _get_arg_parser()
    args = parser.parse_args()
    deploy(args)


if __name__ == "__main__":
    try:
        _main()
    except Exception as e:
        print(f"deploy.py failed: {e}")
        exit(1)