LokiStack/backend/services/log_query.go

package services

// 日志查询代理 — 后端代表前端访问 Loki HTTP API。
// cluster-svc 型经 kube-apiserver service proxy(不依赖 ClusterIP 直通与集群 DNS);
// url 型直连。契约参照 aiops 日志中心(query/labels/values/histogram + cursor 翻页)。
import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"k8s.io/client-go/rest"
	"lokistack/models"
	"lokistack/repository"
)

const maxLogQueryLimit = 1000
const logQueryTimeout = 120 * time.Second

var lokiHTTPClients = [2]*http.Client{
	{Timeout: logQueryTimeout, Transport: newLokiTransport(false)},
	{Timeout: logQueryTimeout, Transport: newLokiTransport(true)},
}

func newLokiTransport(skipTLSVerify bool) *http.Transport {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: skipTLSVerify} //nolint:gosec // 数据源显式配置
	return transport
}

// ==================== Loki HTTP 访问 ====================

// lokiProxyGET 对数据源执行 GET；cluster-svc 使用 client-go 的认证传输访问 API Server Service Proxy。
func lokiProxyGET(ctx context.Context, src *models.LogSource, path string, query url.Values) ([]byte, int, error) {
	if src.TargetType == "cluster-svc" {
		db := repository.GetDB()
		db.RLock()
		var cluster *models.Cluster
		for i := range db.Clusters {
			if src.ClusterID != nil && db.Clusters[i].ID == *src.ClusterID {
				c := db.Clusters[i]
				cluster = &c
				break
			}
		}
		db.RUnlock()
		if cluster == nil {
			return nil, 500, fmt.Errorf("数据源关联的集群不存在")
		}
		if !kubernetesNamespacePattern.MatchString(src.Namespace) || !kubernetesNamespacePattern.MatchString(src.Service) || src.Port < 1 || src.Port > 65535 {
			return nil, 500, fmt.Errorf("Loki 数据源的命名空间、Service 或端口无效")
		}
		config, err := clusterRESTConfig(cluster, logQueryTimeout)
		if err != nil {
			return nil, 500, err
		}
		client, err := rest.HTTPClientFor(config)
		if err != nil {
			return nil, 500, fmt.Errorf("创建 Kubernetes API 客户端失败: %w", err)
		}
		full := strings.TrimRight(config.Host, "/") + fmt.Sprintf("/api/v1/namespaces/%s/services/%s:%d/proxy%s",
			src.Namespace, src.Service, src.Port, path)
		if qs := query.Encode(); qs != "" {
			full += "?" + qs
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, full, nil)
		if err != nil {
			return nil, 500, err
		}
		resp, err := client.Do(req)
		if err != nil {
			return nil, 502, fmt.Errorf("Kubernetes API 代理查询失败: %w", err)
		}
		defer resp.Body.Close()
		body, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
		if err != nil {
			return nil, 502, err
		}
		if resp.StatusCode >= 400 {
			return body, resp.StatusCode, fmt.Errorf("Kubernetes API 代理返回 HTTP %d: %s", resp.StatusCode, truncateStr(string(body), 300))
		}
		return body, resp.StatusCode, nil
	}

	// url 直连
	if src.URL == "" {
		return nil, 500, fmt.Errorf("数据源未配置访问地址")
	}
	full := strings.TrimSuffix(src.URL, "/") + path
	if qs := query.Encode(); qs != "" {
		full += "?" + qs
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, full, nil)
	if err != nil {
		return nil, 500, err
	}
	applyLokiAuth(req, src)
	client := lokiHTTPClients[0]
	if src.SkipTLSVerify {
		client = lokiHTTPClients[1]
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, 502, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
	if err != nil {
		return nil, 502, err
	}
	if resp.StatusCode >= 400 {
		return body, resp.StatusCode, fmt.Errorf("Loki 返回 HTTP %d: %s", resp.StatusCode, truncateStr(string(body), 300))
	}
	return body, resp.StatusCode, nil
}

func applyLokiAuth(req *http.Request, src *models.LogSource) {
	switch src.AuthType {
	case "basic":
		if pass, err := DecryptCred(src.BasicPass); err == nil {
			req.Header.Set("Authorization", "Basic "+
				base64.StdEncoding.EncodeToString([]byte(src.BasicUser+":"+pass)))
		}
	case "bearer":
		if token, err := DecryptCred(src.BearerToken); err == nil {
			req.Header.Set("Authorization", "Bearer "+token)
		}
	}
	for name, encrypted := range src.Headers {
		if validSourceHeader(name) {
			if value, err := DecryptCred(encrypted); err == nil {
				req.Header.Set(name, value)
			}
		}
	}
}

func truncateStr(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}

// ==================== 查询实现(契约对齐 aiops) ====================

type LogEntryRow struct {
	Timestamp   string            `json:"timestamp"`
	TimestampNs string            `json:"timestampNs"`
	Level       string            `json:"level"`
	Message     string            `json:"message"`
	Labels      map[string]string `json:"labels"`
}

type LogQueryResult struct {
	Entries    []LogEntryRow `json:"entries"`
	HasMore    bool          `json:"hasMore"`
	NextCursor string        `json:"nextCursor"`
}

// ContextLogSelector uses stable stream identity labels. Loki can add per-entry
// labels such as detected_level to query results; matching those would hide
// adjacent lines whose detected level differs from the selected line.
func ContextLogSelector(labels map[string]string) (string, error) {
	for key := range labels {
		if !ValidLogLabelName(key) {
			return "", fmt.Errorf("标签名无效")
		}
	}
	keys := []string{"cluster", "namespace", "app", "pod", "container", "filename"}
	if labels["pod"] == "" {
		keys = append(keys, "job", "instance", "host")
	}
	parts := make([]string, 0, len(keys))
	for _, key := range keys {
		if value := labels[key]; value != "" {
			encoded, _ := json.Marshal(value)
			parts = append(parts, key+"="+string(encoded))
		}
	}
	if len(parts) == 0 {
		for key, value := range labels {
			if value == "" || key == "detected_level" || key == "level" || strings.HasPrefix(key, "__") {
				continue
			}
			encoded, _ := json.Marshal(value)
			parts = append(parts, key+"="+string(encoded))
		}
		sort.Strings(parts)
	}
	if len(parts) == 0 {
		return "", fmt.Errorf("缺少可用于定位日志流的标签")
	}
	return "{" + strings.Join(parts, ",") + "}", nil
}

// QueryLokiLogs 日志查询:flatten streams → 统一行;cursor(最后一条纳秒时间戳)向后翻页
func QueryLokiLogs(ctx context.Context, src *models.LogSource, query string, start, end int64, limit int, direction, cursor string) (*LogQueryResult, error) {
	if limit <= 0 {
		limit = 200
	}
	if limit > maxLogQueryLimit {
		limit = maxLogQueryLimit
	}
	if direction == "" {
		direction = "backward"
	}
	if end <= 0 {
		end = time.Now().UnixNano()
	}
	if start <= 0 {
		start = end - int64(time.Hour) // 默认最近 1 小时
	}
	if cursor != "" {
		if c, err := strconv.ParseInt(cursor, 10, 64); err == nil {
			if direction == "forward" {
				start = c + 1
			} else {
				end = c - 1
			}
		}
	}

	q := url.Values{}
	q.Set("query", query)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("limit", strconv.Itoa(limit))
	q.Set("direction", direction)

	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/query_range", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Status string `json:"status"`
		Data   struct {
			ResultType string `json:"resultType"`
			Result     []struct {
				Stream map[string]string `json:"stream"`
				Values [][2]string       `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, fmt.Errorf("解析 Loki 响应失败: %w", err)
	}

	entries := make([]LogEntryRow, 0, limit)
	var lastTs string
	for _, stream := range resp.Data.Result {
		for _, v := range stream.Values {
			entries = append(entries, LogEntryRow{
				Timestamp:   formatLokiTimestamp(v[0]),
				TimestampNs: v[0],
				Level:       detectLevel(stream.Stream, v[1]),
				Message:     v[1],
				Labels:      stream.Stream,
			})
		}
	}
	sort.SliceStable(entries, func(i, j int) bool {
		a, _ := strconv.ParseInt(entries[i].TimestampNs, 10, 64)
		b, _ := strconv.ParseInt(entries[j].TimestampNs, 10, 64)
		if direction == "forward" {
			return a < b
		}
		return a > b
	})
	if len(entries) > limit {
		entries = entries[:limit]
	}
	if len(entries) > 0 {
		lastTs = entries[len(entries)-1].TimestampNs
	}

	nextCursor := ""
	if len(entries) >= limit && lastTs != "" {
		if ts, err := strconv.ParseInt(strings.Split(lastTs, ".")[0], 10, 64); err == nil {
			nextCursor = strconv.FormatInt(ts, 10)
		}
	}
	return &LogQueryResult{Entries: entries, HasMore: nextCursor != "", NextCursor: nextCursor}, nil
}

// formatLokiTimestamp Loki 时间串(纳秒字符串/科学计数法)→ RFC3339
func formatLokiTimestamp(ts string) string {
	if i := strings.IndexAny(ts, ".eE"); i > 0 {
		ts = ts[:i]
	}
	n, err := strconv.ParseInt(ts, 10, 64)
	if err != nil {
		return ts
	}
	// 自适应纳秒/微秒/毫秒/秒
	switch {
	case n > 1e18:
		return time.Unix(0, n).Format(time.RFC3339)
	case n > 1e15:
		return time.Unix(0, n*1e3).Format(time.RFC3339)
	case n > 1e12:
		return time.Unix(0, n*1e6).Format(time.RFC3339)
	case n > 1e9:
		return time.Unix(0, n*1e9).Format(time.RFC3339)
	default:
		return time.Unix(n, 0).Format(time.RFC3339)
	}
}

// detectLevel 优先取 level 标签;否则按行内容推断
func detectLevel(labels map[string]string, line string) string {
	if lv, ok := labels["level"]; ok && lv != "" {
		return strings.ToLower(lv)
	}
	l := strings.ToLower(line)
	for _, p := range []string{"error", "err", "fatal", "panic"} {
		if strings.Contains(l, p) {
			return "error"
		}
	}
	if strings.Contains(l, "warn") {
		return "warn"
	}
	return "info"
}

// LokiLabels 标签名列表
func LokiLabels(ctx context.Context, src *models.LogSource, start, end int64) ([]string, error) {
	q := url.Values{}
	if end > 0 {
		q.Set("end", strconv.FormatInt(end, 10))
	}
	if start > 0 {
		q.Set("start", strconv.FormatInt(start, 10))
	}
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/labels", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data []string `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	// Loki 在某 label 无值时会返回 {"data":null}，Go nil slice 序列化成 JSON null，
	// 前端 setState(null) 后 .map() 直接抛 "Cannot read properties of null"。
	// 统一归一化为空切片，前端拿到的永远是数组。
	if resp.Data == nil {
		resp.Data = []string{}
	}
	return resp.Data, nil
}

// LokiLabelValues 标签值列表(query 只取 {...} selector 部分)
func LokiLabelValues(ctx context.Context, src *models.LogSource, name, selector string, start, end int64) ([]string, error) {
	q := url.Values{}
	if s := extractStreamSelector(selector); s != "" && s != "{}" {
		q.Set("query", s)
	}
	if end > 0 {
		q.Set("end", strconv.FormatInt(end, 10))
	}
	if start > 0 {
		q.Set("start", strconv.FormatInt(start, 10))
	}
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/label/"+url.PathEscape(name)+"/values", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data []string `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	// 同上：label 无值时 Loki 返回 {"data":null}，归一化为空切片，
	// 避免前端 .map() 遇到 null 崩溃（快捷 LogQL 报错的直接原因）。
	if resp.Data == nil {
		resp.Data = []string{}
	}
	return resp.Data, nil
}

// extractStreamSelector 从 LogQL 中提取 {...} 部分(labels API 仅接受 stream selector)
func extractStreamSelector(query string) string {
	i := strings.Index(query, "{")
	if i < 0 {
		return ""
	}
	j := strings.Index(query[i:], "}")
	if j < 0 {
		return ""
	}
	return query[i : i+j+1]
}

// LokiHistogram 日志量直方图:服务端包 sum(count_over_time(q[step]))
// 重查询护栏。设计依据：
//   - 单条 27 天的 count_over_time 曾把 Loki 打到 OOMKilled(2Gi)，
//     进而 Pod 不 Ready → Service 无端点 → 全站查询 502。
//   - 所以"算不完的查询"必须在入口就拒绝，而不是硬跑 120 秒把实例拖垮。
const (
	// 日志量统计最多统计 7 天；更长跨度请分段查（前端会提示缩小范围）
	maxHistogramRange = 7 * 24 * time.Hour
	// 重查询（histogram/metrics）单独的更短超时，避免长时间占用 CPU
	heavyQueryTimeout = 30 * time.Second
	// 全局并发重查询上限：超过直接快速失败，保护单实例 Loki
	maxConcurrentHeavyQueries = 2
)

// ErrTooManyQueries 并发重查询超限（handler 映射为 429）
var ErrTooManyQueries = errors.New("当前有较多日志量统计在进行，请稍后重试")

// heavyQuerySem 全局重查询信号量
var heavyQuerySem = make(chan struct{}, maxConcurrentHeavyQueries)

func acquireHeavyQuery() bool {
	select {
	case heavyQuerySem <- struct{}{}:
		return true
	default:
		return false
	}
}

func releaseHeavyQuery() { <-heavyQuerySem }

func LokiHistogram(ctx context.Context, src *models.LogSource, query string, start, end int64) ([]map[string]interface{}, error) {
	if end <= 0 {
		end = time.Now().UnixNano()
	}
	if start <= 0 {
		start = end - int64(time.Hour)
	}
	span := end - start
	if span > int64(maxHistogramRange) {
		return nil, fmt.Errorf("统计范围过大（%.1f 天），最多支持 %.0f 天；请缩小时间范围后重试",
			float64(span)/float64(24*time.Hour), maxHistogramRange.Hours()/24)
	}
	// 并发护栏：拿不到令牌立刻返回，不排队（排队只会让用户等更久）
	if !acquireHeavyQuery() {
		return nil, ErrTooManyQueries
	}
	defer releaseHeavyQuery()

	// 独立超时：Loki 侧再慢也不会超过 30s，防止单查询长期占用实例
	ctx, cancel := context.WithTimeout(ctx, heavyQueryTimeout)
	defer cancel()

	step := pickHistogramStep(span)

	// ── 首选：Loki 官方的 index/volume_range ──────────────────────────────
	// 它只读 TSDB 索引里的 chunk 元数据（字节数），**不解码日志内容**。
	// 实测同一条件 24 小时：count_over_time 53s 超时失败，volume_range 111ms。
	// 返回单位是「字节」，所以 bucket 里带 unit=bytes，前端据此格式化为 KB/MB/GB。
	q := url.Values{}
	q.Set("query", query)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("step", step)
	q.Set("limit", "50")
	if body, status, err := lokiProxyGET(ctx, src, "/loki/api/v1/index/volume_range", q); err == nil && status < 400 {
		if buckets, ok := parseVolumeBuckets(body); ok {
			return buckets, nil
		}
	}
	// 老版本 Loki 没有 volume_range（404/400）时回退到 count_over_time
	return lokiHistogramByCount(ctx, src, query, step, start, end)
}

// parseVolumeBuckets 解析 index/volume_range 的 matrix 响应。
// 成功但无数据也返回 ok=true（空图），只有结构不对才 ok=false 触发回退。
func parseVolumeBuckets(body []byte) ([]map[string]interface{}, bool) {
	var resp struct {
		Data struct {
			Result []struct {
				Values [][2]interface{} `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, false
	}
	buckets := []map[string]interface{}{}
	if len(resp.Data.Result) > 0 {
		for _, v := range resp.Data.Result[0].Values {
			tsf, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[0]), 64)
			bytes, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[1]), 64)
			buckets = append(buckets, map[string]interface{}{
				"timestamp": time.Unix(int64(tsf), 0).Format(time.RFC3339),
				"count":     int64(bytes),
				"unit":      "bytes",
			})
		}
	}
	return buckets, true
}

// lokiHistogramByCount 回退实现：count_over_time 逐点统计「条数」。
// 代价高（要真读 chunk），只在 volume_range 不可用时使用。
func lokiHistogramByCount(ctx context.Context, src *models.LogSource, query, step string, start, end int64) ([]map[string]interface{}, error) {
	wrapped := fmt.Sprintf("sum(count_over_time(%s[%s]))", query, step)
	q := url.Values{}
	q.Set("query", wrapped)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("step", step)
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/query_range", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data struct {
			Result []struct {
				Values [][2]interface{} `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	buckets := []map[string]interface{}{}
	if len(resp.Data.Result) > 0 {
		for _, v := range resp.Data.Result[0].Values {
			tsf, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[0]), 64)
			cnt, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[1]), 64)
			buckets = append(buckets, map[string]interface{}{
				"timestamp": time.Unix(int64(tsf), 0).Format(time.RFC3339),
				"count":     int64(cnt),
				"unit":      "entries",
			})
		}
	}
	return buckets, nil
}

// MetricSeriesPoint 指标序列上的一个采样点
type MetricSeriesPoint struct {
	Timestamp string  `json:"timestamp"`
	Value     float64 `json:"value"`
}

// MetricSeries 一条按标签分组的指标序列(如 sum by (app)(...) 的每个 app 一条)
type MetricSeries struct {
	Labels map[string]string   `json:"labels"`
	Points []MetricSeriesPoint `json:"points"`
}

// LokiMetrics 通用 LogQL 指标查询(日志分析页用):query 为完整指标表达式,
// 其中所有 [$step] 占位符替换为按时间跨度自适应的步长后走 query_range,
// 返回矩阵多序列(直方图只取单序列,这里按 by (...) 分组整表返回)
func LokiMetrics(ctx context.Context, src *models.LogSource, query string, start, end int64) ([]MetricSeries, error) {
	if end <= 0 {
		end = time.Now().UnixNano()
	}
	if start <= 0 {
		start = end - int64(time.Hour)
	}
	// 与 histogram 共用护栏：并发限 2、单次 30s 超时，避免 analysis 页
	// 多个卡片同时发起 count_over_time 把单实例 Loki 打满。
	if !acquireHeavyQuery() {
		return nil, ErrTooManyQueries
	}
	defer releaseHeavyQuery()
	ctx, cancel := context.WithTimeout(ctx, heavyQueryTimeout)
	defer cancel()

	step := pickHistogramStep(end - start)
	wrapped := strings.ReplaceAll(query, "[$step]", "["+step+"]")

	q := url.Values{}
	q.Set("query", wrapped)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("step", step)
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/query_range", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data struct {
			Result []struct {
				Metric map[string]string `json:"metric"`
				Values [][2]interface{}  `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	series := make([]MetricSeries, 0, len(resp.Data.Result))
	for _, item := range resp.Data.Result {
		points := make([]MetricSeriesPoint, 0, len(item.Values))
		for _, v := range item.Values {
			// metric API 的 ts 为秒级浮点数(如 1.79e+09)
			tsf, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[0]), 64)
			cnt, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[1]), 64)
			points = append(points, MetricSeriesPoint{
				Timestamp: time.Unix(int64(tsf), 0).Format(time.RFC3339),
				Value:     cnt,
			})
		}
		if len(points) == 0 {
			continue
		}
		series = append(series, MetricSeries{Labels: item.Metric, Points: points})
	}
	return series, nil
}

// pickHistogramStep 按时间跨度自适应步长(aiops 同款)
func pickHistogramStep(spanNano int64) string {
	span := time.Duration(spanNano)
	switch {
	case span <= 15*time.Minute:
		return "30s"
	case span <= time.Hour:
		return "1m"
	case span <= 6*time.Hour:
		return "5m"
	case span <= 24*time.Hour:
		return "15m"
	case span <= 7*24*time.Hour:
		return "1h"
	default:
		return "6h"
	}
}

// DeploymentReady 部署级就绪探测:按部署名匹配自动注册的数据源,经其访问 Loki /ready
func DeploymentReady(ctx context.Context, deploymentName string) (bool, string) {
	if deploymentName == "" {
		return false, ""
	}
	db := repository.GetDB()
	db.RLock()
	var src *models.LogSource
	for i := range db.LogSources {
		if db.LogSources[i].Name == deploymentName {
			s := db.LogSources[i]
			src = &s
			break
		}
	}
	db.RUnlock()
	if src == nil {
		return false, "未注册数据源"
	}
	if LokiReady(ctx, src) {
		return true, ""
	}
	return false, "Loki 未就绪"
}

// LokiReady Loki /ready 探活(部署卡片/查询页状态点)
func LokiReady(ctx context.Context, src *models.LogSource) bool {
	body, code, err := lokiProxyGET(ctx, src, "/ready", nil)
	if err != nil || code != 200 {
		return false
	}
	return strings.Contains(strings.ToLower(string(body)), "ready")
}

// TestLogSource 数据源连通性(/loki/api/v1/labels 探活)
func TestLogSource(ctx context.Context, src *models.LogSource) map[string]interface{} {
	labels, err := LokiLabels(ctx, src, 0, 0)
	if err != nil {
		return map[string]interface{}{"success": false, "message": err.Error()}
	}
	return map[string]interface{}{
		"success": true,
		"message": fmt.Sprintf("连接成功,发现 %d 个标签", len(labels)),
		"labels":  labels,
	}
}