"""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) if kubecontext: env["KUBECONTEXT"] = kubecontext env["KUBE_CONTEXT_NAME"] = 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. 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}" )