prole/knoe/core/ops/gke_clusters.py

"""Helpers for explicit app/db GKE cluster targeting in installer flows."""

from __future__ import annotations

import os
import subprocess
from dataclasses import dataclass
from typing import Callable


_LogFn = Callable[[str], None]


@dataclass(frozen=True)
class GkeClusterSpec:
    name: str
    mode: str
    location: str
    machine_type: str = "e2-small"
    node_count: int = 3
    node_pool_name: str = "default-pool"
    boot_disk_type: str = "pd-standard"   # pd-standard keeps boot disks off pd-ssd quota
    boot_disk_size_gb: int = 50


def _run(cmd: list[str], log: _LogFn | None = None) -> subprocess.CompletedProcess:
    if log:
        log("$ " + " ".join(cmd))
    return subprocess.run(cmd, capture_output=True, text=True)


def build_kubectl_env_for_cluster(
    base_env: dict | None,
    kubecontext: str,
    cluster_name: str,
    cluster_role: str,
) -> dict:
    env = dict(base_env or os.environ)
    if kubecontext:
        env["KUBECTL_CONTEXT"] = kubecontext
    env["CLUSTER_NAME"] = cluster_name
    env["KNOE_CLUSTER_ROLE"] = cluster_role
    if cluster_role == "app":
        env["KNOE_APP_CLUSTER_NAME"] = cluster_name
    elif cluster_role == "db":
        env["KNOE_DB_CLUSTER_NAME"] = cluster_name
    return env


def get_cluster_credentials(
    *,
    project_id: str,
    cluster_name: str,
    location: str,
    log: _LogFn | None = None,
) -> str:
    cmd = [
        "gcloud",
        "container",
        "clusters",
        "get-credentials",
        cluster_name,
        "--project",
        project_id,
        "--region",
        location,
        "--quiet",
    ]
    result = _run(cmd, log=log)
    if result.returncode != 0:
        stderr = (result.stderr or "").strip()
        raise RuntimeError(
            f"Failed to get credentials for cluster '{cluster_name}' ({location}): {stderr}"
        )
    return f"gke_{project_id}_{location}_{cluster_name}"


def ensure_app_cluster(
    *,
    project_id: str,
    spec: GkeClusterSpec,
    log: _LogFn | None = None,
) -> None:
    describe_cmd = [
        "gcloud",
        "container",
        "clusters",
        "describe",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--format=value(name)",
        "--quiet",
    ]
    result = _run(describe_cmd, log=log)
    if result.returncode == 0 and (result.stdout or "").strip() == spec.name:
        return

    # Cluster not found — create it as Autopilot.
    # NOTE: gcloud container clusters create-auto does NOT support --disk-type
    # or --disk-size.  Autopilot manages all node infrastructure automatically.
    # Autopilot only provisions physical nodes when pods are scheduled, so SSD
    # quota is not consumed while the cluster is idle.
    create_cmd = [
        "gcloud",
        "container",
        "clusters",
        "create-auto",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--workload-policies=allow-net-admin",
        "--quiet",
    ]
    created = _run(create_cmd, log=log)
    if created.returncode != 0:
        stderr = (created.stderr or "").strip()
        raise RuntimeError(
            f"Failed to create app cluster '{spec.name}' in '{spec.location}': {stderr}"
        )


def ensure_db_cluster(
    *,
    project_id: str,
    spec: GkeClusterSpec,
    log: _LogFn | None = None,
) -> None:
    describe_cmd = [
        "gcloud",
        "container",
        "clusters",
        "describe",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--format=value(name)",
        "--quiet",
    ]
    describe = _run(describe_cmd, log=log)
    if describe.returncode == 0 and (describe.stdout or "").strip() == spec.name:
        return

    if spec.mode.lower() != "standard":
        raise RuntimeError(
            f"Database cluster '{spec.name}' must run in Standard mode, got '{spec.mode}'."
        )

    create_cmd = [
        "gcloud",
        "container",
        "clusters",
        "create",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--num-nodes",
        str(max(3, int(spec.node_count))),
        "--machine-type",
        spec.machine_type,
        "--disk-type",
        spec.boot_disk_type,
        "--disk-size",
        str(spec.boot_disk_size_gb),
        "--enable-ip-alias",
        "--workload-pool",
        f"{project_id}.svc.id.goog",
        "--quiet",
    ]
    created = _run(create_cmd, log=log)
    if created.returncode != 0:
        stderr = (created.stderr or "").strip()
        raise RuntimeError(f"Failed to create DB cluster '{spec.name}': {stderr}")