opsxdb/apps/server/internal/db/ops.go

package db

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"
)

// withDB 统一获取/释放 *sql.DB:
//   - c.SessionID != "" → 会话内按 database 池化复用,不 Close;
//   - 否则 → ApplySSHTunnel(c) + Open + defer Close(短生命周期隧道随请求关闭)。
func withDB(c ConnReq, fn func(*sql.DB) error) error {
	if c.SessionID != "" {
		s, err := GetSession(c.SessionID)
		if err != nil {
			return err
		}
		db, err := s.sqlDB(c.Database)
		if err != nil {
			return err
		}
		return fn(db)
	}
	dial, tun, err := ApplySSHTunnel(c)
	if err != nil {
		return err
	}
	if tun != nil {
		defer tun.Close()
	}
	db, err := Open(dial)
	if err != nil {
		return err
	}
	defer db.Close()
	return fn(db)
}

func unsupportedDriverErr(c ConnReq) error {
	return fmt.Errorf("不支持的驱动「%s」(当前支持: mysql/mariadb/tidb、postgresql、sqlserver、oracle、sqlite、redis；agent 桥: dameng)", c.Driver)
}

func reqCtx(ctx context.Context) context.Context {
	if ctx == nil {
		return context.Background()
	}
	return ctx
}

func ListDatabases(ctx context.Context, c ConnReq) ([]string, error) {
	ctx = reqCtx(ctx)
	if isRedis(c.Driver) {
		return listDatabasesRedis(ctx, c)
	}
	if isES(c.Driver) {
		return esListDatabases(), nil
	}
	if isInflux(c.Driver) {
		return ifxListDatabases(ctx, c)
	}
	if isAgentDriver(c.Driver) {
		return agentListDatabases(ctx, c)
	}
	cc := c
	switch {
	case isMySQLFamily(c.Driver):
		cc.Database = ""
	case isPostgresFamily(c.Driver):
		if cc.Database == "" {
			cc.Database = "postgres"
		}
	case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
		if cc.Database == "" {
			cc.Database = "master"
		}
	}
	var out []string
	err := withDB(cc, func(db *sql.DB) error {
		var err error
		switch {
		case isMySQLFamily(c.Driver):
			out, err = listDatabasesMySQL(ctx, db)
		case isPostgresFamily(c.Driver):
			out, err = listDatabasesPostgres(ctx, db)
		case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
			out, err = listDatabasesMSSQL(ctx, db)
		case driverKey(c.Driver) == "oracle":
			out, err = listDatabasesOracle(ctx, db)
		case driverKey(c.Driver) == "sqlite":
			out, err = listDatabasesSQLite(db)
		default:
			err = unsupportedDriverErr(c)
		}
		return err
	})
	if err != nil {
		return nil, err
	}
	return out, nil
}

func ListTables(ctx context.Context, c ConnReq) ([]string, error) {
	ctx = reqCtx(ctx)
	if isRedis(c.Driver) {
		return listTablesRedis(ctx, c)
	}
	if isES(c.Driver) {
		return esListTables(ctx, c)
	}
	if isInflux(c.Driver) {
		return ifxListTables(ctx, c)
	}
	if isAgentDriver(c.Driver) {
		return agentListTables(ctx, c)
	}
	var out []string
	err := withDB(c, func(db *sql.DB) error {
		var err error
		switch {
		case isMySQLFamily(c.Driver):
			out, err = listTablesMySQL(ctx, db, c.Database)
		case isPostgresFamily(c.Driver):
			out, err = listTablesPostgres(ctx, db, c.Database)
		case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
			out, err = listTablesMSSQL(ctx, db, c.Database)
		case driverKey(c.Driver) == "oracle":
			out, err = listTablesOracle(ctx, db, c.Database)
		case driverKey(c.Driver) == "sqlite":
			out, err = listTablesSQLite(ctx, db, c.Database)
		default:
			err = unsupportedDriverErr(c)
		}
		return err
	})
	if err != nil {
		return nil, err
	}
	return out, nil
}

