prole/knoe/core/ops/openbao.py

from __future__ import annotations

from pathlib import Path

from ._services_common import (
    _LogFn,
    _detect_mode,
    _ensure_namespace,
    _exists,
    _kubectl,
    _log,
    _manifest_path,
    _namespace,
    _to_bool,
    _wait_rollout,
)


def _resource_namespace(namespace: str | None, env: dict | None) -> str:
    if env:
        explicit = str(env.get("OPENBAO_NAMESPACE") or "").strip()
        if explicit:
            return explicit
    return _namespace(namespace, env, default="default")


def _manifest_files(mode: str, project_root: str | Path, include_kerberos: bool) -> list[Path]:
    root = Path(project_root)
    if mode == "k3s":
        files = [
            _manifest_path(root, "k8s", "prole", "openbao-service.yaml"),
            _manifest_path(root, "k8s", "prole", "openbao-statefulset.yaml"),
        ]
    else:
        files = [_manifest_path(root, "k8s", "openbao", "deployment.yaml")]
    if include_kerberos:
        files.append(_manifest_path(root, "k8s", "openbao", "kerberos-configmap.yaml"))
    return files


def _apply(
    *,
    namespace: str,
    mode: str,
    env: dict | None,
    project_root: str | Path,
    log: _LogFn | None,
    include_kerberos_configmap: bool,
) -> None:
    _ensure_namespace(namespace, env)
    for manifest in _manifest_files(mode, project_root, include_kerberos_configmap):
        if not manifest.exists():
            continue
        if manifest.name == "kerberos-configmap.yaml":
            _ensure_namespace("knoe-system", env)
        _log(log, f"[OPENBAO] Applying {manifest}")
        try:
            _kubectl(
                ["-n", namespace, "apply", "-f", str(manifest)],
                env=env,
                timeout=300,
                check=True,
            )
        except RuntimeError as exc:
            msg = str(exc)
            if "does not match the namespace" not in msg:
                raise
            _log(
                log,
                "[OPENBAO] Namespace mismatch in manifest metadata; retrying apply without forced namespace.",
            )
            _kubectl(
                ["apply", "-f", str(manifest)],
                env=env,
                timeout=300,
                check=True,
            )

    if mode == "k3s" or _exists("statefulset", "openbao", namespace, env=env):
        _wait_rollout("statefulset", "openbao", namespace, env=env)
    else:
        _wait_rollout("deployment", "openbao", namespace, env=env)


def initialize(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
    include_kerberos_configmap: bool = True,
) -> None:
    update(
        namespace=namespace,
        env=env,
        project_root=project_root,
        log=log,
        mode=mode,
        include_kerberos_configmap=include_kerberos_configmap,
    )


def start(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
    include_kerberos_configmap: bool = True,
) -> None:
    update(
        namespace=namespace,
        env=env,
        project_root=project_root,
        log=log,
        mode=mode,
        include_kerberos_configmap=include_kerberos_configmap,
    )


def update(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
    include_kerberos_configmap: bool = True,
) -> None:
    effective_mode = _detect_mode(mode, env)
    target_ns = _resource_namespace(namespace, env)
    _apply(
        namespace=target_ns,
        mode=effective_mode,
        env=env,
        project_root=project_root,
        log=log,
        include_kerberos_configmap=include_kerberos_configmap,
    )


def stop(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    log: _LogFn | None = None,
) -> None:
    target_ns = _resource_namespace(namespace, env)
    if _exists("statefulset", "openbao", target_ns, env=env):
        _log(log, f"[OPENBAO] Scaling statefulset/openbao to 0 in namespace {target_ns}")
        _kubectl(
            [
                "-n",
                target_ns,
                "scale",
                "statefulset/openbao",
                "--replicas=0",
            ],
            env=env,
            timeout=90,
            check=True,
        )
        return
    if _exists("deployment", "openbao", target_ns, env=env):
        _log(log, f"[OPENBAO] Scaling deployment/openbao to 0 in namespace {target_ns}")
        _kubectl(
            [
                "-n",
                target_ns,
                "scale",
                "deployment/openbao",
                "--replicas=0",
            ],
            env=env,
            timeout=90,
            check=True,
        )


def restart(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    log: _LogFn | None = None,
) -> None:
    target_ns = _resource_namespace(namespace, env)
    if _exists("statefulset", "openbao", target_ns, env=env):
        _log(log, f"[OPENBAO] Restarting statefulset/openbao in namespace {target_ns}")
        _kubectl(
            ["-n", target_ns, "rollout", "restart", "statefulset/openbao"],
            env=env,
            timeout=120,
            check=True,
        )
        _wait_rollout("statefulset", "openbao", target_ns, env=env)
        return
    if _exists("deployment", "openbao", target_ns, env=env):
        _log(log, f"[OPENBAO] Restarting deployment/openbao in namespace {target_ns}")
        _kubectl(
            ["-n", target_ns, "rollout", "restart", "deployment/openbao"],
            env=env,
            timeout=120,
            check=True,
        )
        _wait_rollout("deployment", "openbao", target_ns, env=env)


def status(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    mode: str | None = None,
) -> bool:
    target_ns = _resource_namespace(namespace, env)
    effective_mode = _detect_mode(mode, env)
    ok = True
    if effective_mode == "k3s" or _exists("statefulset", "openbao", target_ns, env=env):
        sts = _kubectl(
            [
                "-n",
                target_ns,
                "get",
                "statefulset",
                "openbao",
                "-o",
                "jsonpath={.status.readyReplicas}",
            ],
            env=env,
            timeout=20,
        )
        ok = ok and sts.returncode == 0 and (sts.stdout or "0").strip() not in {"", "0"}
    else:
        dep = _kubectl(
            [
                "-n",
                target_ns,
                "get",
                "deployment",
                "openbao",
                "-o",
                "jsonpath={.status.readyReplicas}",
            ],
            env=env,
            timeout=20,
        )
        ok = ok and dep.returncode == 0 and (dep.stdout or "0").strip() not in {"", "0"}
    _kubectl(["-n", target_ns, "get", "svc", "openbao"], env=env, timeout=20)
    return _to_bool(ok)