from __future__ import annotations import hashlib import json from pathlib import Path from typing import Iterable CNPG_PLACEMENT_SCHEMA_VERSION = "v1" def _coerce_positive_int(raw: object, default: int) -> int: try: val = int(str(raw).strip()) if val < 1: return default return val except Exception: return default def _normalize_nodes(nodes: Iterable[str] | None) -> list[str]: normalized = { str(node or "").strip() for node in (nodes or []) if str(node or "").strip() } return sorted(normalized) def _normalize_assignments(raw: object) -> dict[int, str]: if not isinstance(raw, dict): return {} out: dict[int, str] = {} for k, v in raw.items(): try: ordinal = int(str(k).strip()) except Exception: continue if ordinal < 0: continue node = str(v or "").strip() if not node: continue out[ordinal] = node return out def _assign_round_robin( *, desired_instances: int, eligible_nodes: list[str], base: dict[int, str] | None = None ) -> dict[int, str]: assignments = dict(base or {}) if desired_instances < 1 or not eligible_nodes: return {} for ordinal in range(desired_instances): if ordinal in assignments: continue assignments[ordinal] = eligible_nodes[ordinal % len(eligible_nodes)] return {k: assignments[k] for k in sorted(assignments.keys()) if k < desired_instances} def _plan_hash(payload: dict) -> str: canonical = json.dumps(payload, separators=(",", ":"), sort_keys=True) digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() return digest[:16] def _normalize_plan_for_hash( *, cluster_name: str, desired_instances: int, eligible_nodes: list[str], assignments: dict[int, str], ) -> dict: return { "schema_version": CNPG_PLACEMENT_SCHEMA_VERSION, "cluster_name": cluster_name, "desired_instances": desired_instances, "eligible_nodes": eligible_nodes, "assignments": {str(k): assignments[k] for k in sorted(assignments.keys())}, } def _prior_plan_matches_cluster(prior_plan: dict | None, cluster_name: str) -> bool: if not isinstance(prior_plan, dict): return False return str(prior_plan.get("cluster_name") or "").strip() == cluster_name def _prior_plan_is_eligible(prior_assignments: dict[int, str], eligible_nodes: list[str]) -> bool: if not prior_assignments: return False if not eligible_nodes: return False eligible_set = set(eligible_nodes) return all(node in eligible_set for node in prior_assignments.values()) def plan_cnpg_placement( *, cluster_name: str, desired_instances: int, candidate_nodes: Iterable[str] | None, prior_plan: dict | None = None, rebalance: bool = False, ) -> dict: cluster = str(cluster_name or "").strip() or "knoe-db" desired = _coerce_positive_int(desired_instances, 1) eligible_nodes = _normalize_nodes(candidate_nodes) prior_assignments = _normalize_assignments( (prior_plan or {}).get("assignments") if isinstance(prior_plan, dict) else {} ) prior_matches = _prior_plan_matches_cluster(prior_plan, cluster) prior_eligible = _prior_plan_is_eligible(prior_assignments, eligible_nodes) reused = False regenerated = True reason = "new_cluster" if rebalance: assignments = _assign_round_robin( desired_instances=desired, eligible_nodes=eligible_nodes, ) reason = "rebalance_requested" elif prior_matches and prior_eligible: # Reuse existing ordinal bindings and only extend/trim as needed. assignments = _assign_round_robin( desired_instances=desired, eligible_nodes=eligible_nodes, base=prior_assignments, ) reused = True regenerated = False reason = "reused" elif prior_matches and prior_assignments and not prior_eligible: assignments = _assign_round_robin( desired_instances=desired, eligible_nodes=eligible_nodes, ) reason = "assigned_node_no_longer_eligible" else: assignments = _assign_round_robin( desired_instances=desired, eligible_nodes=eligible_nodes, ) payload = _normalize_plan_for_hash( cluster_name=cluster, desired_instances=desired, eligible_nodes=eligible_nodes, assignments=assignments, ) plan_hash = _plan_hash(payload) return { **payload, "plan_id": f"cnpg-placement-{plan_hash}", "plan_hash": plan_hash, "metadata": { "reused": reused, "regenerated": regenerated, "reason": reason, "prior_plan_present": isinstance(prior_plan, dict), }, } def load_cnpg_placement_plan(path: Path) -> dict | None: try: if not path.exists(): return None payload = json.loads(path.read_text()) except Exception: return None if not isinstance(payload, dict): return None return payload def save_cnpg_placement_plan(path: Path, plan: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n")