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"
    node_labels: str = "workload=db"
    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 _nodes_per_location(node_count: int, location: str) -> int:
    """Return GKE --num-nodes value, treating regional locations as per-zone counts.

    For regional clusters, gcloud interprets --num-nodes as per-zone, so convert
    requested total nodes to a per-zone value (ceil(total/3)) for the common
    3-zone regional layout.
    """
    requested = max(1, int(node_count))
    loc = (location or "").strip().lower()
    suffix = loc.rsplit("-", 1)[-1] if "-" in loc else ""
    is_zone = len(suffix) == 1 and suffix.isalpha()
    if is_zone:
        return requested
    return max(1, (requested + 2) // 3)


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)
    role = (cluster_role or "").strip().lower()
    if role not in {"app", "db"}:
        raise ValueError(f"Unsupported cluster role '{cluster_role}'. Expected 'app' or 'db'.")

    required_key = "APP_CLUSTER_KUBECONTEXT" if role == "app" else "DB_CLUSTER_KUBECONTEXT"
    opposite_key = "DB_CLUSTER_KUBECONTEXT" if role == "app" else "APP_CLUSTER_KUBECONTEXT"
    required_ctx = str(env.get(required_key) or "").strip()
    opposite_ctx = str(env.get(opposite_key) or "").strip()
    explicit_ctx = str(kubecontext or "").strip() or required_ctx

    if not explicit_ctx:
        raise ValueError(
            f"Missing required kube context for {role.upper()} cluster targeting. "
            f"Set {required_key} and pass an explicit context."
        )
    if opposite_ctx and explicit_ctx == opposite_ctx:
        raise ValueError(
            f"Refusing {role.upper()} operation: kube context '{explicit_ctx}' matches "
            f"{opposite_key}."
        )
    if required_ctx and explicit_ctx != required_ctx:
        raise ValueError(
            f"Refusing {role.upper()} operation with kube context '{explicit_ctx}': "
            f"expected {required_key}='{required_ctx}'."
        )

    env["KUBECTEXT"] = explicit_ctx
    env["KUBE_CONTEXT_NAME"] = explicit_ctx
    env["KUBECTL_CONTEXT"] = explicit_ctx
    if role == "app":
        env.setdefault("APP_CLUSTER_KUBECONTEXT", explicit_ctx)
    else:
        env.setdefault("DB_CLUSTER_KUBECONTEXT", explicit_ctx)
    env["CLUSTER_NAME"] = cluster_name
    env["KNOE_CLUSTER_ROLE"] = role
    if role == "app":
        env["KNOE_APP_CLUSTER_NAME"] = cluster_name
    elif 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.
    if spec.mode.lower() == "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",
        ]
    else:
        # Standard mode: explicit machine type, pd-standard boot disks (no SSD quota),
        # VPA for dynamic resource adjustment, Workload Identity for GCS access.
        create_cmd = [
            "gcloud",
            "container",
            "clusters",
            "create",
            spec.name,
            "--project",
            project_id,
            "--region",
            spec.location,
            "--cluster-version=latest",
            "--machine-type",
            spec.machine_type,
            "--disk-type",
            spec.boot_disk_type,
            "--disk-size",
            str(spec.boot_disk_size_gb),
            "--num-nodes",
            str(max(1, int(spec.node_count))),
            "--enable-vertical-pod-autoscaling",
            "--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 app cluster '{spec.name}' in '{spec.location}': {stderr}"
        )


def ensure_db_cluster(
    *,
    project_id: str,
    spec: GkeClusterSpec,
    log: _LogFn | None = None,
) -> None:
    if spec.mode.lower() != "standard":
        raise RuntimeError(
            f"Database cluster '{spec.name}' must run in Standard mode, got '{spec.mode}'."
        )

    describe_cmd = [
        "gcloud",
        "container",
        "clusters",
        "describe",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--format=value(name)",
        "--quiet",
    ]
    describe = _run(describe_cmd, log=log)
    cluster_exists = describe.returncode == 0 and (describe.stdout or "").strip() == spec.name

    if not cluster_exists:
        create_cmd = [
            "gcloud",
            "container",
            "clusters",
            "create",
            spec.name,
            "--project",
            project_id,
            "--region",
            spec.location,
            "--num-nodes",
            str(_nodes_per_location(spec.node_count, spec.location)),
            "--machine-type",
            spec.machine_type,
            "--disk-type",
            spec.boot_disk_type,
            "--disk-size",
            str(spec.boot_disk_size_gb),
            "--node-labels",
            spec.node_labels,
            "--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}")
        return

    pool_describe_cmd = [
        "gcloud",
        "container",
        "node-pools",
        "describe",
        spec.node_pool_name,
        "--cluster",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--format=value(name)",
        "--quiet",
    ]
    pool_describe = _run(pool_describe_cmd, log=log)
    pool_exists = (
        pool_describe.returncode == 0
        and (pool_describe.stdout or "").strip() == spec.node_pool_name
    )
    if pool_exists:
        update_labels_cmd = [
            "gcloud",
            "container",
            "node-pools",
            "update",
            spec.node_pool_name,
            "--cluster",
            spec.name,
            "--project",
            project_id,
            "--region",
            spec.location,
            "--node-labels",
            spec.node_labels,
            "--quiet",
        ]
        updated = _run(update_labels_cmd, log=log)
        if updated.returncode != 0:
            stderr = (updated.stderr or "").strip()
            raise RuntimeError(
                f"Failed to label DB node pool '{spec.node_pool_name}' for cluster '{spec.name}': {stderr}"
            )
        return

    create_pool_cmd = [
        "gcloud",
        "container",
        "node-pools",
        "create",
        spec.node_pool_name,
        "--cluster",
        spec.name,
        "--project",
        project_id,
        "--region",
        spec.location,
        "--machine-type",
        spec.machine_type,
        "--disk-type",
        spec.boot_disk_type,
        "--disk-size",
        str(spec.boot_disk_size_gb),
        "--num-nodes",
        str(_nodes_per_location(spec.node_count, spec.location)),
        "--node-labels",
        spec.node_labels,
        "--quiet",
    ]
    pool_created = _run(create_pool_cmd, log=log)
    if pool_created.returncode != 0:
        stderr = (pool_created.stderr or "").strip()
        raise RuntimeError(
            f"Failed to create DB node pool '{spec.node_pool_name}' in cluster '{spec.name}': {stderr}"
        )