507 lines
13 KiB
Go
507 lines
13 KiB
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
|
||
}
|