func ListColumns(ctx context.Context, c ConnReq) ([]ColumnInfo, error) {
	ctx = reqCtx(ctx)
	if c.Table == "" {
		return nil, fmt.Errorf("缺少表名")
	}
	if isRedis(c.Driver) {
		return listColumnsRedis(ctx, c, c.Table)
	}
	if isES(c.Driver) {
		return esListColumns(ctx, c)
	}
	if isInflux(c.Driver) {
		return ifxListColumns(ctx, c)
	}
	if isAgentDriver(c.Driver) {
		return agentListColumns(ctx, c)
	}
	var out []ColumnInfo
	err := withDB(c, func(db *sql.DB) error {
		var err error
		switch {
		case isMySQLFamily(c.Driver):
			out, err = listColumnsMySQL(ctx, db, c.Database, c.Table)
		case isPostgresFamily(c.Driver):
			out, err = listColumnsPostgres(ctx, db, c.Database, c.Table)
		case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
			out, err = listColumnsMSSQL(ctx, db, c.Database, c.Table)
		case driverKey(c.Driver) == "oracle":
			out, err = listColumnsOracle(ctx, db, c.Database, c.Table)
		case driverKey(c.Driver) == "sqlite":
			out, err = listColumnsSQLite(ctx, db, c.Database, c.Table)
		default:
			err = unsupportedDriverErr(c)
		}
		return err
	})
	if err != nil {
		return nil, err
	}
	return out, nil
}

func rowLimit(c ConnReq) int {
	n := c.Limit
	if n <= 0 {
		n = DefaultRowLimit
	}
	if n > MaxRowLimit {
		n = MaxRowLimit
	}
	return n
}

func errMsg(err error) string {
	if err == nil {
		return ""
	}
	if errors.Is(err, context.Canceled) {
		return "已取消"
	}
	// client gone / deadline tied to request cancel
	if errors.Is(err, context.DeadlineExceeded) {
		return err.Error()
	}
	msg := err.Error()
	if strings.Contains(msg, "context canceled") {
		return "已取消"
	}
	return msg
}

func execFail(start time.Time, err error) ExecResult {
	return ExecResult{
		Columns: []string{},
		Rows:    []map[string]interface{}{},
		TimeMs:  time.Since(start).Milliseconds(),
		Error:   errMsg(err),
	}
}

// needConfirm 构造危险语句确认响应(不执行)。
func needConfirm(start time.Time, hits []string) ExecResult {
	return ExecResult{
		Columns:     []string{},
		Rows:        []map[string]interface{}{},
		TimeMs:      time.Since(start).Milliseconds(),
		Error:       "检测到危险语句: " + strings.Join(hits, "; "),
		NeedConfirm: true,
	}
}

// ExecSQL 执行查询 / 非查询语句。ctx 来自 HTTP 请求，客户端断开/Abort 时可取消 DB 调用。
func ExecSQL(ctx context.Context, c ConnReq) ExecResult {
	start := time.Now()
	ctx = reqCtx(ctx)
	if err := ctx.Err(); err != nil {
		return execFail(start, err)
	}
	if strings.TrimSpace(c.SQL) == "" {
		return ExecResult{Columns: []string{}, Rows: []map[string]interface{}{}, Error: "empty sql"}
	}
	if isRedis(c.Driver) {
		return execRedis(ctx, c)
	}
	// 服务端权威危险语句拦截:未确认时返回 needConfirm,不执行
	if !c.ConfirmDanger {
		dangerSQL := c.SQL
		if driverKey(c.Driver) == "mongodb" {
			// Mongo DSL 动词映射进词表：deleteMany 等无 WHERE 概念，按 DELETE/UPDATE 同等管控
			dangerSQL = mongoVerbAsSQL(dangerSQL)
		}
		if hits := CheckDangerousSQL(c.Driver, dangerSQL); len(hits) > 0 {
			return needConfirm(start, hits)
		}
	}
	if isAgentDriver(c.Driver) {
		return agentExecSQL(ctx, c, start)
	}
	if isES(c.Driver) {
		return esExecSQL(ctx, c, start)
	}
	if isInflux(c.Driver) {
		return ifxExecSQL(ctx, c, start)
	}
	var res ExecResult
	err := withDB(c, func(db *sql.DB) error {
		res = execOnDB(ctx, db, start, c)
		return nil
	})
	if err != nil {
		return execFail(start, err)
	}
	return res
}

