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

package db

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"
)

func ListDatabases(c ConnReq) ([]string, error) {
	if isRedis(c.Driver) {
		return listDatabasesRedis(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"
		}
	}
	db, err := Open(cc)
	if err != nil {
		return nil, err
	}
	defer db.Close()
	switch {
	case isMySQLFamily(c.Driver):
		return listDatabasesMySQL(db)
	case isPostgresFamily(c.Driver):
		return listDatabasesPostgres(db)
	case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
		return listDatabasesMSSQL(db)
	case driverKey(c.Driver) == "oracle":
		return listDatabasesOracle(db)
	case driverKey(c.Driver) == "sqlite":
		return listDatabasesSQLite(db)
	default:
		return nil, fmt.Errorf("unsupported driver")
	}
}

func ListTables(c ConnReq) ([]string, error) {
	if isRedis(c.Driver) {
		return listTablesRedis(c)
	}
	db, err := Open(c)
	if err != nil {
		return nil, err
	}
	defer db.Close()
	switch {
	case isMySQLFamily(c.Driver):
		return listTablesMySQL(db, c.Database)
	case isPostgresFamily(c.Driver):
		return listTablesPostgres(db, c.Database)
	case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
		return listTablesMSSQL(db, c.Database)
	case driverKey(c.Driver) == "oracle":
		return listTablesOracle(db, c.Database)
	case driverKey(c.Driver) == "sqlite":
		return listTablesSQLite(db, c.Database)
	default:
		return nil, fmt.Errorf("unsupported driver")
	}
}

func ListColumns(c ConnReq) ([]ColumnInfo, error) {
	if c.Table == "" {
		return nil, fmt.Errorf("缺少表名")
	}
	if isRedis(c.Driver) {
		return listColumnsRedis(c, c.Table)
	}
	db, err := Open(c)
	if err != nil {
		return nil, err
	}
	defer db.Close()
	switch {
	case isMySQLFamily(c.Driver):
		return listColumnsMySQL(db, c.Database, c.Table)
	case isPostgresFamily(c.Driver):
		return listColumnsPostgres(db, c.Database, c.Table)
	case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
		return listColumnsMSSQL(db, c.Database, c.Table)
	case driverKey(c.Driver) == "oracle":
		return listColumnsOracle(db, c.Database, c.Table)
	case driverKey(c.Driver) == "sqlite":
		return listColumnsSQLite(db, c.Database, c.Table)
	default:
		return nil, fmt.Errorf("unsupported driver")
	}
}

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),
	}
}

// ExecSQL 执行查询 / 非查询语句。ctx 来自 HTTP 请求，客户端断开/Abort 时可取消 DB 调用。
func ExecSQL(ctx context.Context, c ConnReq) ExecResult {
	start := time.Now()
	if ctx == nil {
		ctx = context.Background()
	}
	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)
	}
	db, err := Open(c)
	if err != nil {
		return execFail(start, err)
	}
	defer db.Close()

	if !isQuerySQL(c.SQL) {
		res, err := db.ExecContext(ctx, c.SQL)
		if err != nil {
			return execFail(start, err)
		}
		affected, _ := res.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)
	var resultRows []map[string]interface{}
	truncated := false
	for rows.Next() {
		if err := ctx.Err(); err != nil {
			return execFail(start, err)
		}
		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,
	}
}

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
}