from __future__ import annotations import argparse import json import os import sys import urllib.error import urllib.parse import urllib.request from dataclasses import dataclass from datetime import datetime from datetime import timezone from typing import Any DEFAULT_BASE_URL = "https://clinic.pingcap.com" DEFAULT_TIMEOUT_S = 60 class MetricsQueryError(RuntimeError): pass @dataclass(frozen=True) class PodNodeMapping: tikv_pod: str worker_node: str class ClinicMetricsClient: def __init__(self, *, api_key: str, base_url: str, timeout_s: int = DEFAULT_TIMEOUT_S) -> None: self._api_key = api_key self._base_url = base_url.rstrip("/") self._timeout_s = timeout_s def query_instant( self, *, cluster_id: str, promql: str, query_time: int, ) -> dict[str, Any]: return self._get( f"/data-proxy/metrics/{cluster_id}/api/v1/query", params={ "query": promql, "time": str(query_time), }, ) def query_series( self, *, cluster_id: str, match: str, start: int, end: int, ) -> dict[str, Any]: return self._get( f"/data-proxy/metrics/{cluster_id}/api/v1/series", params={ "match[]": [match], "start": str(start), "end": str(end), }, ) def _get(self, path: str, *, params: dict[str, Any]) -> dict[str, Any]: url = f"{self._base_url}{path}?{urllib.parse.urlencode(params, doseq=True)}" request = urllib.request.Request( url, headers={"Authorization": f"Bearer {self._api_key}"}, method="GET", ) try: with urllib.request.urlopen(request, timeout=self._timeout_s) as response: payload = response.read().decode("utf-8") except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise MetricsQueryError(f"HTTP {exc.code} for {path}: {body}") from exc except urllib.error.URLError as exc: raise MetricsQueryError(f"request failed for {path}: {exc}") from exc try: data = json.loads(payload) except json.JSONDecodeError as exc: raise MetricsQueryError(f"invalid JSON response for {path}: {exc}") from exc status = str(data.get("status") or "") if status and status != "success": raise MetricsQueryError(f"unexpected response status for {path}: {status}") return data def _env_api_key() -> str: for key in ("SKILL_CLINIC_API_KEY", "CLINIC_API_KEY", "INVESTIBOOK_CLINIC_API_KEY"): value = os.environ.get(key, "").strip() if value: return value raise MetricsQueryError( "missing Clinic API key in SKILL_CLINIC_API_KEY/CLINIC_API_KEY/INVESTIBOOK_CLINIC_API_KEY" ) def _isoformat_utc(ts: int) -> str: return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat().replace("+00:00", "Z") def _parse_iso8601(value: str) -> int: normalized = value.strip() if normalized.endswith("Z"): normalized = normalized[:-1] + "+00:00" return int(datetime.fromisoformat(normalized).timestamp()) def _namespace_for_cluster(cluster_id: str) -> str: return f"tidb{cluster_id}" def _inventory_match(cluster_id: str) -> str: namespace = _namespace_for_cluster(cluster_id) return f'kube_pod_info{{namespace="{namespace}",pod=~"db-tikv-.*"}}' def _align_snapshot_time(*, now_ts: int, lag_seconds: int) -> int: snapshot_ts = now_ts - lag_seconds if snapshot_ts <= 0: raise MetricsQueryError("resolved snapshot time must be positive") return snapshot_ts def _series_items(payload: dict[str, Any]) -> list[dict[str, Any]]: data = payload.get("data") return [item for item in data if isinstance(item, dict)] if isinstance(data, list) else [] def _result_series(payload: dict[str, Any]) -> list[dict[str, Any]]: data = payload.get("data") if not isinstance(data, dict): return [] result = data.get("result") return [item for item in result if isinstance(item, dict)] if isinstance(result, list) else [] def _pod_sort_key(pod: str) -> tuple[int, str]: try: suffix = int(pod.rsplit("-", 1)[1]) except (IndexError, ValueError): return (sys.maxsize, pod) return (suffix, pod) def _inventory_mappings(payload: dict[str, Any]) -> list[PodNodeMapping]: resolved: dict[str, str] = {} for item in _result_series(payload): metric = item.get("metric") if not isinstance(metric, dict): continue tikv_pod = str(metric.get("pod") or "").strip() worker_node = str(metric.get("node") or "").strip() if not tikv_pod or not worker_node: continue previous = resolved.get(tikv_pod) if previous is None: resolved[tikv_pod] = worker_node continue if previous != worker_node: raise MetricsQueryError( f"pod '{tikv_pod}' maps to multiple worker nodes: '{previous}' and '{worker_node}'" ) return [ PodNodeMapping(tikv_pod=tikv_pod, worker_node=worker_node) for tikv_pod, worker_node in sorted(resolved.items(), key=lambda item: _pod_sort_key(item[0])) ] def _spread_indices(length: int, count: int) -> list[int]: if length <= 0 or count <= 0: return [] if count >= length: return list(range(length)) raw = [round(idx * (length - 1) / (count - 1)) if count > 1 else 0 for idx in range(count)] indices: list[int] = [] seen: set[int] = set() for value in raw: bounded = min(max(int(value), 0), length - 1) if bounded in seen: continue seen.add(bounded) indices.append(bounded) cursor = 0 while len(indices) < count and cursor < length: if cursor not in seen: seen.add(cursor) indices.append(cursor) cursor += 1 return sorted(indices) def _select_target_pods(mappings: list[PodNodeMapping], *, pod_count: int) -> list[PodNodeMapping]: if len(mappings) < pod_count: raise MetricsQueryError(f"expected at least {pod_count} active TiKV pods, got {len(mappings)}") indices = _spread_indices(len(mappings), pod_count) return [mappings[idx] for idx in indices] def _load_snapshot_inventory( client: ClinicMetricsClient, *, cluster_id: str, snapshot_ts: int, ) -> list[PodNodeMapping]: payload = client.query_instant( cluster_id=cluster_id, promql=_inventory_match(cluster_id), query_time=snapshot_ts, ) mappings = _inventory_mappings(payload) if not mappings: raise MetricsQueryError(f"failed to resolve TiKV pod inventory for cluster {cluster_id}") return mappings def _build_case_input(args: argparse.Namespace) -> dict[str, Any]: client = ClinicMetricsClient( api_key=_env_api_key(), base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL, ) snapshot_ts = _align_snapshot_time( now_ts=int(datetime.now(tz=timezone.utc).timestamp()), lag_seconds=args.snapshot_lag_seconds, ) mappings = _load_snapshot_inventory( client, cluster_id=args.cluster_id, snapshot_ts=snapshot_ts, ) selected = _select_target_pods(mappings, pod_count=args.pod_count) pods = [item.tikv_pod for item in selected] return { "id": f"snapshot-{datetime.fromtimestamp(snapshot_ts, tz=timezone.utc).strftime('%Y%m%dT%H%M%SZ')}", "cluster_id": args.cluster_id, "cluster_label": args.cluster_label, "start": _isoformat_utc(max(snapshot_ts - args.window_seconds, 0)), "end": _isoformat_utc(snapshot_ts), "snapshot_time": _isoformat_utc(snapshot_ts), "window_seconds": args.window_seconds, "pod_count": args.pod_count, "pods": pods, "pods_csv": ",".join(pods), "pods_json": json.dumps(pods, ensure_ascii=False), } def _build_reference(args: argparse.Namespace) -> dict[str, Any]: snapshot_ts = _parse_iso8601(args.snapshot_time) target_pods = [item.strip() for item in str(args.pods_csv or "").split(",") if item.strip()] if not target_pods: raise MetricsQueryError("pods-csv must not be empty") client = ClinicMetricsClient( api_key=_env_api_key(), base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL, ) mappings = _load_snapshot_inventory( client, cluster_id=args.cluster_id, snapshot_ts=snapshot_ts, ) mapping_by_pod = {item.tikv_pod: item.worker_node for item in mappings} resolved: list[dict[str, str]] = [] for pod in target_pods: worker_node = mapping_by_pod.get(pod) if not worker_node: return { "cluster_id": args.cluster_id, "cluster_label": args.cluster_label, "snapshot_time": _isoformat_utc(snapshot_ts), "evidence_sufficient": False, "pod_mappings": [], } resolved.append( { "tikv_pod": pod, "worker_node": worker_node, } ) return { "cluster_id": args.cluster_id, "cluster_label": args.cluster_label, "snapshot_time": _isoformat_utc(snapshot_ts), "evidence_sufficient": True, "pod_mappings": resolved, } def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Runtime hooks for saas01 TiKV pod to worker-node mapping cases." ) subparsers = parser.add_subparsers(dest="command", required=True) prepare_parser = subparsers.add_parser("prepare-case-input") prepare_parser.add_argument("--cluster-id", required=True) prepare_parser.add_argument("--cluster-label", required=True) prepare_parser.add_argument("--snapshot-lag-seconds", type=int, required=True) prepare_parser.add_argument("--window-seconds", type=int, required=True) prepare_parser.add_argument("--pod-count", type=int, default=4) build_parser = subparsers.add_parser("build-reference") build_parser.add_argument("--cluster-id", required=True) build_parser.add_argument("--cluster-label", required=True) build_parser.add_argument("--snapshot-time", required=True) build_parser.add_argument("--window-seconds", type=int, required=True) build_parser.add_argument("--pods-csv", required=True) return parser def main(argv: list[str] | None = None) -> int: parser = _build_parser() args = parser.parse_args(argv) if args.command == "prepare-case-input": payload = _build_case_input(args) json.dump(payload, fp=sys.stdout, ensure_ascii=False, indent=2) sys.stdout.write("\n") return 0 if args.command == "build-reference": payload = _build_reference(args) json.dump(payload, fp=sys.stdout, ensure_ascii=False, indent=2) sys.stdout.write("\n") return 0 parser.error(f"unsupported command: {args.command}") return 2 if __name__ == "__main__": raise SystemExit(main())