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 }