Mutual_Fund/nl2sql/metrics.py

"""NL2SQL 查询运行指标，供管理员诊断和后续监控采集。"""
from __future__ import annotations

from collections import Counter
from threading import Lock
import logging
from datetime import datetime

from sqlalchemy import text


logger = logging.getLogger("nl2sql.metrics")

HISTORY_METRICS_SQL = text(
    """
    SELECT
        COUNT(*) AS total,
        COALESCE(SUM(status = 'success'), 0) AS success,
        COALESCE(SUM(status = 'failed'), 0) AS failed,
        COALESCE(SUM(status = 'timeout'), 0) AS timeout,
        COALESCE(AVG(elapsed_ms), 0) AS average_elapsed_ms
    FROM nl2sql_query_history
    WHERE (:since IS NULL OR create_time >= :since)
    """
)


class QueryMetrics:
    """进程内查询计数器，重启后清零。"""

    def __init__(self, *, slow_threshold_ms: float = 500):
        self.slow_threshold_ms = slow_threshold_ms
        self._lock = Lock()
        self._total = 0
        self._success = 0
        self._failed = 0
        self._cache_hits = 0
        self._timeouts = 0
        self._slow = 0
        self._rate_limited = 0
        self._elapsed_total = 0.0
        self._failure_reasons: Counter[str] = Counter()

    def record(
        self,
        *,
        status: str,
        elapsed_ms: float | None = None,
        cache_hit: bool = False,
        failure_reason: str | None = None,
    ) -> None:
        with self._lock:
            self._total += 1
            if status == "success":
                self._success += 1
            else:
                self._failed += 1
                if failure_reason:
                    self._failure_reasons[failure_reason] += 1
            if status == "timeout":
                self._timeouts += 1
            if cache_hit:
                self._cache_hits += 1
            if elapsed_ms is not None and elapsed_ms >= self.slow_threshold_ms:
                self._slow += 1
            if elapsed_ms is not None:
                self._elapsed_total += elapsed_ms

    def record_rate_limited(self) -> None:
        """记录一次因配额、并发或频率限制而拒绝的请求。"""
        with self._lock:
            self._rate_limited += 1

    def snapshot(self) -> dict[str, float | int]:
        with self._lock:
            total = self._total
            return {
                "total": total,
                "success": self._success,
                "failed": self._failed,
                "success_rate": self._success / total if total else 0.0,
                "failure_rate": self._failed / total if total else 0.0,
                "cache_hits": self._cache_hits,
                "cache_hit_rate": self._cache_hits / total if total else 0.0,
                "timeout_count": self._timeouts,
                "slow_query_count": self._slow,
                "rate_limited_count": self._rate_limited,
                "average_elapsed_ms": self._elapsed_total / total if total else 0.0,
                "failure_reasons": dict(sorted(self._failure_reasons.items())),
            }


query_metrics = QueryMetrics()


def render_prometheus(metrics: dict | None = None) -> str:
    """将聚合指标导出为无外部依赖的 Prometheus 文本格式。"""
    snapshot = metrics or query_metrics.snapshot()
    lines = [
        "# HELP nl2sql_queries_total NL2SQL 查询总数",
        "# TYPE nl2sql_queries_total counter",
        f"nl2sql_queries_total {snapshot.get('total', 0)}",
        "# HELP nl2sql_queries_success_total NL2SQL 成功查询数",
        "# TYPE nl2sql_queries_success_total counter",
        f"nl2sql_queries_success_total {snapshot.get('success', 0)}",
        "# HELP nl2sql_queries_failed_total NL2SQL 失败查询数",
        "# TYPE nl2sql_queries_failed_total counter",
        f"nl2sql_queries_failed_total {snapshot.get('failed', 0)}",
        "# HELP nl2sql_queries_timeout_total NL2SQL 超时查询数",
        "# TYPE nl2sql_queries_timeout_total counter",
        f"nl2sql_queries_timeout_total {snapshot.get('timeout_count', 0)}",
        "# HELP nl2sql_queries_cache_hits_total NL2SQL 缓存命中数",
        "# TYPE nl2sql_queries_cache_hits_total counter",
        f"nl2sql_queries_cache_hits_total {snapshot.get('cache_hits', 0)}",
        "# HELP nl2sql_queries_rate_limited_total NL2SQL 限流拒绝数",
        "# TYPE nl2sql_queries_rate_limited_total counter",
        f"nl2sql_queries_rate_limited_total {snapshot.get('rate_limited_count', 0)}",
        "# HELP nl2sql_queries_slow_total NL2SQL 慢查询数",
        "# TYPE nl2sql_queries_slow_total counter",
        f"nl2sql_queries_slow_total {snapshot.get('slow_query_count', 0)}",
        "# HELP nl2sql_queries_average_elapsed_ms NL2SQL 平均耗时毫秒",
        "# TYPE nl2sql_queries_average_elapsed_ms gauge",
        f"nl2sql_queries_average_elapsed_ms {snapshot.get('average_elapsed_ms', 0.0)}",
    ]
    for reason, count in sorted((snapshot.get("failure_reasons") or {}).items()):
        safe_reason = str(reason).replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
        lines.append(f'nl2sql_query_failures_total{{reason="{safe_reason}"}} {count}')
    return "\n".join(lines) + "\n"


async def load_history_metrics(db, *, since: datetime | None = None) -> dict[str, int | float]:
    """从查询归档表聚合指标，不读取问题、SQL 或结果行。"""
    result = await db.execute(HISTORY_METRICS_SQL, {"since": since})
    row = result.mappings().one()
    return {
        "total": int(row.get("total", 0) or 0),
        "success": int(row.get("success", 0) or 0),
        "failed": int(row.get("failed", 0) or 0),
        "timeout": int(row.get("timeout", 0) or 0),
        "average_elapsed_ms": float(row.get("average_elapsed_ms", 0) or 0),
    }


async def load_history_metrics_safely(db, *, since: datetime | None = None) -> dict[str, int | float | bool]:
    """历史指标读取失败时返回明确的不可用状态。"""
    try:
        result = await load_history_metrics(db, since=since)
        return {**result, "available": True}
    except Exception:  # noqa: BLE001 指标故障不能阻断管理接口
        logger.warning("NL2SQL 历史指标读取失败", exc_info=True)
        return {
            "total": 0,
            "success": 0,
            "failed": 0,
            "timeout": 0,
            "average_elapsed_ms": 0.0,
            "available": False,
        }