lingniu-vehicle-ingest/tools/go_native_deploy.py

#!/usr/bin/env python3
"""Build and deploy the Go native vehicle gateway release to ECS."""

from __future__ import annotations

import argparse
import os
import pathlib
import shutil
import subprocess
import sys
import tarfile
import tempfile
from dataclasses import dataclass


DEFAULT_APP_HOST = "115.29.187.205"
DEFAULT_KAFKA_HOST = "114.55.58.251"
DEFAULT_USER = "root"
DEFAULT_RELEASE_ROOT = "/opt/lingniu-go-native"
DEFAULT_GO_DIR = "go/vehicle-gateway"
BINARIES = [
    ("gateway", "./cmd/gateway"),
    ("history-writer", "./cmd/history-writer"),
    ("stat-writer", "./cmd/stat-writer"),
    ("realtime-api", "./cmd/realtime-api"),
]
SERVICES = [
    "lingniu-go-gateway",
    "lingniu-go-history-writer",
    "lingniu-go-stat-writer",
    "lingniu-go-realtime-api",
]


@dataclass(frozen=True)
class BuildCommand:
    output_name: str
    argv: list[str]
    env: list[str]
    cwd: str


def git_short_sha(cwd: str) -> str:
    return run_capture(["git", "rev-parse", "--short=7", "HEAD"], cwd=cwd).strip()


def build_commands(go_dir: str, output_dir: str) -> list[BuildCommand]:
    commands: list[BuildCommand] = []
    env = ["CGO_ENABLED=0", "GOOS=linux", "GOARCH=amd64"]
    for output_name, package in BINARIES:
        commands.append(BuildCommand(
            output_name=output_name,
            argv=[
                "go",
                "build",
                "-trimpath",
                "-ldflags=-s -w",
                "-o",
                str(pathlib.Path(output_dir) / output_name),
                package,
            ],
            env=env,
            cwd=go_dir,
        ))
    return commands


def build_release(go_dir: str, output_dir: str) -> None:
    for command in build_commands(go_dir, output_dir):
        env = os.environ.copy()
        for item in command.env:
            key, value = item.split("=", 1)
            env[key] = value
        run(command.argv, cwd=command.cwd, env=env)


def create_archive(source_dir: str, archive_path: str) -> None:
    with tarfile.open(archive_path, "w:gz") as archive:
        for binary, _ in BINARIES:
            archive.add(pathlib.Path(source_dir) / binary, arcname=binary)


def remote_switch_command(release: str, release_root: str, remote_archive: str) -> str:
    if not safe_release_name(release):
        raise ValueError("release must contain only letters, numbers, dot, underscore, or dash")
    release_dir = release_root.rstrip("/") + "/releases/" + release
    root = release_root.rstrip("/")
    commands = [
        "set -euo pipefail",
        "mkdir -p " + shell_quote(release_dir),
        "tar -xzf " + shell_quote(remote_archive) + " -C " + shell_quote(release_dir),
        "chmod +x " + shell_quote(release_dir) + "/*",
        "ln -sfn " + shell_quote(release_dir) + " " + shell_quote(root + "/current"),
        "systemctl daemon-reload",
    ]
    commands.extend("systemctl restart " + service for service in SERVICES)
    commands.append("rm -f " + shell_quote(remote_archive))
    return " && ".join(commands)


def acceptance_command(app_host: str, kafka_host: str, date: str, timeout: float) -> list[str]:
    command = [
        "python3",
        "tools/go_prod_acceptance.py",
        "--app-host",
        app_host,
        "--kafka-host",
        kafka_host,
        "--timeout",
        str(int(timeout) if timeout == int(timeout) else timeout),
    ]
    if date:
        command.extend(["--date", date])
    return command


def deploy(args: argparse.Namespace) -> None:
    repo = pathlib.Path(args.repo).resolve()
    go_dir = str(repo / args.go_dir)
    release = args.release or git_short_sha(str(repo))
    with tempfile.TemporaryDirectory(prefix="lingniu-go-native-") as tmp:
        output_dir = pathlib.Path(tmp) / "out"
        output_dir.mkdir()
        archive_path = pathlib.Path(tmp) / ("lingniu-go-native-" + release + ".tar.gz")
        build_release(go_dir, str(output_dir))
        create_archive(str(output_dir), str(archive_path))
        remote_archive = "/tmp/" + archive_path.name
        target = ssh_target(args.user, args.app_host)
        run(["scp", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", str(archive_path), target + ":" + remote_archive])
        run([
            "ssh",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            target,
            remote_switch_command(release, args.release_root, remote_archive),
        ])
    if not args.skip_acceptance:
        run(acceptance_command(args.app_host, args.kafka_host, args.date, args.timeout), cwd=str(repo))


def ssh_target(user: str, host: str) -> str:
    return user + "@" + host if user else host


def shell_quote(value: str) -> str:
    return "'" + value.replace("'", "'\"'\"'") + "'"


def safe_release_name(value: str) -> bool:
    value = value.strip()
    if not value:
        return False
    return all(r.isalnum() or r in "._-" for r in value)


def run(argv: list[str], cwd: str | None = None, env: dict[str, str] | None = None) -> None:
    subprocess.run(argv, cwd=cwd, env=env, check=True)


def run_capture(argv: list[str], cwd: str | None = None) -> str:
    return subprocess.run(argv, cwd=cwd, check=True, stdout=subprocess.PIPE, text=True).stdout


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", default=".")
    parser.add_argument("--go-dir", default=DEFAULT_GO_DIR)
    parser.add_argument("--app-host", default=DEFAULT_APP_HOST)
    parser.add_argument("--kafka-host", default=DEFAULT_KAFKA_HOST)
    parser.add_argument("--user", default=DEFAULT_USER)
    parser.add_argument("--release-root", default=DEFAULT_RELEASE_ROOT)
    parser.add_argument("--release", default="")
    parser.add_argument("--date", default="")
    parser.add_argument("--timeout", type=float, default=20.0)
    parser.add_argument("--skip-acceptance", action="store_true")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    try:
        deploy(args)
    except subprocess.CalledProcessError as exc:
        print("command failed: " + " ".join(exc.cmd), file=sys.stderr)
        return exc.returncode or 1
    except (OSError, shutil.Error) as exc:
        print("deploy failed: " + str(exc), file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))