opsxdb/server/internal/db/es.go

package db

// es.go —— Elasticsearch 原生接入（HTTP + X-Pack SQL REST API）。
//
// 为什么不用官方 JDBC：x-pack-sql-jdbc 在 basic 许可下直接拒绝工作
// （"current license is non-compliant for [jdbc]"，实测 8.15.3），而 SQL REST API
// 在 basic 许可可用 —— 走 net/http 零依赖接入，也免除 JVM 负担。
//
// 能力边界（v1，对标 SQL 工作台）：
//   - 对象树：单库节点 "default" → SHOW TABLES 列出索引
//   - 查询：SQL 编辑器里的 ES SQL（SELECT/SHOW/DESC…）直发 /_sql
//   - ES SQL 只读：DDL/DML 由 ES 自身拒绝；危险语句拦截仍在本层先过一遍
//   - 分页：/_sql 为游标式，前端 offset 翻页暂不支持（首屏 fetch_size 截断）

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

func isES(d string) bool { return driverKey(d) == "elasticsearch" }

func esBase(c ConnReq) string {
	host := strings.TrimSpace(c.Host)
	if host == "" {
		host = "127.0.0.1"
	}
	port := c.Port
	if port == 0 {
		port = 9200
	}
	return fmt.Sprintf("http://%s:%d", host, port)
}

// esDo 发请求并在设置了用户时启用 Basic Auth（ES 开启 security 时用）。
func esDo(ctx context.Context, c ConnReq, method, path string, body, out any) error {
	var rd io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return err
		}
		rd = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, esBase(c)+path, rd)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if c.User != "" {
		req.SetBasicAuth(c.User, c.Password)
	}
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return fmt.Errorf("ES HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
	}
	if out != nil {
		return json.NewDecoder(resp.Body).Decode(out)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	return nil
}

func esPing(c ConnReq) error {
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	var info struct {
		Tagline string `json:"tagline"`
	}
	if err := esDo(ctx, c, http.MethodGet, "/", nil, &info); err != nil {
		return err
	}
	if info.Tagline == "" {
		return fmt.Errorf("ES 响应异常：缺少 tagline")
	}
	return nil
}

// esSQLResult /_sql 的 JSON 响应（format=json）。
type esSQLResult struct {
	Columns []struct {
		Name string `json:"name"`
		Type string `json:"type"`
	} `json:"columns"`
	Rows [][]interface{} `json:"rows"`
}

func esRunSQL(ctx context.Context, c ConnReq, sqlText string, fetchSize int) (ExecResult, error) {
	var r esSQLResult
	body := map[string]any{"query": sqlText}
	if fetchSize > 0 {
		body["fetch_size"] = fetchSize
	}
	if err := esDo(ctx, c, http.MethodPost, "/_sql", body, &r); err != nil {
		return ExecResult{}, err
	}
	cols := make([]string, 0, len(r.Columns))
	for _, col := range r.Columns {
		cols = append(cols, col.Name)
	}
	rows := make([]map[string]interface{}, 0, len(r.Rows))
	for _, row := range r.Rows {
		m := make(map[string]interface{}, len(cols))
		for i, v := range row {
			if i < len(cols) {
				m[cols[i]] = v
			}
		}
		rows = append(rows, m)
	}
	return ExecResult{
		Columns:   cols,
		Rows:      rows,
		RowCount:  len(rows),
		Truncated: fetchSize > 0 && len(rows) >= fetchSize,
		Limit:     fetchSize,
	}, nil
}

func esListDatabases() []string { return []string{"default"} }

func esListTables(ctx context.Context, c ConnReq) ([]string, error) {
	res, err := esRunSQL(ctx, c, "SHOW TABLES", 10000)
	if err != nil {
		return nil, err
	}
	out := []string{}
	if len(res.Columns) == 0 {
		return out, nil
	}
	// SHOW TABLES 列为 [catalog, name, type, kind]：取 "name" 列（ES 8 实测），退化为首列
	idx := 0
	for i, col := range res.Columns {
		if col == "name" {
			idx = i
			break
		}
	}
	for _, row := range res.Rows {
		if v, ok := row[res.Columns[idx]].(string); ok && v != "" {
			out = append(out, v)
		}
	}
	return out, nil
}

func esListColumns(ctx context.Context, c ConnReq) ([]ColumnInfo, error) {
	table := strings.TrimSpace(c.Table)
	if table == "" {
		return nil, fmt.Errorf("缺少表名")
	}
	res, err := esRunSQL(ctx, c, "DESCRIBE \""+strings.ReplaceAll(table, "\"", "")+"\"", 1000)
	if err != nil {
		return nil, err
	}
	cols := []ColumnInfo{}
	for _, row := range res.Rows {
		name, _ := row["column"].(string)
		typ, _ := row["type"].(string)
		if name == "" {
			continue
		}
		cols = append(cols, ColumnInfo{Name: name, Type: typ, Nullable: true})
	}
	return cols, nil
}

// esExecSQL 全部按查询处理：ES SQL 只读，写语句由 ES 拒绝；危险语句在 ops.go 入口已拦。
func esExecSQL(ctx context.Context, c ConnReq, start time.Time) ExecResult {
	if c.Offset > 0 {
		return ExecResult{Columns: []string{}, Rows: []map[string]interface{}{}, Error: "ES 游标分页暂不支持翻页（仅首屏，可缩小范围重查）", TimeMs: time.Since(start).Milliseconds()}
	}
	limit := rowLimit(c)
	if limit > 10000 {
		limit = 10000
	}
	res, err := esRunSQL(ctx, c, strings.TrimSpace(c.SQL), limit)
	if err != nil {
		return execFail(start, err)
	}
	res.TimeMs = time.Since(start).Milliseconds()
	return res
}