from __future__ import annotations import argparse import json import math 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 TARGET_ROLES = ("tikv", "tiflash") class MetricsQueryError(RuntimeError): pass @dataclass(frozen=True) class StoreRoleMapping: store: str pod: str component: str address: 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 _get(self, path: str, *, params: dict[str, str]) -> dict[str, Any]: url = f"{self._base_url}{path}?{urllib.parse.urlencode(params)}" 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 _safe_float(value: Any) -> float | None: try: numeric = float(value) except (TypeError, ValueError): return None if not math.isfinite(numeric): return None return numeric 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 _pd_store_address_query() -> str: return "count(pd_scheduler_store_status) by (address, store)" def _align_snapshot_time(*, now_ts: int, step_seconds: int, lag_seconds: int) -> int: effective_now = now_ts - lag_seconds if effective_now <= 0: raise MetricsQueryError("effective current time must be positive") snapshot_ts = (effective_now // step_seconds) * step_seconds if snapshot_ts <= 0: raise MetricsQueryError("resolved snapshot time must be positive") return snapshot_ts def _component_for_pod(pod: str) -> str: if "-tikv-" in pod: return "tikv" if "-compute-tiflash-" in pod or "-tiflash-" in pod: return "tiflash" raise MetricsQueryError(f"unrecognized storage pod role for pod '{pod}'") def _store_role_mappings(payload: dict[str, Any]) -> list[StoreRoleMapping]: pod_to_mapping: dict[str, StoreRoleMapping] = {} store_to_pod: dict[str, str] = {} for item in _result_series(payload): metric = item.get("metric") if not isinstance(metric, dict): continue store = str(metric.get("store") or "").strip() address = str(metric.get("address") or "").strip() if not store or not address: continue pod = address.split(".", 1)[0].strip() mapping = StoreRoleMapping( store=store, pod=pod, component=_component_for_pod(pod), address=address, ) previous = pod_to_mapping.get(pod) if previous is not None and previous.store != mapping.store: raise MetricsQueryError( f"active pod '{pod}' maps to multiple stores: '{previous.store}' and '{mapping.store}'" ) other_pod = store_to_pod.get(store) if other_pod is not None and other_pod != pod: raise MetricsQueryError( f"store '{store}' maps to multiple active pods: '{other_pod}' and '{pod}'" ) pod_to_mapping[pod] = mapping store_to_pod[store] = pod return sorted( pod_to_mapping.values(), key=lambda item: (item.component, item.pod), ) 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_store_samples( mappings: list[StoreRoleMapping], *, stores_per_role: int, ) -> list[StoreRoleMapping]: selected_by_role: dict[str, list[StoreRoleMapping]] = {} for role in TARGET_ROLES: role_items = [item for item in mappings if item.component == role] if len(role_items) < stores_per_role: raise MetricsQueryError( f"expected at least {stores_per_role} active {role} stores, got {len(role_items)}" ) indices = _spread_indices(len(role_items), stores_per_role) selected_by_role[role] = [role_items[idx] for idx in indices] selected: list[StoreRoleMapping] = [] for idx in range(stores_per_role): for role in TARGET_ROLES: selected.append(selected_by_role[role][idx]) return selected 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, ) metrics_cluster_id = args.metrics_cluster_id or args.cluster_id snapshot_ts = _align_snapshot_time( now_ts=int(datetime.now(tz=timezone.utc).timestamp()), step_seconds=args.step_seconds, lag_seconds=args.snapshot_lag_seconds, ) pd_payload = client.query_instant( cluster_id=metrics_cluster_id, promql=_pd_store_address_query(), query_time=snapshot_ts, ) mappings = _store_role_mappings(pd_payload) selected = _select_store_samples(mappings, stores_per_role=args.stores_per_role) store_ids = [item.store 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, "start": _isoformat_utc(snapshot_ts - args.step_seconds), "end": _isoformat_utc(snapshot_ts), "snapshot_time": _isoformat_utc(snapshot_ts), "step_seconds": args.step_seconds, "stores_per_role": args.stores_per_role, "store_ids": store_ids, "store_ids_csv": ",".join(store_ids), "store_ids_json": json.dumps(store_ids, ensure_ascii=False), "selected_stores": [ { "store": item.store, "component": item.component, "pod": item.pod, "address": item.address, } for item in selected ], "inventory_summary": { "tikv": sum(1 for item in mappings if item.component == "tikv"), "tiflash": sum(1 for item in mappings if item.component == "tiflash"), }, "queries": { "pd_store_address_query": _pd_store_address_query(), }, } def _build_reference(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, ) metrics_cluster_id = args.metrics_cluster_id or args.cluster_id snapshot_ts = _parse_iso8601(args.snapshot_time) store_ids = [item.strip() for item in str(args.stores_csv or "").split(",") if item.strip()] if not store_ids: raise MetricsQueryError("resolved empty selected store list") pd_payload = client.query_instant( cluster_id=metrics_cluster_id, promql=_pd_store_address_query(), query_time=snapshot_ts, ) mappings = _store_role_mappings(pd_payload) by_store = {item.store: item for item in mappings} stores: list[dict[str, str]] = [] for store in store_ids: mapping = by_store.get(store) if mapping is None: raise MetricsQueryError(f"selected store '{store}' was not found in the active store inventory") stores.append( { "store": mapping.store, "component": mapping.component, } ) return { "cluster_id": args.cluster_id, "evidence_sufficient": True, "stores": stores, "snapshot_time": _isoformat_utc(snapshot_ts), } def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Runtime hook for store role classification case") subparsers = parser.add_subparsers(dest="command", required=True) input_parser = subparsers.add_parser("prepare-case-input") input_parser.add_argument("--cluster-id", required=True) input_parser.add_argument("--metrics-cluster-id") input_parser.add_argument("--snapshot-lag-seconds", type=int, default=600) input_parser.add_argument("--step-seconds", type=int, default=300) input_parser.add_argument("--stores-per-role", type=int, default=3) reference_parser = subparsers.add_parser("build-reference") reference_parser.add_argument("--cluster-id", required=True) reference_parser.add_argument("--metrics-cluster-id") reference_parser.add_argument("--snapshot-time", required=True) reference_parser.add_argument("--stores-csv", required=True) return parser def main(argv: list[str]) -> int: parser = _build_parser() args = parser.parse_args(argv) if args.command == "prepare-case-input": payload = _build_case_input(args) elif args.command == "build-reference": payload = _build_reference(args) else: raise MetricsQueryError(f"unsupported command: {args.command}") sys.stdout.write(json.dumps(payload, ensure_ascii=False)) sys.stdout.write("\n") return 0 if __name__ == "__main__": raise SystemExit(main(sys.argv[1:]))