func execOnDB(ctx context.Context, db *sql.DB, start time.Time, c ConnReq) ExecResult {
	if !isQuerySQL(c.SQL) {
		r, err := db.ExecContext(ctx, c.SQL)
		if err != nil {
			return execFail(start, err)
		}
		affected, _ := r.RowsAffected()
		return ExecResult{
			Columns:  []string{"affected_rows"},
			Rows:     []map[string]interface{}{{"affected_rows": affected}},
			RowCount: 1,
			TimeMs:   time.Since(start).Milliseconds(),
		}
	}

	rows, err := db.QueryContext(ctx, c.SQL)
	if err != nil {
		return execFail(start, err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return execFail(start, err)
	}
	limit := rowLimit(c)
	offset := c.Offset
	if offset < 0 {
		offset = 0
	}
	var resultRows []map[string]interface{}
	truncated := false
	skipped := 0
	for rows.Next() {
		if err := ctx.Err(); err != nil {
			return execFail(start, err)
		}
		// offset 分页:先丢弃前 offset 行
		if skipped < offset {
			skipped++
			continue
		}
		raw := make([]interface{}, len(cols))
		ptrs := make([]interface{}, len(cols))
		for i := range raw {
			ptrs[i] = &raw[i]
		}
		if err := rows.Scan(ptrs...); err != nil {
			return execFail(start, err)
		}
		m := make(map[string]interface{}, len(cols))
		for i, col := range cols {
			m[col] = normalize(raw[i])
		}
		resultRows = append(resultRows, m)
		if len(resultRows) >= limit {
			truncated = true
			break
		}
	}
	if err := rows.Err(); err != nil {
		return execFail(start, err)
	}
	if resultRows == nil {
		resultRows = []map[string]interface{}{}
	}
	return ExecResult{
		Columns:   cols,
		Rows:      resultRows,
		RowCount:  len(resultRows),
		TimeMs:    time.Since(start).Milliseconds(),
		Truncated: truncated,
		Limit:     limit,
		Offset:    offset,
	}
}

// ExplainSQL 服务端统一的执行计划入口(替代前端方言包装)。
// mysql/pg/sqlite 走行列计划;mssql 用 SHOWPLAN_XML(钉住单连接);oracle 走 DBMS_XPLAN。
func ExplainSQL(ctx context.Context, c ConnReq) ExplainResult {
	start := time.Now()
	ctx = reqCtx(ctx)
	sqlText := strings.TrimSpace(c.SQL)
	if sqlText == "" {
		return ExplainResult{Error: "empty sql", TimeMs: time.Since(start).Milliseconds()}
	}
	sqlText = strings.TrimSuffix(sqlText, ";")
	if isRedis(c.Driver) {
		return ExplainResult{Error: "Redis 不支持执行计划", TimeMs: time.Since(start).Milliseconds()}
	}
	if isES(c.Driver) {
		return ExplainResult{Error: "ES SQL 不支持执行计划", TimeMs: time.Since(start).Milliseconds()}
	}
	if isInflux(c.Driver) {
		return ExplainResult{Error: "InfluxDB 不支持执行计划", TimeMs: time.Since(start).Milliseconds()}
	}
	if isAgentDriver(c.Driver) {
		return agentExplain(ctx, c, sqlText)
	}
	// 危险语句不做计划
	if hits := CheckDangerousSQL(c.Driver, sqlText); len(hits) > 0 {
		return ExplainResult{Error: "检测到危险语句: " + strings.Join(hits, "; "), TimeMs: time.Since(start).Milliseconds()}
	}

	var res ExplainResult
	err := withDB(c, func(db *sql.DB) error {
		res = explainOnDB(ctx, db, c, sqlText)
		return nil
	})
	if err != nil {
		return ExplainResult{Error: errMsg(err), TimeMs: time.Since(start).Milliseconds()}
	}
	res.TimeMs = time.Since(start).Milliseconds()
	return res
}

func explainOnDB(ctx context.Context, db *sql.DB, c ConnReq, sqlText string) ExplainResult {
	fail := func(err error) ExplainResult {
		return ExplainResult{Error: errMsg(err)}
	}
	switch {
	case isMySQLFamily(c.Driver):
		return explainRows(ctx, db, "EXPLAIN "+sqlText)
	case isPostgresFamily(c.Driver):
		return explainRows(ctx, db, "EXPLAIN "+sqlText)
	case driverKey(c.Driver) == "sqlite":
		return explainRows(ctx, db, "EXPLAIN QUERY PLAN "+sqlText)
	case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
		// SET SHOWPLAN_XML 必须与查询同连接,钉住单连接
		conn, err := db.Conn(ctx)
		if err != nil {
			return fail(err)
		}
		defer conn.Close()
		if _, err := conn.ExecContext(ctx, "SET SHOWPLAN_XML ON"); err != nil {
			return fail(err)
		}
		defer func() { _, _ = conn.ExecContext(context.Background(), "SET SHOWPLAN_XML OFF") }()
		rows, err := conn.QueryContext(ctx, sqlText)
		if err != nil {
			return fail(err)
		}
		defer rows.Close()
		var sb strings.Builder
		for rows.Next() {
			var s sql.NullString
			if err := rows.Scan(&s); err != nil {
				return fail(err)
			}
			sb.WriteString(s.String)
		}
		if err := rows.Err(); err != nil {
			return fail(err)
		}
		return ExplainResult{Text: sb.String()}
	case driverKey(c.Driver) == "oracle":
		if _, err := db.ExecContext(ctx, "EXPLAIN PLAN FOR "+sqlText); err != nil {
			return fail(err)
		}
		rows, err := db.QueryContext(ctx, "SELECT plan_table_output FROM TABLE(DBMS_XPLAN.DISPLAY())")
		if err != nil {
			return fail(err)
		}
		defer rows.Close()
		var sb strings.Builder
		for rows.Next() {
			var s sql.NullString
			if err := rows.Scan(&s); err != nil {
				return fail(err)
			}
			sb.WriteString(s.String)
			sb.WriteByte('\n')
		}
		if err := rows.Err(); err != nil {
			return fail(err)
		}
		return ExplainResult{Text: sb.String()}
	default:
		return ExplainResult{Error: errMsg(unsupportedDriverErr(c))}
	}
}

// explainRows 以行列形式读取执行计划(mysql/pg/sqlite)。
func explainRows(ctx context.Context, db *sql.DB, q string) ExplainResult {
	rows, err := db.QueryContext(ctx, q)
	if err != nil {
		return ExplainResult{Error: errMsg(err)}
	}
	defer rows.Close()
	cols, err := rows.Columns()
	if err != nil {
		return ExplainResult{Error: errMsg(err)}
	}
	out := []map[string]interface{}{}
	for rows.Next() {
		raw := make([]interface{}, len(cols))
		ptrs := make([]interface{}, len(cols))
		for i := range raw {
			ptrs[i] = &raw[i]
		}
		if err := rows.Scan(ptrs...); err != nil {
			return ExplainResult{Error: errMsg(err)}
		}
		m := make(map[string]interface{}, len(cols))
		for i, col := range cols {
			m[col] = normalize(raw[i])
		}
		out = append(out, m)
	}
	if err := rows.Err(); err != nil {
		return ExplainResult{Error: errMsg(err)}
	}
	return ExplainResult{Columns: cols, Rows: out}
}

func isQuerySQL(sql string) bool {
	s := stripSQLLead(sql)
	if s == "" {
		return true
	}
	up := strings.ToUpper(s)
	for _, p := range []string{"SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "WITH", "PRAGMA"} {
		if strings.HasPrefix(up, p+" ") || up == p || strings.HasPrefix(up, p+"\t") || strings.HasPrefix(up, p+"\n") {
			return true
		}
	}
	return false
}