from __future__ import annotations import argparse import json import math import os import sys import urllib.error import urllib.request from datetime import UTC from datetime import datetime from datetime import timedelta from typing import Any from zoneinfo import ZoneInfo DEFAULT_BASE_URL = "https://clinic.pingcap.com" DEFAULT_TIMEOUT_S = 60 BEIJING_TZ = ZoneInfo("Asia/Shanghai") class SlowQueryReferenceError(RuntimeError): pass class ClinicDataProxyClient: 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_sql(self, *, cluster_id: str, sql: str, timeout_s: int = DEFAULT_TIMEOUT_S) -> dict[str, Any]: url = f"{self._base_url}/data-proxy/query" payload = json.dumps( { "clusterId": cluster_id, "timeout": timeout_s, "sql": sql, } ).encode("utf-8") request = urllib.request.Request( url, data=payload, headers={ "Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(request, timeout=self._timeout_s) as response: body = response.read().decode("utf-8") except urllib.error.HTTPError as exc: text = exc.read().decode("utf-8", errors="replace") raise SlowQueryReferenceError(f"HTTP {exc.code} for /data-proxy/query: {text}") from exc except urllib.error.URLError as exc: raise SlowQueryReferenceError(f"request failed for /data-proxy/query: {exc}") from exc try: data = json.loads(body) except json.JSONDecodeError as exc: raise SlowQueryReferenceError(f"invalid JSON response for /data-proxy/query: {exc}") from exc if isinstance(data, dict) and data.get("rows") is None: data["rows"] = [] error_text = str(data.get("error") or "").strip() if isinstance(data, dict) else "" if error_text: raise SlowQueryReferenceError(error_text) 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 SlowQueryReferenceError( "missing Clinic API key in SKILL_CLINIC_API_KEY/CLINIC_API_KEY/INVESTIBOOK_CLINIC_API_KEY" ) 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 _isoformat_utc(ts: int) -> str: return datetime.fromtimestamp(ts, tz=UTC).isoformat().replace("+00:00", "Z") 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 _safe_int(value: Any) -> int | None: numeric = _safe_float(value) if numeric is None: return None return int(numeric) def _sql_string_literal(value: str) -> str: return value.replace("'", "''") def _compact_sql(value: str, *, limit: int = 200) -> str: compacted = " ".join(str(value or "").strip().split()) return compacted[:limit] def _partition_literals(start_ts: int, end_ts: int) -> list[str]: start_date = datetime.fromtimestamp(start_ts, tz=UTC).date() last_ts = max(start_ts, end_ts - 1) end_date = datetime.fromtimestamp(last_ts, tz=UTC).date() current = start_date partitions: list[str] = [] while current <= end_date: partitions.append(current.strftime("%Y%m%d")) current += timedelta(days=1) return partitions def _partition_predicate(start_ts: int, end_ts: int) -> str: partitions = _partition_literals(start_ts, end_ts) if len(partitions) == 1: return f"date = '{partitions[0]}'" values = ", ".join(f"'{part}'" for part in partitions) return f"date IN ({values})" def _rows_as_dicts(payload: dict[str, Any]) -> list[dict[str, Any]]: columns = payload.get("columns") rows = payload.get("rows") if not isinstance(columns, list) or not isinstance(rows, list): raise SlowQueryReferenceError("SQL response is missing columns or rows") column_names = [str(item) for item in columns] normalized_rows: list[dict[str, Any]] = [] for row in rows: if not isinstance(row, list): continue normalized_rows.append(dict(zip(column_names, row))) return normalized_rows def _latest_completed_beijing_window( *, local_hour: int, local_minute: int, duration_minutes: int, ) -> tuple[int, int, str]: now_local = datetime.now(tz=BEIJING_TZ) candidate_start = now_local.replace( hour=local_hour, minute=local_minute, second=0, microsecond=0, ) candidate_end = candidate_start + timedelta(minutes=duration_minutes) if now_local < candidate_end: candidate_start = candidate_start - timedelta(days=1) candidate_end = candidate_start + timedelta(minutes=duration_minutes) return ( int(candidate_start.astimezone(UTC).timestamp()), int(candidate_end.astimezone(UTC).timestamp()), candidate_start.date().isoformat(), ) def _build_top_digest_sql(*, start_ts: int, end_ts: int, limit: int) -> str: return ( "SELECT LOWER(digest) AS sql_digest, MIN(query) AS sample_sql, COUNT(*) AS occurrence_count " "FROM clinic_data_proxy.slow_query_logs " f"WHERE {_partition_predicate(start_ts, end_ts)} " f"AND CAST(time AS DOUBLE) >= {start_ts} AND CAST(time AS DOUBLE) < {end_ts} " "GROUP BY LOWER(digest) " "ORDER BY occurrence_count DESC, sql_digest ASC " f"LIMIT {int(limit)}" ) def _top_rows(payload: dict[str, Any]) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for row in _rows_as_dicts(payload): sql_digest = str(row.get("sql_digest") or "").strip().lower() occurrence_count = _safe_int(row.get("occurrence_count")) if not sql_digest or occurrence_count is None: continue rows.append( { "sql_digest": sql_digest, "sql_statement": _compact_sql(str(row.get("sample_sql") or "")), "occurrence_count": occurrence_count, } ) rows.sort(key=lambda item: (-int(item["occurrence_count"]), str(item["sql_digest"]))) return rows def _resolve_window(args: argparse.Namespace) -> dict[str, Any]: current_start, current_end, local_date = _latest_completed_beijing_window( local_hour=args.local_hour, local_minute=args.local_minute, duration_minutes=args.duration_minutes, ) previous_start = current_start - args.duration_minutes * 60 previous_end = current_start return { "id": f"bjt-{local_date}-{args.local_hour:02d}{args.local_minute:02d}-{args.duration_minutes}m", "cluster_id": args.cluster_id, "cluster_label": args.cluster_label, "local_timezone": "Asia/Shanghai", "local_date": local_date, "duration_minutes": args.duration_minutes, "top_n": args.top_n, "previous_window": { "start": _isoformat_utc(previous_start), "end": _isoformat_utc(previous_end), "start_ts": previous_start, "end_ts": previous_end, }, "current_window": { "start": _isoformat_utc(current_start), "end": _isoformat_utc(current_end), "start_ts": current_start, "end_ts": current_end, }, } def _build_reference(args: argparse.Namespace) -> dict[str, Any]: current_start_ts = _parse_iso8601(args.start) current_end_ts = _parse_iso8601(args.end) if current_end_ts <= current_start_ts: raise SlowQueryReferenceError("reference window end must be greater than start") window_seconds = current_end_ts - current_start_ts previous_start_ts = current_start_ts - window_seconds previous_end_ts = current_start_ts client = ClinicDataProxyClient( api_key=_env_api_key(), base_url=os.environ.get("CLINIC_ENDPOINT", DEFAULT_BASE_URL).strip() or DEFAULT_BASE_URL, ) current_sql = _build_top_digest_sql( start_ts=current_start_ts, end_ts=current_end_ts, limit=args.top_n, ) previous_sql = _build_top_digest_sql( start_ts=previous_start_ts, end_ts=previous_end_ts, limit=args.top_n, ) current_top = _top_rows(client.query_sql(cluster_id=args.cluster_id, sql=current_sql)) previous_top = _top_rows(client.query_sql(cluster_id=args.cluster_id, sql=previous_sql)) evidence_sufficient = len(current_top) >= args.top_n and len(previous_top) >= args.top_n if not evidence_sufficient: return { "cluster_id": args.cluster_id, "comparison": { "previous_window": { "start": _isoformat_utc(previous_start_ts), "end": _isoformat_utc(previous_end_ts), }, "current_window": { "start": _isoformat_utc(current_start_ts), "end": _isoformat_utc(current_end_ts), }, }, "evidence_sufficient": False, "new_sql_digests": [], "debug": { "current_query_sql": current_sql, "previous_query_sql": previous_sql, "current_top_count": len(current_top), "previous_top_count": len(previous_top), }, } previous_digests = {str(item["sql_digest"]) for item in previous_top} new_sql_digests = [ { "sql_digest": str(item["sql_digest"]), "sql_statement": str(item["sql_statement"]), "occurrence_count": int(item["occurrence_count"]), } for item in current_top if str(item["sql_digest"]) not in previous_digests ] new_sql_digests.sort( key=lambda item: (-int(item["occurrence_count"]), str(item["sql_digest"])) ) return { "cluster_id": args.cluster_id, "comparison": { "previous_window": { "start": _isoformat_utc(previous_start_ts), "end": _isoformat_utc(previous_end_ts), }, "current_window": { "start": _isoformat_utc(current_start_ts), "end": _isoformat_utc(current_end_ts), }, }, "evidence_sufficient": True, "new_sql_digests": new_sql_digests, "debug": { "current_query_sql": current_sql, "previous_query_sql": previous_sql, "current_top_count": len(current_top), "previous_top_count": len(previous_top), }, } def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Runtime hook for adjacent slow-query top10 comparison in saas01" ) subparsers = parser.add_subparsers(dest="command", required=True) resolve_parser = subparsers.add_parser("resolve-window") resolve_parser.add_argument("--cluster-id", required=True) resolve_parser.add_argument("--cluster-label", required=True) resolve_parser.add_argument("--local-hour", type=int, default=11) resolve_parser.add_argument("--local-minute", type=int, default=0) resolve_parser.add_argument("--duration-minutes", type=int, default=10) resolve_parser.add_argument("--top-n", type=int, default=10) reference_parser = subparsers.add_parser("build-reference") reference_parser.add_argument("--cluster-id", required=True) reference_parser.add_argument("--cluster-label", required=True) reference_parser.add_argument("--start", required=True) reference_parser.add_argument("--end", required=True) reference_parser.add_argument("--top-n", type=int, default=10) return parser def main(argv: list[str]) -> int: parser = _build_parser() args = parser.parse_args(argv) if args.command == "resolve-window": payload = _resolve_window(args) elif args.command == "build-reference": payload = _build_reference(args) else: raise SlowQueryReferenceError(f"unsupported command: {args.command}") sys.stdout.write(json.dumps(payload, ensure_ascii=False, indent=2)) sys.stdout.write("\n") return 0 if __name__ == "__main__": raise SystemExit(main(sys.argv[1:]))