prole/knoe/tools/cleanup_cnpg_storage.py

from __future__ import annotations

import argparse
from dataclasses import dataclass, field
from typing import Any


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_cli_args(argv: list[str] | None = None) -> argparse.Namespace:
    # shared args for all subcommands
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--dry-run", action="store_true", default=True)
    common.add_argument("--database-namespace", required=True)
    common.add_argument("--cluster-name", default=None)

    # separate parent without --cluster-name for cleanup-cluster (which makes it required)
    common_no_cluster = argparse.ArgumentParser(add_help=False)
    common_no_cluster.add_argument("--dry-run", action="store_true", default=True)
    common_no_cluster.add_argument("--database-namespace", required=True)

    parser = argparse.ArgumentParser(description="CNPG storage cleanup tool")
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("inspect", parents=[common])

    cleanup_cluster = sub.add_parser("cleanup-cluster", parents=[common_no_cluster])
    cleanup_cluster.add_argument("--cluster-name", required=True, dest="cluster_name")

    ns = parser.parse_args(argv)
    if not hasattr(ns, "dry_run") or ns.dry_run is None:
        ns.dry_run = True
    if not hasattr(ns, "cluster_name"):
        ns.cluster_name = None
    return ns


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass
class ResourceItem:
    name: str
    raw: dict
    reasons: list[str] = field(default_factory=list)


@dataclass
class ResourceSnapshot:
    clusters: list[dict] = field(default_factory=list)
    pvcs: list[dict] = field(default_factory=list)
    pvs: list[dict] = field(default_factory=list)
    pods: list[dict] = field(default_factory=list)
    jobs: list[dict] = field(default_factory=list)


@dataclass
class SelectedResources:
    clusters: list[ResourceItem] = field(default_factory=list)
    pvcs: list[ResourceItem] = field(default_factory=list)
    pvs: list[ResourceItem] = field(default_factory=list)
    pods: list[ResourceItem] = field(default_factory=list)
    jobs: list[ResourceItem] = field(default_factory=list)


@dataclass
class MatchFilters:
    namespace: str
    cluster_name: str | None = None


# ---------------------------------------------------------------------------
# Selection logic
# ---------------------------------------------------------------------------

def _meta_name(obj: dict) -> str:
    return obj.get("metadata", {}).get("name", "")


def _meta_namespace(obj: dict) -> str:
    return obj.get("metadata", {}).get("namespace", "")


def _meta_labels(obj: dict) -> dict:
    return obj.get("metadata", {}).get("labels", {})


def select_resources(snapshot: ResourceSnapshot, filters: MatchFilters) -> SelectedResources:
    result = SelectedResources()

    # clusters
    for c in snapshot.clusters:
        ns = _meta_namespace(c)
        if ns == filters.namespace:
            reasons = [f"metadata.namespace={ns}"]
            if filters.cluster_name and _meta_name(c) == filters.cluster_name:
                reasons.append(f"metadata.name={filters.cluster_name}")
            result.clusters.append(ResourceItem(name=_meta_name(c), raw=c, reasons=reasons))

    # pvcs
    for pvc in snapshot.pvcs:
        ns = _meta_namespace(pvc)
        labels = _meta_labels(pvc)
        reasons: list[str] = []
        if ns == filters.namespace:
            reasons.append(f"metadata.namespace={ns}")
        if filters.cluster_name and labels.get("cnpg.io/cluster") == filters.cluster_name:
            reasons.append(f"label cnpg.io/cluster={filters.cluster_name}")
        if reasons:
            result.pvcs.append(ResourceItem(name=_meta_name(pvc), raw=pvc, reasons=reasons))

    # pvs
    for pv in snapshot.pvs:
        labels = _meta_labels(pv)
        spec = pv.get("spec", {})
        claim_ref = spec.get("claimRef", {})
        reasons: list[str] = []
        if filters.namespace and claim_ref.get("namespace") == filters.namespace:
            reasons.append(f"spec.claimRef.namespace={filters.namespace}")
        if filters.cluster_name and labels.get("knoe.io/cluster") == filters.cluster_name:
            reasons.append(f"label knoe.io/cluster={filters.cluster_name}")
        if reasons:
            result.pvs.append(ResourceItem(name=_meta_name(pv), raw=pv, reasons=reasons))

    # pods
    for pod in snapshot.pods:
        ns = _meta_namespace(pod)
        labels = _meta_labels(pod)
        reasons: list[str] = []
        if ns == filters.namespace:
            reasons.append(f"metadata.namespace={ns}")
        if filters.cluster_name and labels.get("cnpg.io/cluster") == filters.cluster_name:
            reasons.append(f"label cnpg.io/cluster={filters.cluster_name}")
        if reasons:
            result.pods.append(ResourceItem(name=_meta_name(pod), raw=pod, reasons=reasons))

    # jobs
    for job in snapshot.jobs:
        ns = _meta_namespace(job)
        labels = _meta_labels(job)
        reasons: list[str] = []
        if ns == filters.namespace:
            reasons.append(f"metadata.namespace={ns}")
        if filters.cluster_name and labels.get("cnpg.io/cluster") == filters.cluster_name:
            reasons.append(f"label cnpg.io/cluster={filters.cluster_name}")
        if reasons:
            result.jobs.append(ResourceItem(name=_meta_name(job), raw=job, reasons=reasons))

    return result


# ---------------------------------------------------------------------------
# PV helpers
# ---------------------------------------------------------------------------

def _should_patch_pv_claim_ref(pv: dict, *, delete_pv: bool) -> bool:
    if delete_pv:
        return False
    phase = pv.get("status", {}).get("phase", "")
    return phase == "Released" and bool(pv.get("spec", {}).get("claimRef"))


def _should_delete_pv(pv: dict, *, delete_pv: bool, force: bool) -> bool:
    if not delete_pv:
        return False
    phase = pv.get("status", {}).get("phase", "")
    if phase != "Released" and not force:
        return False
    return True