chore: 恢复旧布局基线(迁移改动丢失后重建,apps/packages 布局暂存)

This commit is contained in:
logan committed 2026-10-14 20:31:17 +08:00
1 parent a9510235bb
commit 456fa6954b
203 files changed
+18787 -48

No files matched your search

File diff suppressed because it is too large. Load diff
+379
View File
@@ -0,0 +1,379 @@
package db
// agent_e2e_test.go —— JDBC agent 桥端到端验证(真实 JVM + 真实 JDBC 驱动)。
// 前置:本机 JDK、agent jar(OPSXDB_AGENT_JAR 或默认 ~/.opsxdb/agents/jdbc/)。
// 触发:OPSXDB_E2E_JDBC=1 才运行(CI 无 JVM 时自动跳过)。
// 覆盖:StartAgent 全链路(gson 校验/JVM 启动/握手)→ open → DDL/DML → query → 元数据。
import (
"context"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"testing"
"time"
)
func TestE2EAgentSQLiteBridge(t *testing.T) {
if os.Getenv("OPSXDB_E2E_JDBC") != "1" {
t.Skip("置 OPSXDB_E2E_JDBC=1 并备好 JDK 与 agent jar 后运行")
}
home, _ := os.UserHomeDir()
jars := filepath.Join(home, ".opsxdb", "agents", "jdbc", "drivers", "sqlite-e2e")
if _, err := os.Stat(filepath.Join(jars, "sqlite-jdbc-3.46.1.3.jar")); err != nil {
t.Skipf("缺 sqlite-jdbc jar: %s", jars)
}
t.Setenv("OPSXDB_E2E_JARS", jars)
spec := AgentSpec{
DriverKey: "e2e-sqlite",
DriverClass: "org.sqlite.JDBC",
URLTemplate: "jdbc:sqlite:{{.Database}}",
JarsDirEnv: "OPSXDB_E2E_JARS",
}
agentDriverSpecs[spec.DriverKey] = spec // 注册进表,AgentConfigured 才能按 JarsDirEnv 解析
t.Cleanup(func() { delete(agentDriverSpecs, spec.DriverKey) })
cli, err := StartAgent(spec)
if err != nil {
t.Fatalf("StartAgent: %v", err)
}
t.Cleanup(func() { cli.Shutdown(context.Background()) })
dbPath := filepath.Join(t.TempDir(), "e2e.db")
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
handle := ""
{
var res struct {
Handle string `json:"handle"`
}
if err := cli.call(ctx, "open", map[string]any{"url": "jdbc:sqlite:" + dbPath, "user": "", "password": ""}, &res); err != nil {
t.Fatalf("open: %v (stderr=%s)", err, cli.stderrTail.String())
}
handle = res.Handle
}
mustExec := func(sql string) {
t.Helper()
var r struct {
AffectedRows int64 `json:"affectedRows"`
}
if err := cli.call(ctx, "exec", map[string]any{"handle": handle, "sql": sql, "timeoutSec": 30}, &r); err != nil {
t.Fatalf("exec %q: %v (stderr=%s)", sql, err, cli.stderrTail.String())
}
}
mustExec("CREATE TABLE demo(id INTEGER PRIMARY KEY, name TEXT)")
mustExec("INSERT INTO demo(name) VALUES ('opsx-e2e')")
{
var res struct {
Columns []string `json:"columns"`
Rows []map[string]interface{} `json:"rows"`
}
if err := cli.call(ctx, "query", map[string]any{"handle": handle, "sql": "SELECT id, name FROM demo", "maxRows": 100, "timeoutSec": 30}, &res); err != nil {
t.Fatalf("query: %v", err)
}
if len(res.Rows) != 1 || res.Rows[0]["name"] != "opsx-e2e" {
t.Fatalf("rows = %+v", res.Rows)
}
}
{
var res struct {
Names []string `json:"names"`
}
if err := cli.call(ctx, "metaTables", map[string]any{"handle": handle, "schema": ""}, &res); err != nil {
t.Fatalf("metaTables: %v", err)
}
found := false
for _, n := range res.Names {
if n == "demo" {
found = true
}
}
if !found {
t.Fatalf("metaTables 未包含 demo: %v", res.Names)
}
}
{
var res struct {
Columns []ColumnInfo `json:"columns"`
}
if err := cli.call(ctx, "metaColumns", map[string]any{"handle": handle, "schema": "", "table": "demo"}, &res); err != nil {
t.Fatalf("metaColumns: %v", err)
}
if len(res.Columns) != 2 {
t.Fatalf("columns = %+v", res.Columns)
}
}
}
// TestE2EElasticsearch —— 真实 ES 8(docker,安全关闭)+ 官方 x-pack-sql-jdbc 走 agent 桥全链路。
// 触发:OPSXDB_E2E_ES=1;前置 ES :9200 可连、~/.opsxdb/agents/jdbc/drivers/elasticsearch/ 有官方 JDBC jar。
func TestE2EElasticsearch(t *testing.T) {
if os.Getenv("OPSXDB_E2E_ES") != "1" {
t.Skip("置 OPSXDB_E2E_ES=1 并启动 ES :9200 后运行")
}
// 种子文档(让 SHOW TABLES / SELECT 有真实对象)
req, _ := http.NewRequest(http.MethodPut, "http://127.0.0.1:9200/e2e_smoke/_doc/1?refresh=true", strings.NewReader(`{"name":"opsx","n":42}`))
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Skipf("ES :9200 不可达: %v", err)
}
_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()
if resp.StatusCode >= 400 {
t.Fatalf("种子文档写入失败: HTTP %d", resp.StatusCode)
}
c := ConnReq{Driver: "elasticsearch", Host: "127.0.0.1", Port: 9200}
if err := TestConnection(c); err != nil {
t.Fatalf("TestConnection: %v", err)
}
dbs, err := ListDatabases(context.Background(), c)
if err != nil || len(dbs) != 1 || dbs[0] != "default" {
t.Fatalf("ListDatabases = %v, %v", dbs, err)
}
cc := c
cc.Database = "default"
tables, err := ListTables(context.Background(), cc)
if err != nil {
t.Fatalf("ListTables: %v", err)
}
found := false
for _, n := range tables {
if n == "e2e_smoke" {
found = true
}
}
if !found {
t.Fatalf("tables 未含 e2e_smoke: %v", tables)
}
r := ExecSQL(context.Background(), ConnReq{Driver: "elasticsearch", Host: "127.0.0.1", Port: 9200, Database: "default", SQL: "SELECT name, n FROM e2e_smoke", Limit: 100})
if r.Error != "" {
t.Fatalf("ExecSQL: %s (stderr 见测试日志)", r.Error)
}
if r.RowCount < 1 {
t.Fatalf("rows = %+v", r.Rows)
}
}
// TestE2EZooKeeper —— 真实 ZK(docker)+ OpsxZkAgent 全链路:连接/浏览/读写/删除。
// 触发:OPSXDB_E2E_ZK=1;前置 ZK :2181、~/.opsxdb/agents/jdbc/drivers/zookeeper/ 有官方 jar。
func TestE2EZooKeeper(t *testing.T) {
if os.Getenv("OPSXDB_E2E_ZK") != "1" {
t.Skip("置 OPSXDB_E2E_ZK=1 并启动 ZK :2181 后运行")
}
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
c := ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181}
if err := TestConnection(c); err != nil {
t.Skipf("ZK :2181 不可达: %v", err)
}
dbs, err := ListDatabases(ctx, c)
if err != nil || len(dbs) == 0 {
t.Fatalf("ListDatabases = %v, %v", dbs, err)
}
hasZk := false
for _, n := range dbs {
if n == "zookeeper" {
hasZk = true
}
}
if !hasZk {
t.Fatalf("根下应含 zookeeper: %v", dbs)
}
cc := c
cc.Database = "zookeeper"
tables, err := ListTables(ctx, cc)
if err != nil {
t.Fatalf("ListTables: %v", err)
}
hasConfig := false
for _, n := range tables {
if n == "config" {
hasConfig = true
}
}
if !hasConfig {
t.Fatalf("/zookeeper 下应含 config: %v", tables)
}
// 读:GET
r := ExecSQL(ctx, ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181, Database: "zookeeper", SQL: "GET /zookeeper/config", Limit: 100})
if r.Error != "" {
t.Fatalf("GET: %s", r.Error)
}
if r.RowCount != 1 || r.Rows[0]["path"] != "/zookeeper/config" {
t.Fatalf("GET rows = %+v", r.Rows)
}
// 写:CREATE → GET → SET → DELETE(DELETE 无 WHERE 属危险语句,走确认通道)
w := func(sql string) {
t.Helper()
res := ExecSQL(ctx, ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181, SQL: sql, ConfirmDanger: true})
if res.Error != "" {
t.Fatalf("%s → %s", sql, res.Error)
}
}
w("CREATE /opsxe2e hello-opsx")
r = ExecSQL(ctx, ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181, SQL: "GET /opsxe2e", Limit: 100})
if r.Rows[0]["data"] != "hello-opsx" {
t.Fatalf("CREATE 后 data = %v", r.Rows[0]["data"])
}
w("SET /opsxe2e updated")
r = ExecSQL(ctx, ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181, SQL: "GET /opsxe2e", Limit: 100})
if r.Rows[0]["data"] != "updated" {
t.Fatalf("SET 后 data = %v", r.Rows[0]["data"])
}
// 危险语句未确认时必须先返回 NeedConfirm
guard := ExecSQL(ctx, ConnReq{Driver: "zookeeper", Host: "127.0.0.1", Port: 2181, SQL: "DELETE /opsxe2e"})
if !guard.NeedConfirm {
t.Fatalf("DELETE 未确认应返回 NeedConfirm,got %+v", guard)
}
w("DELETE /opsxe2e")
}
// TestE2EInfluxDB —— 真实 InfluxDB 2.7(docker)原生 HTTP 全链路:bucket/measurement/Flux 查询。
// 触发:OPSXDB_E2E_INFLUX=1;前置 InfluxDB :8086(org opsx / token opsxdb-demo-token / bucket demo)。
func TestE2EInfluxDB(t *testing.T) {
if os.Getenv("OPSXDB_E2E_INFLUX") != "1" {
t.Skip("置 OPSXDB_E2E_INFLUX=1 并启动 InfluxDB :8086 后运行")
}
// 种子点(行协议写入)
wreq, _ := http.NewRequest(http.MethodPost, "http://127.0.0.1:8086/api/v2/write?org=opsx&bucket=demo&precision=s",
strings.NewReader("e2e_m,site=local v=42"))
wreq.Header.Set("Authorization", "Token opsxdb-demo-token")
wreq.Header.Set("Content-Type", "text/plain; charset=utf-8")
wresp, err := http.DefaultClient.Do(wreq)
if err != nil {
t.Skipf("InfluxDB :8086 不可达: %v", err)
}
_, _ = io.Copy(io.Discard, wresp.Body)
wresp.Body.Close()
if wresp.StatusCode >= 400 {
t.Fatalf("种子写入失败: HTTP %d", wresp.StatusCode)
}
time.Sleep(500 * time.Millisecond) // 存储引擎落盘
c := ConnReq{Driver: "influxdb", Host: "127.0.0.1", Port: 8086, User: "opsx", Password: "opsxdb-demo-token"}
if err := TestConnection(c); err != nil {
t.Fatalf("TestConnection: %v", err)
}
dbs, err := ListDatabases(context.Background(), c)
if err != nil {
t.Fatalf("ListDatabases: %v", err)
}
hasDemo := false
for _, n := range dbs {
if n == "demo" {
hasDemo = true
}
}
if !hasDemo {
t.Fatalf("buckets 未含 demo: %v", dbs)
}
cc := c
cc.Database = "demo"
tables, err := ListTables(context.Background(), cc)
if err != nil {
t.Fatalf("ListTables: %v", err)
}
hasM := false
for _, n := range tables {
if n == "e2e_m" {
hasM = true
}
}
if !hasM {
t.Fatalf("measurements 未含 e2e_m: %v", tables)
}
r := ExecSQL(context.Background(), ConnReq{
Driver: "influxdb", Host: "127.0.0.1", Port: 8086, User: "opsx", Password: "opsxdb-demo-token",
SQL: `from(bucket: "demo") |> range(start: -1h) |> filter(fn: (r) => r._measurement == "e2e_m")`,
Limit: 100,
})
if r.Error != "" {
t.Fatalf("Flux 查询: %s", r.Error)
}
if r.RowCount < 1 {
t.Fatalf("rows = %+v", r.Rows)
}
}
// TestE2EMongoDB —— 真实 MongoDB 7 + OpsxMongoAgent 全链路:库/集合浏览、DSL 查改增删、危险确认。
// 触发:OPSXDB_E2E_MONGO=1;前置 MongoDB :27017、~/.opsxdb/agents/jdbc/drivers/mongodb/ 有官方 jar。
func TestE2EMongoDB(t *testing.T) {
if os.Getenv("OPSXDB_E2E_MONGO") != "1" {
t.Skip("置 OPSXDB_E2E_MONGO=1 并启动 MongoDB :27017 后运行")
}
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
c := ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017}
if err := TestConnection(c); err != nil {
t.Skipf("MongoDB :27017 不可达: %v", err)
}
// 种子:插入文档(确认通道)
seed := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.insertOne({"name":"opsx","n":7})`, ConfirmDanger: false})
if seed.Error != "" {
t.Fatalf("insertOne: %s", seed.Error)
}
dbs, err := ListDatabases(ctx, c)
if err != nil || len(dbs) == 0 {
t.Fatalf("ListDatabases = %v, %v", dbs, err)
}
hasDb := false
for _, n := range dbs {
if n == "opsxe2e" {
hasDb = true
}
}
if !hasDb {
t.Fatalf("库列表未含 opsxe2e: %v", dbs)
}
cc := c
cc.Database = "opsxe2e"
tables, err := ListTables(ctx, cc)
if err != nil || len(tables) == 0 {
t.Fatalf("ListTables = %v, %v", tables, err)
}
found := false
for _, n := range tables {
if n == "c1" {
found = true
}
}
if !found {
t.Fatalf("集合列表未含 c1: %v", tables)
}
cols, err := ListColumns(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", Table: "c1"})
if err != nil || len(cols) == 0 {
t.Fatalf("ListColumns = %v, %v", cols, err)
}
// 查询
q := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.find({"name":"opsx"})`, Limit: 100})
if q.Error != "" || q.RowCount < 1 {
t.Fatalf("find: err=%s rows=%+v", q.Error, q.Rows)
}
// 危险确认:deleteMany 未确认必须先返回 NeedConfirm
guard := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.deleteMany({"name":"opsx"})`})
if !guard.NeedConfirm {
t.Fatalf("deleteMany 未确认应 NeedConfirm,got %+v", guard)
}
// 更新 + 删除(确认通道)
u := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.updateOne({"name":"opsx"}, {"$set":{"n":8}})`, ConfirmDanger: true})
if u.Error != "" {
t.Fatalf("updateOne: %s", u.Error)
}
d := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.deleteMany({"name":"opsx"})`, ConfirmDanger: true})
if d.Error != "" {
t.Fatalf("deleteMany: %s", d.Error)
}
cnt := ExecSQL(ctx, ConnReq{Driver: "mongodb", Host: "127.0.0.1", Port: 27017, Database: "opsxe2e", SQL: `db.c1.count({"name":"opsx"})`})
if fmt.Sprintf("%v", cnt.Rows[0]["n"]) != "0" {
t.Fatalf("count after delete = %v", cnt.Rows[0]["n"])
}
}
+394
View File
@@ -0,0 +1,394 @@
package db
// agent_test.go —— JDBC agent 桥的协议层测试。
// 用 helper process 扮演 JVM agent(TestMain 拦截 -agent-stub),
// 无需本机安装 Java 即可验证:握手、请求/响应路由、超时、进程退出、并发串行写。
import (
"archive/zip"
"bufio"
"context"
"crypto/sha1"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"testing"
"time"
)
// TestMain 拦截 stub 模式:测试二进制被当成 agent 子进程启动时直接进入协议循环。
func TestMain(m *testing.M) {
for _, a := range os.Args {
if a == "-agent-stub" {
runAgentStub()
return
}
}
os.Exit(m.Run())
}
// runAgentStub 扮演 OpsxJdbcAgent 的最小协议面(与 agents/jdbc/OpsxJdbcAgent.java 对应)。
func runAgentStub() {
out := bufio.NewWriter(os.Stdout)
announce := agentResponse{ID: 0, OK: true}
announce.Result = json.RawMessage(`{"agent":"opsx-agent-jdbc","version":"0.1.0","protocol":1}`)
write := func(r agentResponse) {
b, _ := json.Marshal(r)
_, _ = out.WriteString(string(b) + "\n")
out.Flush()
}
write(announce)
in := bufio.NewScanner(os.Stdin)
in.Buffer(make([]byte, 1024*1024), 1024*1024)
for in.Scan() {
var req agentRequest
if json.Unmarshal(in.Bytes(), &req) != nil {
continue
}
args := req.Args
fail := func(msg string) agentResponse { return agentResponse{ID: req.ID, OK: false, Error: msg} }
switch req.Op {
case "hello":
r := agentResponse{ID: req.ID, OK: true}
r.Result = json.RawMessage(`{"agent":"opsx-agent-jdbc","version":"0.1.0","protocol":1}`)
write(r)
case "open":
if args["url"] == "" || args["url"] == nil {
write(fail("missing url"))
break
}
r := agentResponse{ID: req.ID, OK: true}
r.Result = json.RawMessage(`{"handle":"h1"}`)
write(r)
case "ping":
write(agentResponse{ID: req.ID, OK: true})
case "query":
if strings.Contains(fmt.Sprintf("%v", args["sql"]), "SLOW") {
time.Sleep(2 * time.Second)
}
r := agentResponse{ID: req.ID, OK: true}
r.Result = json.RawMessage(`{"columns":["id","name"],"rows":[{"id":"1","name":"demo"}],"truncated":false}`)
write(r)
case "exec":
r := agentResponse{ID: req.ID, OK: true}
r.Result = json.RawMessage(`{"affectedRows":3}`)
write(r)
case "metaDatabases":
r := agentResponse{ID: req.ID, OK: true}
r.Result = json.RawMessage(`{"names":["MAIN","SYS"]}`)
write(r)
case "close":
write(agentResponse{ID: req.ID, OK: true})
case "shutdown":
write(agentResponse{ID: req.ID, OK: true})
os.Exit(0)
default:
write(fail("unknown op: " + req.Op))
}
}
os.Exit(0)
}
// TestAgentStubHelper 占位:stub 子进程以 -test.run=TestAgentStubHelper 启动时
// 不会走到这里(TestMain 已拦截),但名字必须存在以匹配 -test.run。
func TestAgentStubHelper(t *testing.T) {}
func startStubAgent(t *testing.T) *AgentClient {
t.Helper()
cmd := exec.Command(os.Args[0], "-test.run=TestAgentStubHelper", "--", "-agent-stub")
cli, err := StartAgentFromCmd(cmd, agentDriverSpecs["dameng"])
if err != nil {
t.Fatalf("StartAgentFromCmd: %v", err)
}
t.Cleanup(func() { cli.Shutdown(context.Background()) })
return cli
}
func TestAgentClientRoundTrip(t *testing.T) {
cli := startStubAgent(t)
ctx := context.Background()
// open → query → exec → metaDatabases → ping
handle := ""
{
var res struct {
Handle string `json:"handle"`
}
if err := cli.call(ctx, "open", map[string]any{"url": "jdbc:dm://127.0.0.1:5236", "user": "SYSDBA"}, &res); err != nil {
t.Fatalf("open: %v", err)
}
if res.Handle != "h1" {
t.Fatalf("handle = %q, want h1", res.Handle)
}
handle = res.Handle
}
{
var res struct {
Columns []string `json:"columns"`
Rows []map[string]interface{} `json:"rows"`
}
if err := cli.call(ctx, "query", map[string]any{"handle": handle, "sql": "SELECT 1", "maxRows": 100}, &res); err != nil {
t.Fatalf("query: %v", err)
}
if len(res.Columns) != 2 || res.Rows[0]["name"] != "demo" {
t.Fatalf("query result mismatch: %+v", res)
}
}
{
var res struct {
AffectedRows int64 `json:"affectedRows"`
}
if err := cli.call(ctx, "exec", map[string]any{"handle": handle, "sql": "UPDATE t SET a=1"}, &res); err != nil {
t.Fatalf("exec: %v", err)
}
if res.AffectedRows != 3 {
t.Fatalf("affectedRows = %d, want 3", res.AffectedRows)
}
}
{
var res struct {
Names []string `json:"names"`
}
if err := cli.call(ctx, "metaDatabases", map[string]any{"handle": handle}, &res); err != nil {
t.Fatalf("metaDatabases: %v", err)
}
if len(res.Names) != 2 || res.Names[0] != "MAIN" {
t.Fatalf("names = %v", res.Names)
}
}
if err := cli.call(ctx, "ping", map[string]any{"handle": handle}, nil); err != nil {
t.Fatalf("ping: %v", err)
}
}
func TestAgentCallTimeout(t *testing.T) {
cli := startStubAgent(t)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
var res map[string]any
err := cli.call(ctx, "query", map[string]any{"handle": "h1", "sql": "SELECT SLOW", "maxRows": 10}, &res)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("err = %v, want DeadlineExceeded", err)
}
}
func TestAgentProcessExit(t *testing.T) {
cli := startStubAgent(t)
cli.Kill() // 杀进程后再调用,必须报进程退出而非挂死
var res map[string]any
err := cli.call(context.Background(), "ping", map[string]any{"handle": "h1"}, &res)
if err == nil {
t.Fatal("期望进程退出错误,got nil")
}
}
func TestAgentConcurrentCalls(t *testing.T) {
cli := startStubAgent(t)
ctx := context.Background()
if err := cli.call(ctx, "open", map[string]any{"url": "jdbc:dm://h:5236"}, nil); err != nil {
t.Fatalf("open: %v", err)
}
const n = 24
var wg sync.WaitGroup
errs := make([]error, n)
for i := 0; i < n; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
var res struct {
Columns []string `json:"columns"`
}
errs[i] = cli.call(ctx, "query", map[string]any{"handle": "h1", "sql": "SELECT 1", "maxRows": 1}, &res)
}(i)
}
wg.Wait()
for i, err := range errs {
if err != nil {
t.Fatalf("call %d: %v", i, err)
}
}
}
func TestRenderAgentURL(t *testing.T) {
spec := agentDriverSpecs["dameng"]
url, err := renderAgentURL(spec, ConnReq{Host: "127.0.0.1", Port: 5236})
if err != nil {
t.Fatalf("render: %v", err)
}
if url != "jdbc:dm://127.0.0.1:5236" {
t.Fatalf("url = %q", url)
}
if _, err := renderAgentURL(spec, ConnReq{Host: "", Port: 5236}); err == nil {
t.Fatal("空主机必须报错")
}
if _, err := renderAgentURL(spec, ConnReq{Host: "127.0.0.1", Port: 0}); err == nil {
t.Fatal("零端口必须报错")
}
}
func TestDownloadPinnedChecksum(t *testing.T) {
payload := []byte("fake-gson-bytes")
good := sha1.Sum(payload)
cases := []struct {
name string
body []byte
want string
wantErr bool
}{
{"校验通过", payload, hex.EncodeToString(good[:]), false},
{"校验失败拒绝落盘", []byte("tampered"), hex.EncodeToString(good[:]), true},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write(tc.body)
}))
defer srv.Close()
dest := t.TempDir() + "/gson.jar"
err := downloadPinned(srv.URL, dest, tc.want)
if tc.wantErr && err == nil {
t.Fatal("期望校验失败")
}
if !tc.wantErr {
if err != nil {
t.Fatalf("download: %v", err)
}
if _, statErr := os.Stat(dest); statErr != nil {
t.Fatalf("文件未落盘: %v", statErr)
}
} else {
if _, statErr := os.Stat(dest); statErr == nil {
t.Fatal("校验失败时文件不应存在")
}
}
})
}
}
func TestAgentConfigured(t *testing.T) {
dir := t.TempDir()
t.Setenv("OPSXDB_DAMENG_JARS", dir)
if ok, _ := AgentConfigured("dameng"); ok {
t.Fatal("空目录不应判为已配置")
}
if err := os.WriteFile(dir+"/DmJdbcDriver-8.1.3.140.jar", []byte("x"), 0o644); err != nil {
t.Fatal(err)
}
if ok, _ := AgentConfigured("dameng"); !ok {
t.Fatal("放置 jar 后应判为已配置")
}
if IsAgentDriver("mysql") {
t.Fatal("mysql 不应是 agent 驱动")
}
if !IsAgentDriver("dameng") {
t.Fatal("dameng 应是 agent 驱动")
}
}
// TestEnsureAgentRuntimeOfflinePack —— 网络不可用(镜像指向无效 + sha256 不匹配)时,
// ensureAgentRuntime 必须从离线驱动包导入 jar(pack.sh 产物)。
func TestEnsureAgentRuntimeOfflinePack(t *testing.T) {
home := t.TempDir()
t.Setenv("OPSXDB_HOME", home)
t.Setenv("OPSXDB_AGENT_JAR", filepath.Join(home, "agent", "opsx-agent-jdbc.jar"))
// 网络必失败:镜像指向本地空端口 + 错误的 sha256
t.Setenv("MAVEN_MIRROR", "http://127.0.0.1:1/none")
badSHA := sha1.Sum([]byte("bad"))
spec := AgentSpec{
DriverKey: "e2e-offline",
DriverClass: "x",
URLTemplate: "jdbc:x://{{.Host}}:{{.Port}}",
JarsDirEnv: "OPSXDB_OFFLINE_JARS",
}
agentDriverSpecs[spec.DriverKey] = spec
t.Cleanup(func() { delete(agentDriverSpecs, spec.DriverKey) })
agentDriverArtifacts[spec.DriverKey] = []agentArtifact{
{"demo.jar", "com.nope:demo:1.0", hex.EncodeToString(badSHA[:])},
}
// 构造离线包:manifest + 正确的 jar 内容(sha256 以包内 manifest 为准)
jarBytes := []byte("fake-jar-content")
files := map[string]string{"drivers/e2e-offline/demo.jar": string(jarBytes)}
manifestFiles := []map[string]string{{"path": "drivers/e2e-offline/demo.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(jarBytes))}}
zipPath := filepath.Join(home, "offline.zip")
if err := writeTestZip(zipPath, manifestFiles, files); err != nil {
t.Fatal(err)
}
t.Setenv("OPSXDB_OFFLINE_ZIP", zipPath)
// agent jar 也在包里(确保 ensureAgentJar 的离线路径不拦路)
if err := importFromOfflinePack(map[string]string{"pack/agent/opsx-agent-jdbc.jar": filepath.Join(home, "agent", "opsx-agent-jdbc.jar")}); err != nil {
// agent jar 不在此包属预期——单独构造一份最小 agent 条目
agentBytes := []byte("fake-agent")
mf2 := []map[string]string{{"path": "agent/opsx-agent-jdbc.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(agentBytes))}}
f2 := map[string]string{"agent/opsx-agent-jdbc.jar": string(agentBytes)}
if werr := writeTestZip(zipPath, append(manifestFiles, mf2...), mergeMaps(files, f2)); werr != nil {
t.Fatal(werr)
}
}
if err := ensureAgentRuntime(spec); err != nil {
t.Fatalf("ensureAgentRuntime 应走离线包成功: %v", err)
}
got, err := os.ReadFile(filepath.Join(home, "drivers", "e2e-offline", "demo.jar"))
if err != nil || string(got) != string(jarBytes) {
t.Fatalf("离线导入内容不符: %v", err)
}
// 负例:manifest hash 不符必须拒
jarBytes2 := []byte("tampered")
mf3 := []map[string]string{{"path": "drivers/e2e-offline/demo.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(jarBytes2))}}
if err := writeTestZip(zipPath, mf3, map[string]string{"drivers/e2e-offline/demo.jar": "not-the-same-content"}); err != nil {
t.Fatal(err)
}
os.Remove(filepath.Join(home, "drivers", "e2e-offline", "demo.jar"))
if err := ensureAgentRuntime(spec); err == nil {
t.Fatal("manifest 与实际内容不符时必须失败")
}
}
func writeTestZip(path string, manifestFiles []map[string]string, files map[string]string) error {
zf, err := os.Create(path)
if err != nil {
return err
}
defer zf.Close()
zw := zip.NewWriter(zf)
mf, err := zw.Create("pack/manifest.json")
if err != nil {
return err
}
b, _ := json.Marshal(map[string]any{"bundleVersion": 1, "files": manifestFiles})
if _, err := mf.Write(b); err != nil {
return err
}
for rel, content := range files {
w, err := zw.Create("pack/" + rel)
if err != nil {
return err
}
if _, err := w.Write([]byte(content)); err != nil {
return err
}
}
return zw.Close()
}
func mergeMaps(a, b map[string]string) map[string]string {
out := map[string]string{}
for k, v := range a {
out[k] = v
}
for k, v := range b {
out[k] = v
}
return out
}
+107
View File
@@ -0,0 +1,107 @@
package db
import (
"regexp"
"strings"
)
// 危险语句规则:词边界匹配,作用于 StripComments 之后的纯代码文本。
// 词表:DROP/TRUNCATE/ALTER/GRANT/REVOKE/CREATE USER|ROLE/RENAME TABLE/
// REPLACE INTO/LOAD DATA/INTO OUTFILE|DUMPFILE/xp_cmdshell/COPY..PROGRAM/EXEC/CALL。
var dangerPatterns = []struct {
re *regexp.Regexp
desc string
}{
{regexp.MustCompile(`\bDROP\b`), "DROP"},
{regexp.MustCompile(`\bTRUNCATE\b`), "TRUNCATE"},
{regexp.MustCompile(`\bALTER\b`), "ALTER"},
{regexp.MustCompile(`\bGRANT\b`), "GRANT"},
{regexp.MustCompile(`\bREVOKE\b`), "REVOKE"},
{regexp.MustCompile(`\bCREATE\s+(USER|ROLE)\b`), "CREATE USER/ROLE"},
{regexp.MustCompile(`\bRENAME\s+TABLE\b`), "RENAME TABLE"},
{regexp.MustCompile(`\bREPLACE\s+INTO\b`), "REPLACE INTO"},
{regexp.MustCompile(`\bLOAD\s+DATA\b`), "LOAD DATA"},
{regexp.MustCompile(`\bINTO\s+(OUTFILE|DUMPFILE)\b`), "INTO OUTFILE/DUMPFILE"},
{regexp.MustCompile(`(?i)\bxp_cmdshell\b`), "xp_cmdshell"},
{regexp.MustCompile(`\bCOPY\b[^;]*\b(FROM|TO)\s+PROGRAM\b`), "COPY FROM/TO PROGRAM"},
}
var (
reWriteNoWhere = regexp.MustCompile(`\b(DELETE|UPDATE)\b`)
reWhere = regexp.MustCompile(`\bWHERE\b`)
reExecCallHead = regexp.MustCompile(`^(EXEC|EXECUTE|CALL)\b`)
)
// CheckDangerousSQL 服务端权威的危险语句判定(前端预检仅为 UX)。
// 逐语句检查;注释/字符串内的关键字不计;MySQL /*! */ 可执行注释内容参与检查;
// WITH 包裹的写操作因全文扫描 CTE 体同样命中。
// 返回命中的危险原因清单(去重)。
func CheckDangerousSQL(driver, sqlText string) []string {
_ = driver // 词表当前各驱动通用,保留参数便于后续按方言细化
seen := map[string]struct{}{}
var hits []string
add := func(desc, stmt string) {
if _, ok := seen[desc]; ok {
return
}
seen[desc] = struct{}{}
hits = append(hits, desc+": "+truncateSQL(stmt, 80))
}
for _, stmt := range SplitStatements(sqlText) {
code := strings.ToUpper(StripComments(stmt))
code = strings.TrimSpace(code)
if code == "" {
continue
}
for _, p := range dangerPatterns {
if p.re.MatchString(code) {
add(p.desc, stmt)
}
}
// DELETE/UPDATE 无 WHERE(含 WITH CTE 体内的写操作)
if reWriteNoWhere.MatchString(code) && !reWhere.MatchString(code) {
kw := reWriteNoWhere.FindString(code)
add(kw+" 无 WHERE", stmt)
}
// 存储过程/动态执行仅匹配语句开头,降低误报
if m := reExecCallHead.FindString(code); m != "" {
add(m, stmt)
}
}
return hits
}
// IsDangerousSQL 兼容旧调用:命中任意危险规则即 true。
func IsDangerousSQL(sql string) bool {
return len(CheckDangerousSQL("", sql)) > 0
}
func truncateSQL(s string, n int) string {
s = strings.Join(strings.Fields(s), " ")
if len(s) <= n {
return s
}
return s[:n] + "…"
}
func stripSQLLead(sql string) string {
s := strings.TrimSpace(sql)
for {
if strings.HasPrefix(s, "--") {
if i := strings.IndexByte(s, '\n'); i >= 0 {
s = strings.TrimSpace(s[i+1:])
continue
}
return ""
}
if strings.HasPrefix(s, "/*") {
if i := strings.Index(s, "*/"); i >= 0 {
s = strings.TrimSpace(s[i+2:])
continue
}
return ""
}
break
}
return s
}
+76
View File
@@ -0,0 +1,76 @@
package db
import "testing"
// 危险语句命中/放行矩阵,含四类历史绕过:多语句、/*! */ 可执行注释、
// 注释吞掉 WHERE、WITH CTE 包裹写操作。
func TestCheckDangerousSQLMatrix(t *testing.T) {
dangerous := []string{
"DROP TABLE users",
"drop table users", // 大小写
"TRUNCATE TABLE t", // TRUNCATE
"ALTER TABLE t ADD COLUMN c int", // ALTER(误报方向维持现状)
"GRANT ALL ON *.* TO 'u'", // GRANT
"REVOKE ALL ON *.* FROM 'u'", // REVOKE
"CREATE USER 'u'@'%' IDENTIFIED BY 'p'", // CREATE USER
"CREATE ROLE app", // CREATE ROLE
"RENAME TABLE a TO b", // RENAME TABLE
"REPLACE INTO t VALUES (1)", // REPLACE INTO
"LOAD DATA INFILE '/tmp/x' INTO TABLE t", // LOAD DATA
"SELECT * INTO OUTFILE '/tmp/x' FROM t", // INTO OUTFILE
"SELECT * INTO DUMPFILE '/tmp/x' FROM t", // INTO DUMPFILE
"EXEC xp_cmdshell 'dir'", // xp_cmdshell
"EXEC sp_configure 'x'", // EXEC 开头
"CALL do_thing()", // CALL 开头
"COPY t FROM PROGRAM 'cmd'", // COPY FROM PROGRAM
"DELETE FROM users", // DELETE 无 WHERE
"UPDATE users SET x=1", // UPDATE 无 WHERE
"SELECT 1; DROP TABLE t", // 多语句绕过
"/*!50000 DROP TABLE t*/", // MySQL 可执行注释
"SELECT 1; /*!40101 SET @x = 1*/; DROP TABLE t", // 可执行注释 + 多语句
"UPDATE t SET x=1 -- WHERE id=2", // 注释吞掉 WHERE
"UPDATE t SET x=1 # WHERE id=2", // # 注释吞掉 WHERE
"UPDATE t SET x=1 /* WHERE id=2 */", // 块注释吞掉 WHERE
"WITH x AS (SELECT 1) DELETE FROM t", // CTE 包裹 DELETE 无 WHERE
"WITH x AS (SELECT 1) UPDATE t SET a=1", // CTE 包裹 UPDATE 无 WHERE
}
for _, sql := range dangerous {
if hits := CheckDangerousSQL("mysql", sql); len(hits) == 0 {
t.Errorf("应命中但未命中: %q", sql)
}
}
safe := []string{
"SELECT * FROM users",
"SELECT * FROM users WHERE id = 1",
"DELETE FROM users WHERE id = 1",
"UPDATE users SET x = 1 WHERE id = 2",
"UPDATE users SET x = 1\nWHERE id = 2",
"INSERT INTO t VALUES (1)",
"CREATE TABLE t (id int)",
"SELECT 'DROP TABLE x'", // 字符串内关键字
"SELECT * FROM t WHERE note = 'truncate it'", // 字符串内关键字
"-- DROP TABLE t\nSELECT 1", // 注释内关键字
"SELECT * FROM t /* DROP TABLE x */ WHERE 1=1", // 块注释内关键字
"SELECT * FROM updated WHERE id = 1", // 词边界:updated 不算 UPDATE
"SELECT 1; SELECT 2",
"WITH x AS (SELECT 1) SELECT * FROM x",
"SHOW TABLES",
"EXPLAIN SELECT 1",
}
for _, sql := range safe {
if hits := CheckDangerousSQL("mysql", sql); len(hits) > 0 {
t.Errorf("误报: %q → %v", sql, hits)
}
}
}
// 旧接口兼容:IsDangerousSQL 委托 CheckDangerousSQL。
func TestIsDangerousSQLCompat(t *testing.T) {
if !IsDangerousSQL("DROP TABLE users") {
t.Fatal("drop")
}
if IsDangerousSQL("SELECT * FROM users") {
t.Fatal("select")
}
}
+121
View File
@@ -0,0 +1,121 @@
package db
import (
"context"
"database/sql"
"fmt"
"strings"
)
// GetDDL 返回表/视图的建表语句。mysql 族走 SHOW CREATE TABLE,sqlite 读 sqlite_master,
// pg/mssql/oracle 用 ListColumns 元数据在 Go 侧拼装 CREATE TABLE。
func GetDDL(ctx context.Context, c ConnReq) (string, error) {
ctx = reqCtx(ctx)
if c.Table == "" {
return "", fmt.Errorf("缺少表名")
}
if isRedis(c.Driver) {
return "", fmt.Errorf("该驱动不支持 DDL 查看")
}
// pg/mssql/oracle:复用列元数据拼装,避免再查一次库
switch {
case isPostgresFamily(c.Driver), isMSSQL(c.Driver), driverKey(c.Driver) == "oracle":
cols, err := ListColumns(ctx, c)
if err != nil {
return "", err
}
return buildDDLFromColumns(c.Driver, c.Table, cols)
}
var out string
err := withDB(c, func(db *sql.DB) error {
var err error
switch {
case isMySQLFamily(c.Driver):
out, err = getDDLMySQL(ctx, db, c)
case driverKey(c.Driver) == "sqlite":
out, err = getDDLSQLite(ctx, db, c.Table)
default:
err = unsupportedDriverErr(c)
}
return err
})
if err != nil {
return "", err
}
return out, nil
}
// getDDLMySQL SHOW CREATE TABLE 第二列即 DDL(列名随对象类型变化:Create Table/Create View)。
func getDDLMySQL(ctx context.Context, dbh *sql.DB, c ConnReq) (string, error) {
ident := quoteMySQLIdent(c.Table)
if c.Database != "" {
ident = quoteMySQLIdent(c.Database) + "." + ident
}
rows, err := dbh.QueryContext(ctx, "SHOW CREATE TABLE "+ident)
if err != nil {
return "", err
}
defer rows.Close()
if rows.Next() {
var name string
var ddl sql.NullString
if err := rows.Scan(&name, &ddl); err != nil {
return "", err
}
if !ddl.Valid {
return "", fmt.Errorf("无法读取 %s 的 DDL", c.Table)
}
return ddl.String, nil
}
if err := rows.Err(); err != nil {
return "", err
}
return "", fmt.Errorf("表 %s 不存在", c.Table)
}
// getDDLSQLite sqlite_master 的 sql 列即建表语句。
func getDDLSQLite(ctx context.Context, dbh *sql.DB, table string) (string, error) {
var ddl sql.NullString
err := dbh.QueryRowContext(ctx, `SELECT sql FROM sqlite_master WHERE name = ?`, table).Scan(&ddl)
if err == sql.ErrNoRows {
return "", fmt.Errorf("表 %s 不存在", table)
}
if err != nil {
return "", err
}
if !ddl.Valid {
return "", fmt.Errorf("表 %s 无 DDL 记录", table)
}
return ddl.String, nil
}
// buildDDLFromColumns 用列元数据拼装 CREATE TABLE(类型直接用 ColumnInfo.Type)。
func buildDDLFromColumns(driver, table string, cols []ColumnInfo) (string, error) {
if len(cols) == 0 {
return "", fmt.Errorf("表 %s 不存在或无列", table)
}
var pks []string
for _, col := range cols {
if col.Key == "PRI" {
pks = append(pks, quoteIdentFor(driver, col.Name))
}
}
var b strings.Builder
fmt.Fprintf(&b, "CREATE TABLE %s (\n", quoteTableNameFor(driver, table))
for i, col := range cols {
line := " " + quoteIdentFor(driver, col.Name) + " " + col.Type
if !col.Nullable {
line += " NOT NULL"
}
// 非最后一列、或后面还有 PRIMARY KEY 子句时需要逗号
if i < len(cols)-1 || len(pks) > 0 {
line += ","
}
b.WriteString(line + "\n")
}
if len(pks) > 0 {
b.WriteString(" PRIMARY KEY (" + strings.Join(pks, ", ") + ")\n")
}
b.WriteString(");\n")
return b.String(), nil
}
+94
View File
@@ -0,0 +1,94 @@
package db
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
)
// newSQLiteFile 建一个 sqlite 测试库文件并执行建表/插数语句,返回 ConnReq。
func newSQLiteFile(t *testing.T, stmts ...string) ConnReq {
t.Helper()
path := filepath.Join(t.TempDir(), "test.db")
if err := os.WriteFile(path, []byte{}, 0600); err != nil {
t.Fatal(err)
}
for _, s := range stmts {
res := ExecSQL(context.Background(), ConnReq{Driver: "sqlite", Database: path, SQL: s})
if res.Error != "" {
t.Fatalf("初始化语句失败 %q: %s", s, res.Error)
}
}
return ConnReq{Driver: "sqlite", Database: path}
}
// R2-B5:sqlite 返回 sqlite_master 中的原始建表语句。
func TestGetDDLSQLite(t *testing.T) {
c := newSQLiteFile(t, `CREATE TABLE t_person (id INTEGER PRIMARY KEY, name TEXT NOT NULL, age INTEGER)`)
c.Table = "t_person"
ddl, err := GetDDL(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(ddl, "CREATE TABLE t_person") || !strings.Contains(ddl, "name TEXT") {
t.Fatalf("DDL 内容不符: %s", ddl)
}
c.Table = "no_such_table"
if _, err := GetDDL(context.Background(), c); err == nil {
t.Fatal("不存在的表应报错")
}
}
func TestGetDDLRedis(t *testing.T) {
_, err := GetDDL(context.Background(), ConnReq{Driver: "redis", Table: "k"})
if err == nil || !strings.Contains(err.Error(), "不支持") {
t.Fatalf("redis 应报不支持: %v", err)
}
}
func TestGetDDLMissingTable(t *testing.T) {
_, err := GetDDL(context.Background(), ConnReq{Driver: "sqlite", Table: ""})
if err == nil {
t.Fatal("空表名应报错")
}
}
// R2-B5:pg/mssql/oracle 族用列元数据拼装,各 family 引号不同。
func TestBuildDDLFromColumns(t *testing.T) {
cols := []ColumnInfo{
{Name: "id", Type: "integer", Nullable: false, Key: "PRI"},
{Name: "na\"me", Type: "text", Nullable: true},
}
// pg:双引号加倍
ddl, err := buildDDLFromColumns("postgres", "public.t", cols)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(ddl, `"public"."t"`) || !strings.Contains(ddl, `"na""me" text`) ||
!strings.Contains(ddl, `"id" integer NOT NULL,`) || !strings.Contains(ddl, `PRIMARY KEY ("id")`) {
t.Fatalf("pg DDL 不符:\n%s", ddl)
}
// mssql:中括号
ddl, err = buildDDLFromColumns("mssql", "dbo.t", cols)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(ddl, `[id] integer NOT NULL`) || !strings.Contains(ddl, `PRIMARY KEY ([id])`) {
t.Fatalf("mssql DDL 不符:\n%s", ddl)
}
// oracle:原名
ddl, err = buildDDLFromColumns("oracle", "T", cols)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(ddl, `CREATE TABLE T (`) || !strings.Contains(ddl, `PRIMARY KEY (id)`) {
t.Fatalf("oracle DDL 不符:\n%s", ddl)
}
// 无列报错
if _, err := buildDDLFromColumns("postgres", "t", nil); err == nil {
t.Fatal("无列应报错")
}
}
+193
View File
@@ -0,0 +1,193 @@
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
}
+252
View File
@@ -0,0 +1,252 @@
package db
import (
"bufio"
"bytes"
"context"
"database/sql"
"encoding/csv"
"encoding/json"
"fmt"
"io"
"strings"
"time"
)
const (
DefaultExportLimit = 100000
MaxExportLimit = 500000
exportFlushRows = 512
)
// ExportAllowedQuery 导出专用的只读校验:首个关键词须为 SELECT/WITH/SHOW/DESCRIBE/EXPLAIN。
func ExportAllowedQuery(sqlText string) bool {
s := stripSQLLead(strings.TrimSpace(sqlText))
if s == "" {
return false
}
up := strings.ToUpper(s)
for _, p := range []string{"SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN"} {
if strings.HasPrefix(up, p+" ") || up == p || strings.HasPrefix(up, p+"\t") || strings.HasPrefix(up, p+"\n") {
return true
}
}
return false
}
// SelectAllSQL 表模式导出的取数 SQL:按 family 引用(支持 schema.table 多段)。
func SelectAllSQL(driver, table string) string {
return "SELECT * FROM " + quoteTableNameFor(driver, table)
}
// ExportQuery 把查询结果按 c.Format(csv/json/insert)流式编码写入 w,避免全量载入内存。
// limit <= 0 用默认上限 100000,硬顶 MaxExportLimit。行值经 normalize 保真(同 ExecSQL)。
func ExportQuery(ctx context.Context, c ConnReq, w io.Writer) error {
ctx = reqCtx(ctx)
format := strings.ToLower(strings.TrimSpace(c.Format))
sqlText := strings.TrimSpace(c.SQL)
if sqlText == "" {
return fmt.Errorf("缺少导出 SQL")
}
limit := c.Limit
if limit <= 0 {
limit = DefaultExportLimit
}
if limit > MaxExportLimit {
limit = MaxExportLimit
}
return withDB(c, func(db *sql.DB) error {
rows, err := db.QueryContext(ctx, sqlText)
if err != nil {
return err
}
defer rows.Close()
cols, err := rows.Columns()
if err != nil {
return err
}
bw := bufio.NewWriterSize(w, 64*1024)
enc := newExportEncoder(bw, format, c.Driver, exportTargetTable(c), cols)
if err := enc.begin(); err != nil {
return err
}
n := 0
for rows.Next() {
if err := ctx.Err(); err != nil {
return 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 err
}
vals := make([]interface{}, len(cols))
for i := range raw {
vals[i] = normalize(raw[i])
}
if err := enc.row(vals); err != nil {
return err
}
n++
if n%exportFlushRows == 0 {
if err := bw.Flush(); err != nil {
return err
}
}
if n >= limit {
break
}
}
if err := rows.Err(); err != nil {
return err
}
if err := enc.end(); err != nil {
return err
}
return bw.Flush()
})
}
// exportTargetTable insert 格式的目标表名:显式 table 优先,否则用通用名。
func exportTargetTable(c ConnReq) string {
if strings.TrimSpace(c.Table) != "" {
return c.Table
}
return "result"
}
// exportEncoder 按格式把行流式编码进 buffered writer。
type exportEncoder struct {
bw *bufio.Writer
format string
driver string
table string
cols []string
csvW *csv.Writer
// jsonPending 上一行 JSON,延迟一行写出,使行间逗号紧跟行尾而非独占一行
jsonPending string
n int
}
func newExportEncoder(bw *bufio.Writer, format, driver, table string, cols []string) *exportEncoder {
return &exportEncoder{bw: bw, format: format, driver: driver, table: table, cols: cols}
}
func (e *exportEncoder) begin() error {
switch e.format {
case "csv":
e.csvW = csv.NewWriter(e.bw)
return e.csvW.Write(e.cols)
case "json":
_, err := e.bw.WriteString("[\n")
return err
case "insert":
return nil
default:
return fmt.Errorf("不支持的导出格式: %s", e.format)
}
}
func (e *exportEncoder) row(vals []interface{}) error {
e.n++
switch e.format {
case "csv":
rec := make([]string, len(vals))
for i, v := range vals {
rec[i] = csvCell(v)
}
return e.csvW.Write(rec)
case "json":
m := make(map[string]interface{}, len(e.cols))
for i, col := range e.cols {
m[col] = vals[i]
}
var rb bytes.Buffer
if err := json.NewEncoder(&rb).Encode(m); err != nil {
return err
}
line := strings.TrimSuffix(rb.String(), "\n")
if e.jsonPending != "" {
if _, err := e.bw.WriteString(e.jsonPending + ",\n"); err != nil {
return err
}
}
e.jsonPending = line
return nil
case "insert":
var b strings.Builder
b.WriteString("INSERT INTO ")
b.WriteString(quoteTableNameFor(e.driver, e.table))
b.WriteString(" (")
for i, col := range e.cols {
if i > 0 {
b.WriteString(", ")
}
b.WriteString(quoteIdentFor(e.driver, col))
}
b.WriteString(") VALUES (")
for i, v := range vals {
if i > 0 {
b.WriteString(", ")
}
b.WriteString(sqlLiteral(v))
}
b.WriteString(");\n")
_, err := e.bw.WriteString(b.String())
return err
}
return nil
}
func (e *exportEncoder) end() error {
switch e.format {
case "csv":
e.csvW.Flush()
return e.csvW.Error()
case "json":
if e.jsonPending != "" {
if _, err := e.bw.WriteString(e.jsonPending + "\n"); err != nil {
return err
}
e.jsonPending = ""
}
_, err := e.bw.WriteString("]\n")
return err
}
return nil
}
func csvCell(v interface{}) string {
switch t := v.(type) {
case nil:
return ""
case string:
return t
default:
return fmt.Sprintf("%v", t)
}
}
// sqlLiteral insert 字面量:NULL/数字原样,字符串单引号 doubling。
func sqlLiteral(v interface{}) string {
switch t := v.(type) {
case nil:
return "NULL"
case string:
return "'" + strings.ReplaceAll(t, "'", "''") + "'"
case bool:
if t {
return "1"
}
return "0"
case time.Time:
return "'" + t.Format(time.RFC3339Nano) + "'"
case []byte:
return "'" + strings.ReplaceAll(string(t), "'", "''") + "'"
default:
return fmt.Sprintf("%v", t)
}
}
+139
View File
@@ -0,0 +1,139 @@
package db
import (
"bytes"
"context"
"encoding/csv"
"encoding/json"
"strings"
"testing"
)
func newExportConn(t *testing.T) ConnReq {
return newSQLiteFile(t,
`CREATE TABLE t_export (id INTEGER PRIMARY KEY, name TEXT, note TEXT)`,
`INSERT INTO t_export (id, name, note) VALUES (1, 'alice', 'hi, "q"')`,
`INSERT INTO t_export (id, name, note) VALUES (2, 'bo''b', NULL)`,
`INSERT INTO t_export (id, name, note) VALUES (3, 'carol', 'line1`+"\n"+`line2')`,
)
}
// R2-B7:csv 输出形状(表头 + csv 引号转义由 encoding/csv 负责)。
func TestExportCSV(t *testing.T) {
c := newExportConn(t)
c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
c.Format = "csv"
var buf bytes.Buffer
if err := ExportQuery(context.Background(), c, &buf); err != nil {
t.Fatal(err)
}
out := buf.String()
if !strings.HasPrefix(out, "id,name,note\n") {
t.Fatalf("csv 表头错误: %q", out)
}
// 含逗号/引号的字段整体加引号、内部引号 doubling
if !strings.Contains(out, `"hi, ""q"""`) {
t.Fatalf("csv 转义错误: %q", out)
}
// 单引号 doubling 是 SQL 字面量层;csv 层原样输出 'bo''b' 存好的 bo'b
if !strings.Contains(out, "2,bo'b,") {
t.Fatalf("csv 内容错误: %q", out)
}
// 多行字面量会使单条记录跨物理行,用 csv.Reader 按记录解析
recs, err := csv.NewReader(strings.NewReader(out)).ReadAll()
if err != nil {
t.Fatalf("csv 输出应可解析: %v", err)
}
if len(recs) != 4 { // 表头 + 3 行
t.Fatalf("应为表头+3 行: %q", out)
}
}
func TestExportJSON(t *testing.T) {
c := newExportConn(t)
c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
c.Format = "json"
var buf bytes.Buffer
if err := ExportQuery(context.Background(), c, &buf); err != nil {
t.Fatal(err)
}
var rows []map[string]interface{}
if err := json.Unmarshal(buf.Bytes(), &rows); err != nil {
t.Fatalf("json 输出应可解析: %v\n%s", err, buf.String())
}
if len(rows) != 3 {
t.Fatalf("应为 3 行: %v", rows)
}
if rows[0]["name"] != "alice" || rows[1]["note"] != nil {
t.Fatalf("行值错误: %v", rows)
}
}
func TestExportInsert(t *testing.T) {
c := newExportConn(t)
c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
c.Format = "insert"
c.Table = "t_export"
var buf bytes.Buffer
if err := ExportQuery(context.Background(), c, &buf); err != nil {
t.Fatal(err)
}
out := buf.String()
if !strings.Contains(out, `INSERT INTO "t_export" ("id", "name", "note") VALUES (1, 'alice', 'hi, "q"');`) {
t.Fatalf("insert 首行错误: %q", out)
}
// 字符串单引号 doubling;NULL 原样
if !strings.Contains(out, `VALUES (2, 'bo''b', NULL);`) {
t.Fatalf("insert 字面量错误: %q", out)
}
// 多行字面量内换行原样保留
if !strings.Contains(out, "line1\nline2") {
t.Fatalf("insert 换行字面量错误: %q", out)
}
}
// limit 生效(含默认上限收敛)。
func TestExportLimit(t *testing.T) {
c := newExportConn(t)
c.SQL = `SELECT id FROM t_export ORDER BY id`
c.Format = "csv"
c.Limit = 2
var buf bytes.Buffer
if err := ExportQuery(context.Background(), c, &buf); err != nil {
t.Fatal(err)
}
rows := strings.Split(strings.TrimSpace(buf.String()), "\n")
if len(rows) != 3 { // 表头 + 2 行
t.Fatalf("limit=2 应只导 2 行: %q", buf.String())
}
}
func TestExportAllowedQuery(t *testing.T) {
for _, ok := range []string{
"SELECT * FROM t", "WITH x AS (SELECT 1) SELECT * FROM x",
"SHOW TABLES", "DESCRIBE t", "EXPLAIN SELECT 1",
"-- c\nSELECT 1",
} {
if !ExportAllowedQuery(ok) {
t.Fatalf("应允许只读: %q", ok)
}
}
for _, bad := range []string{
"DELETE FROM t", "INSERT INTO t VALUES (1)", "DROP TABLE t",
"UPDATE t SET a=1", "", " ",
} {
if ExportAllowedQuery(bad) {
t.Fatalf("不应允许: %q", bad)
}
}
}
func TestExportBadFormat(t *testing.T) {
c := newExportConn(t)
c.SQL = "SELECT 1"
c.Format = "xml"
var buf bytes.Buffer
if err := ExportQuery(context.Background(), c, &buf); err == nil {
t.Fatal("未知格式应报错")
}
}
+72
View File
@@ -0,0 +1,72 @@
package db
// family_e2e_test.go —— 协议兼容族真实例验收(MySQL 族:TiDB;PG 族:openGauss)。
// 触发:OPSXDB_E2E_TIDB=1 / OPSXDB_E2E_GAUSS=1;验证族内真实例后注册表才可翻 verified。
import (
"context"
"os"
"testing"
"time"
)
// TestE2ETiDB —— 真实 TiDB(pingcap/tidb 单容器,MySQL 协议族)。
func TestE2ETiDB(t *testing.T) {
if os.Getenv("OPSXDB_E2E_TIDB") != "1" {
t.Skip("置 OPSXDB_E2E_TIDB=1 并启动 TiDB :4000 后运行(docker run -p 4000:4000 pingcap/tidb:v7.5.0)")
}
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
c := ConnReq{Driver: "tidb", Host: "127.0.0.1", Port: 4000, User: "root"}
if err := TestConnection(c); err != nil {
t.Skipf("TiDB :4000 不可达: %v", err)
}
dbs, err := ListDatabases(ctx, c)
if err != nil || len(dbs) == 0 {
t.Fatalf("ListDatabases = %v, %v", dbs, err)
}
// 建库建表写读
must := func(sql string, db string) ExecResult {
t.Helper()
return ExecSQL(ctx, ConnReq{Driver: "tidb", Host: "127.0.0.1", Port: 4000, User: "root", Database: db, SQL: sql, ConfirmDanger: true})
}
if r := must("CREATE DATABASE IF NOT EXISTS opsxe2e", ""); r.Error != "" {
t.Fatalf("create db: %s", r.Error)
}
if r := must("CREATE TABLE IF NOT EXISTS fam_t(id INT PRIMARY KEY, name VARCHAR(32))", "opsxe2e"); r.Error != "" {
t.Fatalf("create table: %s", r.Error)
}
if r := must(`INSERT INTO fam_t VALUES (1, 'tidb-ok')`, "opsxe2e"); r.Error != "" {
t.Fatalf("insert: %s", r.Error)
}
r := must("SELECT * FROM fam_t", "opsxe2e")
if r.Error != "" || r.RowCount != 1 || r.Rows[0]["name"] != "tidb-ok" {
t.Fatalf("select: err=%s rows=%+v", r.Error, r.Rows)
}
}
// TestE2EOpenGauss —— 真实 openGauss(PG 协议族;Apple Silicon 需 --platform linux/amd64)。
func TestE2EOpenGauss(t *testing.T) {
if os.Getenv("OPSXDB_E2E_GAUSS") != "1" {
t.Skip("置 OPSXDB_E2E_GAUSS=1 并启动 openGauss :5432 后运行")
}
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
c := ConnReq{Driver: "gaussdb", Host: "127.0.0.1", Port: 5432, User: "omm", Password: "Demo@123", Database: "postgres"}
if err := TestConnection(c); err != nil {
t.Skipf("openGauss :5432 不可达: %v", err)
}
dbs, err := ListDatabases(ctx, c)
if err != nil || len(dbs) == 0 {
t.Fatalf("ListDatabases = %v, %v", dbs, err)
}
q := ExecSQL(ctx, ConnReq{Driver: "gaussdb", Host: "127.0.0.1", Port: 5432, User: "omm", Password: "Demo@123", Database: "postgres", SQL: "SELECT 1 AS one"})
if q.Error != "" {
t.Fatalf("select: %s", q.Error)
}
if q.RowCount != 1 {
t.Fatalf("rows = %+v", q.Rows)
}
}
+49
View File
@@ -0,0 +1,49 @@
package db
import "strings"
// quoteMySQLIdent 反引号标识符转义:内部反引号加倍(“),与前端 sql.ts quoteIdent 对齐。
func quoteMySQLIdent(s string) string {
return "`" + strings.ReplaceAll(s, "`", "``") + "`"
}
func quotePGIdent(s string) string {
return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
}
// quoteIdentFor 按驱动族引用标识符:mysql 反引号加倍、pg/sqlite 双引号加倍、
// mssql 中括号加倍、oracle 沿用原名(大写约定)。
func quoteIdentFor(driver, name string) string {
switch {
case isMySQLFamily(driver):
return quoteMySQLIdent(name)
case isMSSQL(driver):
return "[" + strings.ReplaceAll(name, "]", "]]") + "]"
case driverKey(driver) == "oracle":
return name
case driverKey(driver) == "sqlite":
return quoteSQLiteIdent(name)
default:
return quotePGIdent(name)
}
}
// quoteTableNameFor 引用表名(支持 schema.table 多段,oracle 除外沿用原名)。
func quoteTableNameFor(driver, table string) string {
if driverKey(driver) == "oracle" {
return table
}
parts := strings.Split(table, ".")
for i := range parts {
parts[i] = quoteIdentFor(driver, parts[i])
}
return strings.Join(parts, ".")
}
func splitSchemaTable(name string) (schema, table string) {
name = strings.TrimSpace(name)
if i := strings.LastIndex(name, "."); i > 0 {
return name[:i], name[i+1:]
}
return "", name
}
+46
View File
@@ -0,0 +1,46 @@
package db
import (
"context"
"sync"
)
// inflightCancels 在执行的请求注册表:requestId → cancel,供 /api/db/cancel 远程取消。
// 请求完成(正常/异常)必须 UnregisterInflight,避免 map 泄漏。
var inflightCancels = struct {
sync.Mutex
m map[string]context.CancelFunc
}{m: make(map[string]context.CancelFunc)}
// RegisterInflight 注册一个可按 requestId 取消的在执行请求。
// requestID 为空或 cancel 为 nil 时忽略。
func RegisterInflight(requestID string, cancel context.CancelFunc) {
if requestID == "" || cancel == nil {
return
}
inflightCancels.Lock()
defer inflightCancels.Unlock()
inflightCancels.m[requestID] = cancel
}
// UnregisterInflight 请求结束时注销(幂等)。
func UnregisterInflight(requestID string) {
inflightCancels.Lock()
defer inflightCancels.Unlock()
delete(inflightCancels.m, requestID)
}
// CancelInflight 按 requestId 取消在执行的请求;不存在则返回 false。
// 取消后即注销,重复取消返回 false。
func CancelInflight(requestID string) bool {
inflightCancels.Lock()
cancel, ok := inflightCancels.m[requestID]
if ok {
delete(inflightCancels.m, requestID)
}
inflightCancels.Unlock()
if ok {
cancel()
}
return ok
}
+51
View File
@@ -0,0 +1,51 @@
package db
import (
"context"
"testing"
)
// R2-B4:注册→取消→CancelFunc 被调用;注销后 CancelInflight 返回 false。
func TestInflightCancel(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
const id = "req-cancel-1"
RegisterInflight(id, cancel)
if !CancelInflight(id) {
t.Fatal("已注册应能取消")
}
if ctx.Err() != context.Canceled {
t.Fatal("CancelFunc 应已被调用")
}
// 取消后即注销,重复取消返回 false
if CancelInflight(id) {
t.Fatal("注销后 CancelInflight 应返回 false")
}
}
func TestInflightUnregister(t *testing.T) {
_, cancel := context.WithCancel(context.Background())
defer cancel()
const id = "req-unregister-1"
RegisterInflight(id, cancel)
UnregisterInflight(id)
if CancelInflight(id) {
t.Fatal("注销后不应再取消到")
}
// 幂等:重复注销不 panic
UnregisterInflight(id)
}
// 空 id / nil cancel 注册应被忽略。
func TestInflightRegisterGuards(t *testing.T) {
RegisterInflight("", func() {})
if CancelInflight("") {
t.Fatal("空 requestId 不应注册成功")
}
RegisterInflight("req-nil-cancel", nil)
if CancelInflight("req-nil-cancel") {
t.Fatal("nil cancel 不应注册成功")
}
}
+271
View File
@@ -0,0 +1,271 @@
package db
// influx.go —— InfluxDB 2.x 原生接入(HTTP API,零第三方依赖)。
//
// 字段映射(连接表单):
//
// 主机/端口 → 服务地址(默认 :8086)
// 用户 → org(如 opsx)
// 密码 → admin token(如 opsxdb-demo-token)
// 库名 → bucket(如 demo)
//
// 能力:对象树 bucket → measurement(Flux schema.measurements);
// SQL 编辑器直写 Flux(from(bucket: ...) |> range(...) ...),结果按注解 CSV 解析为表格。
// Flux 只读(写走 /write 专用通道),危险语句拦截在 ops.go 入口统一过一遍。
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
)
func isInflux(d string) bool { return driverKey(d) == "influxdb" }
func ifxBase(c ConnReq) string {
host := strings.TrimSpace(c.Host)
if host == "" {
host = "127.0.0.1"
}
port := c.Port
if port == 0 {
port = 8086
}
return fmt.Sprintf("http://%s:%d", host, port)
}
func ifxToken(c ConnReq) string { return c.Password }
// ifxDo 统一请求:Token 鉴权 + org 查询参数。
func ifxDo(ctx context.Context, c ConnReq, method, path string, body io.Reader, out any) error {
req, err := http.NewRequestWithContext(ctx, method, ifxBase(c)+path, body)
if err != nil {
return err
}
if ifxToken(c) != "" {
req.Header.Set("Authorization", "Token "+ifxToken(c))
}
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
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("InfluxDB 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 ifxOrg(c ConnReq) (string, error) {
org := strings.TrimSpace(c.User)
if org == "" {
return "", fmt.Errorf("InfluxDB 2.x 需要填「用户」= org 名(如 opsx)")
}
return org, nil
}
func ifxPing(c ConnReq) error {
org, err := ifxOrg(c)
if err != nil {
return err
}
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
defer cancel()
var out struct {
Buckets []struct {
Name string `json:"name"`
} `json:"buckets"`
}
return ifxDo(ctx, c, http.MethodGet, "/api/v2/buckets?org="+org, nil, &out)
}
func ifxListDatabases(ctx context.Context, c ConnReq) ([]string, error) {
org, err := ifxOrg(c)
if err != nil {
return nil, err
}
var out struct {
Buckets []struct {
Name string `json:"name"`
} `json:"buckets"`
}
if err := ifxDo(ctx, c, http.MethodGet, "/api/v2/buckets?org="+org, nil, &out); err != nil {
return nil, err
}
names := []string{}
for _, b := range out.Buckets {
names = append(names, b.Name)
}
return names, nil
}
// ifxQueryFlux 跑 Flux 并把注解 CSV 解析为表格。
func ifxQueryFlux(ctx context.Context, c ConnReq, flux string, maxRows int) (ExecResult, error) {
org, err := ifxOrg(c)
if err != nil {
return ExecResult{}, err
}
body := map[string]string{"query": flux, "type": "flux"}
b, err := json.Marshal(body)
if err != nil {
return ExecResult{}, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, ifxBase(c)+"/api/v2/query?org="+org, strings.NewReader(string(b)))
if err != nil {
return ExecResult{}, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/csv")
if ifxToken(c) != "" {
req.Header.Set("Authorization", "Token "+ifxToken(c))
}
resp, err := (&http.Client{Timeout: 60 * time.Second}).Do(req)
if err != nil {
return ExecResult{}, err
}
defer resp.Body.Close()
raw, err := io.ReadAll(resp.Body)
if err != nil {
return ExecResult{}, err
}
if resp.StatusCode >= 400 {
return ExecResult{}, fmt.Errorf("InfluxDB HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
}
return parseAnnotatedCSV(string(raw), maxRows)
}
// parseAnnotatedCSV 解析 InfluxDB 注解 CSV:
// #group/#datatype/#default 注释行 + 空行 + header 行 + 数据行。
func parseAnnotatedCSV(text string, maxRows int) (ExecResult, error) {
lines := strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n")
headerIdx := -1
var header []string
for i, ln := range lines {
s := strings.TrimSpace(ln)
if s == "" || strings.HasPrefix(s, "#") {
continue
}
headerIdx = i
header = strings.Split(ln, ",")
break
}
if headerIdx < 0 {
return ExecResult{Columns: []string{}, Rows: []map[string]interface{}{}}, nil
}
res := ExecResult{Columns: header, Rows: []map[string]interface{}{}}
for _, ln := range lines[headerIdx+1:] {
if strings.TrimSpace(ln) == "" {
continue
}
cells := strings.Split(ln, ",")
row := make(map[string]interface{}, len(header))
for i, col := range header {
v := ""
if i < len(cells) {
v = cells[i]
}
row[col] = v
}
res.Rows = append(res.Rows, row)
if len(res.Rows) >= maxRows {
res.Truncated = true
break
}
}
res.RowCount = len(res.Rows)
return res, nil
}
func ifxListTables(ctx context.Context, c ConnReq) ([]string, error) {
bucket := strings.TrimSpace(c.Database)
if bucket == "" {
return nil, fmt.Errorf("InfluxDB 列表需要库名(bucket)")
}
flux := fmt.Sprintf("import \"influxdata/influxdb/schema\"\nschema.measurements(bucket: \"%s\")", strings.ReplaceAll(bucket, "\"", ""))
res, err := ifxQueryFlux(ctx, c, flux, 10000)
if err != nil {
return nil, err
}
out := []string{}
if len(res.Columns) == 0 {
return out, nil
}
// schema.measurements 输出列为 [,result,table,_value]:优先 _value 列,退化为最后一列
idx := len(res.Columns) - 1
for i, col := range res.Columns {
if col == "_value" {
idx = i
break
}
}
col := res.Columns[idx]
for _, row := range res.Rows {
if v, ok := row[col].(string); ok && v != "" {
out = append(out, v)
}
}
return out, nil
}
// ifxExecSQL 编辑器内容按 Flux 处理(只读)。
func ifxExecSQL(ctx context.Context, c ConnReq, start time.Time) ExecResult {
if c.Offset > 0 {
return ExecResult{Columns: []string{}, Rows: []map[string]interface{}{}, Error: "InfluxDB Flux 结果暂不支持翻页(可缩小 range 重查)", TimeMs: time.Since(start).Milliseconds()}
}
limit := rowLimit(c)
if limit > MaxRowLimit {
limit = MaxRowLimit
}
res, err := ifxQueryFlux(ctx, c, strings.TrimSpace(c.SQL), limit)
if err != nil {
return execFail(start, err)
}
res.TimeMs = time.Since(start).Milliseconds()
return res
}
// ifxListColumns measurement 的字段探查:Flux schema.fieldKeys + tagKeys 合并。
func ifxListColumns(ctx context.Context, c ConnReq) ([]ColumnInfo, error) {
bucket := strings.TrimSpace(c.Database)
table := strings.TrimSpace(c.Table)
if bucket == "" || table == "" {
return nil, fmt.Errorf("InfluxDB 列探查需要库名与 measurement")
}
safeB := strings.ReplaceAll(bucket, "\"", "")
safeM := strings.ReplaceAll(table, "\"", "")
flux := fmt.Sprintf(`import "influxdata/influxdb/schema"
union(tables: [
schema.fieldKeys(bucket: "%s", predicate: (r) => r._measurement == "%s"),
schema.tagKeys(bucket: "%s", predicate: (r) => r._measurement == "%s")
])`, safeB, safeM, safeB, safeM)
res, err := ifxQueryFlux(ctx, c, flux, 1000)
if err != nil {
return nil, err
}
cols := []ColumnInfo{}
seen := map[string]bool{}
if len(res.Columns) == 0 {
return cols, nil
}
last := res.Columns[len(res.Columns)-1]
for _, row := range res.Rows {
name, _ := row[last].(string)
if name == "" || seen[name] {
continue
}
seen[name] = true
cols = append(cols, ColumnInfo{Name: name, Type: "tag/field", Nullable: true})
}
return cols, nil
}
+160
View File
@@ -0,0 +1,160 @@
package db
import (
"context"
"database/sql"
"fmt"
"net/url"
"strings"
_ "github.com/microsoft/go-mssqldb"
)
func openMSSQL(c ConnReq) (*sql.DB, error) {
dsn, err := mssqlDSN(c)
if err != nil {
return nil, err
}
return sql.Open("sqlserver", dsn)
}
// mssqlDSN 拼 sqlserver:// DSN;账号密码走 url.UserPassword 编码,
// 修复含空格 / '+' 等特殊字符密码连接失败的问题(url.QueryEscape 不适用 userinfo)。
func mssqlDSN(c ConnReq) (string, error) {
host := c.Host
if host == "" {
host = "127.0.0.1"
}
port := c.Port
if port == 0 {
port = 1433
}
user := c.User
if user == "" {
user = "sa"
}
dbName := c.Database
if dbName == "" {
dbName = "master"
}
mode := EffectiveSSLMode(c)
q := url.Values{}
q.Set("database", dbName)
q.Set("connection timeout", "5")
switch mode {
case "disable":
q.Set("encrypt", "disable")
case "skip-verify":
q.Set("encrypt", "true")
q.Set("TrustServerCertificate", "true")
if ca := strings.TrimSpace(c.SSLCA); ca != "" {
q.Set("certificate", ca)
}
case "require", "verify-ca":
q.Set("encrypt", "true")
q.Set("TrustServerCertificate", "false")
ca := strings.TrimSpace(c.SSLCA)
if mode == "verify-ca" && ca == "" {
return "", fmt.Errorf("verify-ca 模式需要填写 CA 证书路径")
}
if ca != "" {
// go-mssqldb msdsn certificate= builds RootCAs from PEM/DER
q.Set("certificate", ca)
}
if host != "" {
q.Set("hostnameincertificate", host)
}
default:
return "", fmt.Errorf("未知 SSL 模式: %s", mode)
}
u := &url.URL{
Scheme: "sqlserver",
User: url.UserPassword(user, c.Password),
Host: fmt.Sprintf("%s:%d", host, port),
RawQuery: q.Encode(),
}
return u.String(), nil
}
func listDatabasesMSSQL(ctx context.Context, db *sql.DB) ([]string, error) {
rows, err := db.QueryContext(ctx, `SELECT name FROM sys.databases WHERE name NOT IN ('master','tempdb','model','msdb') ORDER BY name`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listTablesMSSQL(ctx context.Context, db *sql.DB, database string) ([]string, error) {
_ = database
rows, err := db.QueryContext(ctx, `
SELECT TABLE_SCHEMA + '.' + TABLE_NAME
FROM INFORMATION_SCHEMA.TABLES
WHERE TABLE_TYPE = 'BASE TABLE'
ORDER BY 1`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listColumnsMSSQL(ctx context.Context, dbh *sql.DB, _, table string) ([]ColumnInfo, error) {
schema, t := splitSchemaTable(table)
if schema == "" {
schema = "dbo"
}
// 关联 sys.indexes / sys.index_columns 补主键标记(命中则 Key='PRI')
rows, err := dbh.QueryContext(ctx, `
SELECT c.COLUMN_NAME, c.DATA_TYPE,
ISNULL(CAST(c.CHARACTER_MAXIMUM_LENGTH AS varchar), ''),
c.IS_NULLABLE, c.COLUMN_DEFAULT,
CASE WHEN pk.name IS NOT NULL THEN 'PRI' ELSE '' END
FROM INFORMATION_SCHEMA.COLUMNS c
LEFT JOIN sys.tables t ON t.name = c.TABLE_NAME AND SCHEMA_NAME(t.schema_id) = c.TABLE_SCHEMA
LEFT JOIN sys.indexes i ON i.object_id = t.object_id AND i.is_primary_key = 1
LEFT JOIN sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id
LEFT JOIN sys.columns pk ON pk.object_id = ic.object_id AND pk.column_id = ic.column_id AND pk.name = c.COLUMN_NAME
WHERE c.TABLE_SCHEMA = @p1 AND c.TABLE_NAME = @p2
ORDER BY c.ORDINAL_POSITION`, schema, t)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ColumnInfo
for rows.Next() {
var name, typ, extra, nullable, key string
var def sql.NullString
if err := rows.Scan(&name, &typ, &extra, &nullable, &def, &key); err != nil {
return nil, err
}
if extra != "" && extra != "-1" {
typ = typ + "(" + extra + ")"
}
col := ColumnInfo{Name: name, Type: typ, Nullable: nullable == "YES", Key: key}
if def.Valid {
col.Default = def.String
}
out = append(out, col)
}
if out == nil {
out = []ColumnInfo{}
}
return out, rows.Err()
}
+165
View File
@@ -0,0 +1,165 @@
package db
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"fmt"
"net"
"strings"
"time"
mysqldriver "github.com/go-sql-driver/mysql"
)
func openMySQL(c ConnReq) (*sql.DB, error) {
host := c.Host
if host == "" {
host = "127.0.0.1"
}
port := c.Port
if port == 0 {
port = 3306
}
user := c.User
if user == "" {
user = "root"
}
cfg := mysqldriver.NewConfig()
cfg.User = user
cfg.Passwd = c.Password
cfg.Net = "tcp"
cfg.Addr = net.JoinHostPort(host, fmt.Sprintf("%d", port))
cfg.DBName = c.Database
cfg.ParseTime = true
cfg.Timeout = 5 * time.Second
cfg.ReadTimeout = 30 * time.Second
cfg.WriteTimeout = 30 * time.Second
cfg.Loc = time.Local
cfg.Params = map[string]string{"charset": "utf8mb4"}
mode := EffectiveSSLMode(c)
switch mode {
case "disable":
// no TLS
case "skip-verify":
if !hasTLSMaterial(c) {
cfg.TLSConfig = "skip-verify"
} else {
name, err := registerMySQLTLS(c)
if err != nil {
return nil, err
}
cfg.TLSConfig = name
}
case "require", "verify-ca":
if !hasTLSMaterial(c) && mode == "require" {
cfg.TLSConfig = "true"
} else {
name, err := registerMySQLTLS(c)
if err != nil {
return nil, err
}
cfg.TLSConfig = name
}
default:
return nil, fmt.Errorf("未知 SSL 模式: %s", mode)
}
return sql.Open("mysql", cfg.FormatDSN())
}
// hasTLSMaterial 是否提供了任何 CA/客户端证书材料(路径或 PEM 内容)。
func hasTLSMaterial(c ConnReq) bool {
return strings.TrimSpace(c.SSLCA) != "" ||
strings.TrimSpace(c.SSLCert) != "" ||
strings.TrimSpace(c.SSLKey) != "" ||
strings.TrimSpace(c.SSLCAContent) != "" ||
strings.TrimSpace(c.SSLCertContent) != "" ||
strings.TrimSpace(c.SSLKeyContent) != ""
}
func registerMySQLTLS(c ConnReq) (string, error) {
tlsCfg, err := tlsConfigFrom(c)
if err != nil {
return "", err
}
if tlsCfg == nil {
return "", fmt.Errorf("SSL 已启用但未能构建 TLS 配置")
}
sum := sha256.Sum256([]byte(EffectiveSSLMode(c) + "|" + c.SSLCA + "|" + c.SSLCert + "|" + c.SSLKey +
"|" + c.SSLCAContent + "|" + c.SSLCertContent + "|" + c.SSLKeyContent + "|" + c.TLSServerName + "|" + c.Host))
name := "opsxdb-" + hex.EncodeToString(sum[:8])
mysqldriver.DeregisterTLSConfig(name)
if err := mysqldriver.RegisterTLSConfig(name, tlsCfg); err != nil {
return "", fmt.Errorf("注册 MySQL TLS 配置失败: %w", err)
}
return name, nil
}
func listDatabasesMySQL(ctx context.Context, db *sql.DB) ([]string, error) {
rows, err := db.QueryContext(ctx, "SHOW DATABASES")
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listTablesMySQL(ctx context.Context, db *sql.DB, database string) ([]string, error) {
q := "SHOW TABLES"
if database != "" {
q = "SHOW TABLES FROM " + quoteMySQLIdent(database)
}
rows, err := db.QueryContext(ctx, q)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listColumnsMySQL(ctx context.Context, db *sql.DB, database, table string) ([]ColumnInfo, error) {
ident := quoteMySQLIdent(table)
if database != "" {
ident = quoteMySQLIdent(database) + "." + ident
}
rows, err := db.QueryContext(ctx, "SHOW FULL COLUMNS FROM "+ident)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ColumnInfo
for rows.Next() {
var field, typ, null, key string
var collation, extra, priv, comment sql.NullString
var def sql.NullString
if err := rows.Scan(&field, &typ, &collation, &null, &key, &def, &extra, &priv, &comment); err != nil {
return nil, err
}
col := ColumnInfo{Name: field, Type: typ, Nullable: null == "YES", Key: key, Comment: comment.String}
if def.Valid {
col.Default = def.String
}
out = append(out, col)
}
if out == nil {
out = []ColumnInfo{}
}
return out, rows.Err()
}
+114
View File
@@ -0,0 +1,114 @@
package db
import (
"database/sql"
"fmt"
"strings"
"time"
)
func driverKey(d string) string {
return strings.ToLower(strings.TrimSpace(d))
}
func isMySQLFamily(d string) bool {
switch driverKey(d) {
case "mysql", "mariadb", "tidb", "oceanbase", "polardb", "tdsql":
return true
default:
return false
}
}
func isPostgresFamily(d string) bool {
switch driverKey(d) {
case "postgresql", "postgres", "gaussdb", "opengauss", "kingbase":
return true
default:
return false
}
}
func isMSSQL(d string) bool {
switch driverKey(d) {
case "sqlserver", "mssql":
return true
default:
return false
}
}
// Open 按驱动打开 *sql.DB
func Open(c ConnReq) (*sql.DB, error) {
var (
db *sql.DB
err error
)
switch {
case isMySQLFamily(c.Driver):
db, err = openMySQL(c)
case isPostgresFamily(c.Driver):
db, err = openPostgres(c)
case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
db, err = openMSSQL(c)
case driverKey(c.Driver) == "oracle":
db, err = openOracle(c)
case driverKey(c.Driver) == "sqlite":
db, err = openSQLite(c)
default:
return nil, fmt.Errorf("驱动「%s」尚未实现(当前支持: mysql/mariadb/tidb、postgresql、sqlserver、oracle、sqlite、redis)", c.Driver)
}
if err != nil {
return nil, err
}
db.SetConnMaxLifetime(2 * time.Minute)
db.SetConnMaxIdleTime(5 * time.Minute) // 防空闲连接堆积
db.SetMaxOpenConns(5)
db.SetMaxIdleConns(2)
return db, nil
}
func TestConnection(c ConnReq) error {
if isAgentDriver(c.Driver) {
// JDBC agent 桥的 TCP 由 JVM 侧发起,SSH 隧道(本地端口转发)对 agent 不生效,明确拒绝
if c.SSHEnabled {
return fmt.Errorf("%s 走 JDBC agent,暂不支持 SSH 隧道", c.Driver)
}
if err := agentPing(c); err != nil {
return fmt.Errorf("数据库: %w", err)
}
return nil
}
dial, tun, err := ApplySSHTunnel(c)
if err != nil {
return err
}
if tun != nil {
defer tun.Close()
}
if err := testDial(dial); err != nil {
return fmt.Errorf("数据库: %w", err)
}
return nil
}
func testDial(c ConnReq) error {
if isRedis(c.Driver) {
return testRedis(c)
}
if isES(c.Driver) {
return esPing(c)
}
if isInflux(c.Driver) {
return ifxPing(c)
}
if isAgentDriver(c.Driver) {
return agentPing(c)
}
db, err := Open(c)
if err != nil {
return err
}
defer db.Close()
return db.Ping()
}
+506
View File
@@ -0,0 +1,506 @@
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
}
+39
View File
@@ -0,0 +1,39 @@
package db
import (
"context"
"testing"
)
func TestIsQuerySQL(t *testing.T) {
if !isQuerySQL("SELECT 1") {
t.Fatal("select")
}
if isQuerySQL("DELETE FROM t") {
t.Fatal("delete is not query")
}
if !isQuerySQL("-- hi\nSELECT 1") {
t.Fatal("comment then select")
}
}
func TestIsDangerousSQL(t *testing.T) {
if !IsDangerousSQL("DROP TABLE users") {
t.Fatal("drop")
}
if !IsDangerousSQL("DELETE FROM users") {
t.Fatal("delete no where")
}
if IsDangerousSQL("DELETE FROM users WHERE id=1") {
t.Fatal("delete with where is ok")
}
if IsDangerousSQL("SELECT * FROM users") {
t.Fatal("select")
}
}
func TestErrMsgCanceled(t *testing.T) {
if got := errMsg(context.Canceled); got != "已取消" {
t.Fatalf("got %q", got)
}
}
+157
View File
@@ -0,0 +1,157 @@
package db
import (
"context"
"database/sql"
"fmt"
"net/url"
ora "github.com/sijms/go-ora/v2"
)
func openOracle(c ConnReq) (*sql.DB, error) {
host := c.Host
if host == "" {
host = "127.0.0.1"
}
port := c.Port
if port == 0 {
port = 1521
}
user := c.User
if user == "" {
user = "system"
}
service := c.Service
if service == "" {
if c.Database != "" {
service = c.Database
} else {
service = "XEPDB1"
}
}
dsn := fmt.Sprintf("oracle://%s:%s@%s:%d/%s", url.PathEscape(user), url.PathEscape(c.Password), host, port, service)
mode := EffectiveSSLMode(c)
if mode != "disable" {
q := url.Values{}
q.Set("SSL", "true")
if mode == "skip-verify" {
q.Set("SSL VERIFY", "false")
} else {
q.Set("SSL VERIFY", "true")
}
dsn = dsn + "?" + q.Encode()
}
tlsCfg, err := tlsConfigFrom(c)
if err != nil {
return nil, err
}
if tlsCfg != nil {
connector := ora.NewConnector(dsn)
oc, ok := connector.(*ora.OracleConnector)
if !ok {
return nil, fmt.Errorf("oracle: connector 类型异常")
}
oc.WithTLSConfig(tlsCfg)
return sql.OpenDB(oc), nil
}
return sql.Open("oracle", dsn)
}
func listDatabasesOracle(ctx context.Context, db *sql.DB) ([]string, error) {
// 返回当前用户可见 schema
rows, err := db.QueryContext(ctx, `SELECT username FROM all_users ORDER BY username`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listTablesOracle(ctx context.Context, db *sql.DB, schema string) ([]string, error) {
var (
rows *sql.Rows
err error
)
if schema == "" {
rows, err = db.QueryContext(ctx, `SELECT table_name FROM user_tables ORDER BY table_name`)
} else {
rows, err = db.QueryContext(ctx, `SELECT table_name FROM all_tables WHERE owner = UPPER(:1) ORDER BY table_name`, schema)
}
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listColumnsOracle(ctx context.Context, dbh *sql.DB, schema, table string) ([]ColumnInfo, error) {
owner, t := splitSchemaTable(table)
if owner == "" {
owner = schema
}
var (
rows *sql.Rows
err error
)
// 关联约束元数据补主键标记(命中则 Key='PRI');沿用既有 UPPER() 绑定约定
if owner == "" {
rows, err = dbh.QueryContext(ctx, `
SELECT c.column_name, c.data_type, c.nullable, c.data_default,
CASE WHEN acc.column_name IS NOT NULL THEN 'PRI' ELSE '' END
FROM user_tab_columns c
LEFT JOIN user_constraints ac ON ac.table_name = UPPER(:1) AND ac.constraint_type = 'P'
LEFT JOIN user_cons_columns acc ON acc.constraint_name = ac.constraint_name
AND acc.table_name = ac.table_name AND acc.column_name = c.column_name
WHERE c.table_name = UPPER(:2)
ORDER BY c.column_id`, t, t)
} else {
rows, err = dbh.QueryContext(ctx, `
SELECT c.column_name, c.data_type, c.nullable, c.data_default,
CASE WHEN acc.column_name IS NOT NULL THEN 'PRI' ELSE '' END
FROM all_tab_columns c
LEFT JOIN all_constraints ac ON ac.owner = UPPER(:1) AND ac.table_name = UPPER(:2) AND ac.constraint_type = 'P'
LEFT JOIN all_cons_columns acc ON acc.owner = ac.owner AND acc.constraint_name = ac.constraint_name
AND acc.column_name = c.column_name
WHERE c.owner = UPPER(:3) AND c.table_name = UPPER(:4)
ORDER BY c.column_id`, owner, t, owner, t)
}
if err != nil {
return nil, err
}
defer rows.Close()
var out []ColumnInfo
for rows.Next() {
var name, typ, nullable, key string
var def sql.NullString
if err := rows.Scan(&name, &typ, &nullable, &def, &key); err != nil {
return nil, err
}
col := ColumnInfo{Name: name, Type: typ, Nullable: nullable == "Y", Key: key}
if def.Valid {
col.Default = def.String
}
out = append(out, col)
}
if out == nil {
out = []ColumnInfo{}
}
return out, rows.Err()
}
+155
View File
@@ -0,0 +1,155 @@
package db
import (
"context"
"database/sql"
"fmt"
"strings"
_ "github.com/jackc/pgx/v5/stdlib"
)
func openPostgres(c ConnReq) (*sql.DB, error) {
host := c.Host
if host == "" {
host = "127.0.0.1"
}
port := c.Port
if port == 0 {
port = 5432
}
user := c.User
if user == "" {
user = "postgres"
}
dbName := c.Database
if dbName == "" {
dbName = "postgres"
}
mode := EffectiveSSLMode(c)
sslmode := "disable"
switch mode {
case "disable":
sslmode = "disable"
case "require", "skip-verify":
// libpq require ≈ encrypt without CA/hostname verify
sslmode = "require"
case "verify-ca":
sslmode = "verify-ca"
if strings.TrimSpace(c.SSLCA) == "" {
if strings.TrimSpace(c.SSLCAContent) != "" {
return nil, fmt.Errorf("verify-ca 模式需要 CA 证书路径(pg 驱动暂不支持证书内容,请改用 sslCA)")
}
return nil, fmt.Errorf("verify-ca 模式需要填写 CA 证书路径")
}
default:
return nil, fmt.Errorf("未知 SSL 模式: %s", mode)
}
var b strings.Builder
fmt.Fprintf(&b, "host=%s port=%d user=%s password=%s dbname=%s sslmode=%s connect_timeout=5",
pqEscape(host), port, pqEscape(user), pqEscape(c.Password), pqEscape(dbName), sslmode)
if ca := strings.TrimSpace(c.SSLCA); ca != "" {
fmt.Fprintf(&b, " sslrootcert=%s", pqEscape(ca))
}
if cert := strings.TrimSpace(c.SSLCert); cert != "" {
fmt.Fprintf(&b, " sslcert=%s", pqEscape(cert))
}
if key := strings.TrimSpace(c.SSLKey); key != "" {
fmt.Fprintf(&b, " sslkey=%s", pqEscape(key))
}
return sql.Open("pgx", b.String())
}
func pqEscape(v string) string {
if v == "" {
return "''"
}
needs := strings.ContainsAny(v, " '\\")
if !needs {
return v
}
return "'" + strings.ReplaceAll(strings.ReplaceAll(v, `\`, `\\`), `'`, `\'`) + "'"
}
func listDatabasesPostgres(ctx context.Context, db *sql.DB) ([]string, error) {
rows, err := db.QueryContext(ctx, `SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY 1`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listTablesPostgres(ctx context.Context, db *sql.DB, database string) ([]string, error) {
// 当前连接库下的 public 表;切库需重新 Open
_ = database
rows, err := db.QueryContext(ctx, `
SELECT CASE WHEN schemaname = 'public' THEN tablename ELSE schemaname || '.' || tablename END
FROM pg_catalog.pg_tables
WHERE schemaname NOT IN ('pg_catalog','information_schema')
ORDER BY 1`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listColumnsPostgres(ctx context.Context, dbh *sql.DB, _, table string) ([]ColumnInfo, error) {
schema, t := splitSchemaTable(table)
if schema == "" {
schema = "public"
}
// to_regclass 解析失败(表不存在/跨 schema)返回 NULL 而非报错,主键标记自然落空
rows, err := dbh.QueryContext(ctx, `
SELECT c.column_name, c.data_type,
COALESCE(c.character_maximum_length::text, c.numeric_precision::text, ''),
c.is_nullable, c.column_default,
CASE WHEN a.attname IS NOT NULL THEN 'PRI' ELSE '' END
FROM information_schema.columns c
LEFT JOIN pg_index i ON i.indrelid = to_regclass(quote_ident($1)||'.'||quote_ident($2)) AND i.indisprimary
LEFT JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attname = c.column_name AND a.attnum = ANY(i.indkey)
WHERE c.table_schema = $1 AND c.table_name = $2
ORDER BY c.ordinal_position`, schema, t)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ColumnInfo
for rows.Next() {
var name, typ, extra, nullable, key string
var def sql.NullString
if err := rows.Scan(&name, &typ, &extra, &nullable, &def, &key); err != nil {
return nil, err
}
if extra != "" {
typ = typ + "(" + extra + ")"
}
col := ColumnInfo{Name: name, Type: typ, Nullable: nullable == "YES", Key: key}
if def.Valid {
col.Default = def.String
}
out = append(out, col)
}
if out == nil {
out = []ColumnInfo{}
}
return out, rows.Err()
}
File diff suppressed because it is too large. Load diff
+164
View File
@@ -0,0 +1,164 @@
package db
import (
"context"
"net"
"strconv"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
)
func redisTestReq(t *testing.T, mr *miniredis.Miniredis) ConnReq {
t.Helper()
host, portStr, err := net.SplitHostPort(mr.Addr())
if err != nil {
t.Fatal(err)
}
port, _ := strconv.Atoi(portStr)
return ConnReq{Driver: "redis", Host: host, Port: port, Database: "0"}
}
// 写命令白名单:SET EX / HSET / EXPIRE / MSET / INCR 等应直接可用。
func TestExecRedisWriteWhitelist(t *testing.T) {
mr, err := miniredis.Run()
if err != nil {
t.Fatal(err)
}
defer mr.Close()
base := redisTestReq(t, mr)
ctx := context.Background()
run := func(sql string) ExecResult {
c := base
c.SQL = sql
return execRedis(ctx, c)
}
if r := run(`SET k1 "v with space" EX 100`); r.Error != "" {
t.Fatalf("SET EX: %s", r.Error)
}
if v, _ := mr.Get("k1"); v != "v with space" {
t.Fatalf("SET 引号分组值被破坏: %q", v)
}
if ttl := mr.TTL("k1"); ttl != 100*time.Second {
t.Fatalf("SET EX TTL: %v", ttl)
}
if r := run("HSET h1 f1 v1 f2 v2"); r.Error != "" {
t.Fatalf("HSET: %s", r.Error)
}
if mr.HGet("h1", "f2") != "v2" {
t.Fatal("HSET 未生效")
}
if r := run("EXPIRE h1 60"); r.Error != "" {
t.Fatalf("EXPIRE: %s", r.Error)
}
if ttl := mr.TTL("h1"); ttl != 60*time.Second {
t.Fatalf("EXPIRE 未生效: %v", ttl)
}
if r := run("PERSIST h1"); r.Error != "" {
t.Fatalf("PERSIST: %s", r.Error)
}
if ttl := mr.TTL("h1"); ttl != 0 {
t.Fatalf("PERSIST 未生效: %v", ttl)
}
if r := run("MSET a 1 b 2"); r.Error != "" {
t.Fatalf("MSET: %s", r.Error)
}
if r := run("MGET a b"); r.Error != "" || r.RowCount != 2 {
t.Fatalf("MGET: %#v", r)
}
if r := run("INCR n1"); r.Error != "" {
t.Fatalf("INCR: %s", r.Error)
}
if r := run("SADD s1 m1 m2"); r.Error != "" {
t.Fatalf("SADD: %s", r.Error)
}
if r := run("ZADD z1 1.5 zm"); r.Error != "" {
t.Fatalf("ZADD: %s", r.Error)
}
if r := run("LPUSH l1 x y"); r.Error != "" {
t.Fatalf("LPUSH: %s", r.Error)
}
if r := run("RENAME a a2"); r.Error != "" {
t.Fatalf("RENAME: %s", r.Error)
}
if !mr.Exists("a2") {
t.Fatal("RENAME 未生效")
}
}
// 危险命令二段确认:FLUSHALL/KEYS 未确认不执行,确认后放行。
func TestExecRedisDangerConfirm(t *testing.T) {
mr, err := miniredis.Run()
if err != nil {
t.Fatal(err)
}
defer mr.Close()
base := redisTestReq(t, mr)
ctx := context.Background()
_ = mr.Set("k", "v")
run := func(sql string, confirm bool) ExecResult {
c := base
c.SQL = sql
c.ConfirmDanger = confirm
return execRedis(ctx, c)
}
if r := run("FLUSHALL", false); !r.NeedConfirm {
t.Fatalf("FLUSHALL 未确认应 needConfirm: %#v", r)
}
if !mr.Exists("k") {
t.Fatal("未确认的 FLUSHALL 不应执行")
}
if r := run("KEYS *", false); !r.NeedConfirm {
t.Fatalf("KEYS 未确认应 needConfirm: %#v", r)
}
if r := run("KEYS *", true); r.Error != "" || r.RowCount != 1 {
t.Fatalf("KEYS 确认后应放行: %#v", r)
}
if r := run("FLUSHDB", true); r.Error != "" {
t.Fatalf("FLUSHDB 确认后应放行: %#v", r)
}
if mr.Exists("k") {
t.Fatal("确认后的 FLUSHDB 应已执行")
}
}
func TestExecRedisGenericDoAndScan(t *testing.T) {
mr, err := miniredis.Run()
if err != nil {
t.Fatal(err)
}
defer mr.Close()
base := redisTestReq(t, mr)
ctx := context.Background()
_ = mr.Set("user:1", "a")
_ = mr.Set("user:2", "b")
_ = mr.Set("other", "c")
run := func(sql string) ExecResult {
c := base
c.SQL = sql
return execRedis(ctx, c)
}
if r := run("EXISTS user:1"); r.Error != "" || r.RowCount != 1 {
t.Fatalf("EXISTS via Do: %#v", r)
}
if r := run("SCAN 0 MATCH user:* COUNT 50"); r.Error != "" {
t.Fatalf("SCAN MATCH: %#v", r)
} else {
if r.RowCount != 2 {
t.Fatalf("SCAN MATCH want 2 got %d %#v", r.RowCount, r.Rows)
}
}
if r := run("ECHO hello"); r.Error != "" {
t.Fatalf("ECHO via Do: %#v", r)
}
}
+90
View File
@@ -0,0 +1,90 @@
package db
import (
"context"
"net"
"strconv"
"testing"
"github.com/alicebob/miniredis/v2"
)
func TestListGetDeleteRedisKeys(t *testing.T) {
mr, err := miniredis.Run()
if err != nil {
t.Fatal(err)
}
defer mr.Close()
_ = mr.Set("str:a", "hello")
mr.HSet("hash:b", "f1", "v1")
mr.Push("list:c", "x", "y")
mr.SAdd("set:d", "m1", "m2")
_, _ = mr.ZAdd("zset:e", 1.5, "zm")
host, portStr, err := net.SplitHostPort(mr.Addr())
if err != nil {
t.Fatal(err)
}
port, _ := strconv.Atoi(portStr)
base := ConnReq{Driver: "redis", Host: host, Port: port, Database: "0"}
list := ListRedisKeys(context.Background(), base)
if list.Error != "" {
t.Fatal(list.Error)
}
if len(list.Keys) < 5 {
t.Fatalf("want >=5 keys, got %d %#v", len(list.Keys), list.Keys)
}
byName := map[string]RedisKeyInfo{}
for _, k := range list.Keys {
byName[k.Name] = k
}
if byName["str:a"].Type != "string" {
t.Fatalf("str type %s", byName["str:a"].Type)
}
if byName["hash:b"].Type != "hash" {
t.Fatalf("hash type %s", byName["hash:b"].Type)
}
// Real Redis returns -1 (no expire); miniredis may report 0.
if byName["str:a"].TTL < -1 {
t.Fatalf("unexpected ttl %d", byName["str:a"].TTL)
}
gv := GetRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "str:a"})
if gv.Error != "" || gv.RowCount != 1 || gv.Rows[0]["value"] != "hello" {
t.Fatalf("get string %#v", gv)
}
gh := GetRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "hash:b"})
if gh.Error != "" || gh.RowCount < 1 || gh.Columns[0] != "field" {
t.Fatalf("get hash %#v", gh)
}
if gh.Rows[0]["field"] != "f1" {
t.Fatalf("hash fields should be sorted, got %#v", gh.Rows)
}
gl := GetRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "list:c"})
if gl.Error != "" || gl.RowCount < 2 {
t.Fatalf("get list %#v", gl)
}
gs := GetRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "set:d"})
if gs.Error != "" || gs.RowCount < 2 {
t.Fatalf("get set %#v", gs)
}
gz := GetRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "zset:e"})
if gz.Error != "" || gz.RowCount != 1 {
t.Fatalf("get zset %#v", gz)
}
pat := ListRedisKeys(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Pattern: "str:*"})
if pat.Error != "" || len(pat.Keys) != 1 || pat.Keys[0].Name != "str:a" {
t.Fatalf("pattern %#v", pat)
}
del := DeleteRedisKey(context.Background(), ConnReq{Driver: "redis", Host: host, Port: port, Database: "0", Key: "str:a"})
if del.Error != "" {
t.Fatal(del.Error)
}
if mr.Exists("str:a") {
t.Fatal("key should be deleted")
}
}
+28
View File
@@ -0,0 +1,28 @@
package db
import (
"strings"
"testing"
)
func TestParseRedisAddrs(t *testing.T) {
got := parseRedisAddrs("127.0.0.1:26379\n127.0.0.1:26380, 10.0.0.1:26379")
if len(got) != 3 {
t.Fatalf("want 3 addrs, got %#v", got)
}
}
func TestOpenRedisModeValidation(t *testing.T) {
_, err := openRedis(ConnReq{Driver: "redis", RedisMode: "sentinel"})
if err == nil || !strings.Contains(err.Error(), "Sentinel") {
t.Fatalf("sentinel without nodes: %v", err)
}
_, err = openRedis(ConnReq{Driver: "redis", RedisMode: "sentinel", RedisNodes: "127.0.0.1:26379"})
if err == nil || !strings.Contains(err.Error(), "masterName") {
t.Fatalf("sentinel without master: %v", err)
}
_, err = openRedis(ConnReq{Driver: "redis", RedisMode: "cluster"})
if err == nil || !strings.Contains(err.Error(), "Cluster") {
t.Fatalf("cluster without nodes: %v", err)
}
}
+29
View File
@@ -0,0 +1,29 @@
package db
import (
"testing"
"time"
)
func TestTTLSecondsSentinels(t *testing.T) {
cases := []struct {
in time.Duration
want int64
}{
{time.Duration(-1), -1},
{-1 * time.Nanosecond, -1},
{-1 * time.Second, -1},
{time.Duration(-2), -2},
{-2 * time.Nanosecond, -2},
{-2 * time.Second, -2},
{0, 0},
{5 * time.Second, 5},
{1500 * time.Millisecond, 1},
}
for _, c := range cases {
got := ttlSeconds(c.in)
if got != c.want {
t.Fatalf("ttlSeconds(%v)=%d want %d", c.in, got, c.want)
}
}
}
+131
View File
@@ -0,0 +1,131 @@
package db
import (
"context"
"database/sql"
"strings"
)
// ObjectHit 对象搜索结果条目。
type ObjectHit struct {
Schema string `json:"schema"`
Name string `json:"name"`
Type string `json:"type"` // table | view
}
// SearchObjects 按关键词模糊搜索表/视图,最多返回 50 条。redis 返回空。
func SearchObjects(ctx context.Context, c ConnReq) ([]ObjectHit, error) {
ctx = reqCtx(ctx)
if isRedis(c.Driver) {
return []ObjectHit{}, nil
}
var out []ObjectHit
err := withDB(c, func(db *sql.DB) error {
var err error
switch {
case isMySQLFamily(c.Driver):
out, err = searchObjectsMySQL(ctx, db, c.Database, c.Keyword)
case isPostgresFamily(c.Driver):
out, err = searchObjectsPostgres(ctx, db, c.Keyword)
case isMSSQL(c.Driver):
out, err = searchObjectsMSSQL(ctx, db, c.Keyword)
case driverKey(c.Driver) == "oracle":
out, err = searchObjectsOracle(ctx, db, c.Database, c.User, c.Keyword)
case driverKey(c.Driver) == "sqlite":
out, err = searchObjectsSQLite(ctx, db, c.Keyword)
default:
err = unsupportedDriverErr(c)
}
return err
})
if err != nil {
return nil, err
}
if out == nil {
out = []ObjectHit{}
}
return out, nil
}
func likePattern(keyword string) string {
return "%" + keyword + "%"
}
// hitType information_schema 的 table_type 含 VIEW 即为视图。
func hitType(tableType string) string {
if strings.Contains(strings.ToUpper(tableType), "VIEW") {
return "view"
}
return "table"
}
func searchObjectsMySQL(ctx context.Context, dbh *sql.DB, database, keyword string) ([]ObjectHit, error) {
// database 为空时用当前库 DATABASE()
q := `SELECT table_schema, table_name, table_type FROM information_schema.TABLES
WHERE table_schema = DATABASE() AND table_name LIKE ? LIMIT 50`
args := []interface{}{likePattern(keyword)}
if strings.TrimSpace(database) != "" {
q = `SELECT table_schema, table_name, table_type FROM information_schema.TABLES
WHERE table_schema = ? AND table_name LIKE ? LIMIT 50`
args = []interface{}{database, likePattern(keyword)}
}
return scanObjectHits(ctx, dbh, q, args, "mysql")
}
func searchObjectsPostgres(ctx context.Context, dbh *sql.DB, keyword string) ([]ObjectHit, error) {
return scanObjectHits(ctx, dbh,
`SELECT table_schema, table_name, table_type FROM information_schema.tables
WHERE table_schema NOT IN ('pg_catalog','information_schema') AND table_name LIKE $1 LIMIT 50`,
[]interface{}{likePattern(keyword)}, "pg")
}
func searchObjectsMSSQL(ctx context.Context, dbh *sql.DB, keyword string) ([]ObjectHit, error) {
// MSSQL 无 LIMIT,用 TOP 50
return scanObjectHits(ctx, dbh,
`SELECT TOP 50 TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES
WHERE TABLE_NAME LIKE @p1`,
[]interface{}{likePattern(keyword)}, "mssql")
}
func searchObjectsOracle(ctx context.Context, dbh *sql.DB, database, user, keyword string) ([]ObjectHit, error) {
owner := strings.ToUpper(strings.TrimSpace(database))
if owner == "" {
owner = strings.ToUpper(strings.TrimSpace(user))
}
return scanObjectHits(ctx, dbh,
`SELECT owner, table_name, 'table' FROM all_tables WHERE owner = :1 AND table_name LIKE :2
UNION ALL
SELECT owner, view_name, 'view' FROM all_views WHERE owner = :3 AND view_name LIKE :4
FETCH FIRST 50 ROWS ONLY`,
[]interface{}{owner, likePattern(keyword), owner, likePattern(keyword)}, "oracle")
}
func searchObjectsSQLite(ctx context.Context, dbh *sql.DB, keyword string) ([]ObjectHit, error) {
return scanObjectHits(ctx, dbh,
`SELECT '', name, type FROM sqlite_master WHERE name LIKE ? AND type IN ('table','view') LIMIT 50`,
[]interface{}{likePattern(keyword)}, "sqlite")
}
func scanObjectHits(ctx context.Context, dbh *sql.DB, q string, args []interface{}, dialect string) ([]ObjectHit, error) {
rows, err := dbh.QueryContext(ctx, q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
out := []ObjectHit{}
for rows.Next() {
var schema, name, typ string
if err := rows.Scan(&schema, &name, &typ); err != nil {
return nil, err
}
t := "table"
switch dialect {
case "sqlite":
t = typ // sqlite_master.type 直接是 table/view
default:
t = hitType(typ)
}
out = append(out, ObjectHit{Schema: schema, Name: name, Type: t})
}
return out, rows.Err()
}
+55
View File
@@ -0,0 +1,55 @@
package db
import (
"context"
"testing"
)
// R2-B6:sqlite 搜索命中表与视图,类型区分正确。
func TestSearchObjectsSQLite(t *testing.T) {
c := newSQLiteFile(t,
`CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT)`,
`CREATE TABLE orders (id INTEGER PRIMARY KEY)`,
`CREATE VIEW v_users AS SELECT id, name FROM users`,
)
c.Keyword = "user"
hits, err := SearchObjects(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if len(hits) != 2 {
t.Fatalf("应命中 2 个(user 前缀的表和视图): %+v", hits)
}
byName := map[string]ObjectHit{}
for _, h := range hits {
byName[h.Name] = h
}
if byName["users"].Type != "table" {
t.Fatalf("users 应为 table: %+v", byName["users"])
}
if byName["v_users"].Type != "view" {
t.Fatalf("v_users 应为 view: %+v", byName["v_users"])
}
// 不命中
c.Keyword = "nothing_match"
hits, err = SearchObjects(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if len(hits) != 0 {
t.Fatalf("不应命中: %+v", hits)
}
}
// redis 返回空结果而非报错。
func TestSearchObjectsRedis(t *testing.T) {
c := ConnReq{Driver: "redis", Keyword: "k"}
hits, err := SearchObjects(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if len(hits) != 0 {
t.Fatalf("redis 应返回空: %+v", hits)
}
}
+370
View File
@@ -0,0 +1,370 @@
package db
import (
"context"
"crypto/rand"
"database/sql"
"encoding/hex"
"fmt"
"strconv"
"sync"
"time"
"github.com/redis/go-redis/v9"
)
const sessionTTL = 45 * time.Minute
const maxSessions = 64
// maxPoolHandles 每会话缓存的 *sql.DB 句柄上限(按 database 分键),超出时关闭最久未用。
const maxPoolHandles = 8
type Session struct {
ID string
Conn ConnReq
mu sync.Mutex
sqlPool map[string]*sql.DB
poolUsed map[string]time.Time // sqlPool 各键最近使用时间,供 LRU 轮换
redisClis map[int]redis.UniversalClient // 按 DB 序号缓存的 Redis 客户端(DB 是连接级参数)
// agentHandles: agent 驱动(如达梦)按 database 缓存的 JDBC 句柄;
// JVM 进程走全局池(agentPool),会话关闭只关句柄、不杀进程。
agentHandles map[string]string
tunnel *SSHTunnel
LastUsed time.Time
}
var (
sessionsMu sync.Mutex
sessions = map[string]*Session{}
)
func init() {
go func() {
t := time.NewTicker(2 * time.Minute)
defer t.Stop()
for range t.C {
expireSessions()
}
}()
}
// newSessionID 生成 128bit 随机会话 ID;crypto/rand 失败时返回错误,
// 不再静默降级为低熵时间戳。
func newSessionID() (string, error) {
var b [16]byte
if _, err := rand.Read(b[:]); err != nil {
return "", fmt.Errorf("生成会话 ID 失败: %w", err)
}
return hex.EncodeToString(b[:]), nil
}
func OpenSession(c ConnReq) (string, error) {
c.SQL = ""
c.Table = ""
c.Key = ""
c.Pattern = ""
c.SessionID = ""
// agent 驱动(JDBC 桥)不走本地 SSH 端口转发(转发对 JVM 侧 TCP 无效)
if isAgentDriver(c.Driver) && c.SSHEnabled {
return "", fmt.Errorf("%s 走 JDBC agent,暂不支持 SSH 隧道", c.Driver)
}
dial, tun, err := ApplySSHTunnel(c)
if err != nil {
return "", err
}
if err := testDial(dial); err != nil {
if tun != nil {
tun.Close()
}
return "", fmt.Errorf("数据库: %w", err)
}
sessionsMu.Lock()
if len(sessions) >= maxSessions {
sessionsMu.Unlock()
if tun != nil {
tun.Close()
}
return "", fmt.Errorf("连接会话已满,请关闭不用的连接后重试")
}
id, err := newSessionID()
if err != nil {
sessionsMu.Unlock()
if tun != nil {
tun.Close()
}
return "", err
}
// Persist dialed endpoints (127.0.0.1:local) so later ops reuse the tunnel.
stored := dial
stored.SSHEnabled = false
s := &Session{
ID: id,
Conn: stored,
sqlPool: map[string]*sql.DB{},
poolUsed: map[string]time.Time{},
agentHandles: map[string]string{},
tunnel: tun,
LastUsed: time.Now(),
}
sessions[id] = s
sessionsMu.Unlock()
return id, nil
}
func GetSession(id string) (*Session, error) {
if id == "" {
return nil, fmt.Errorf("缺少 sessionId")
}
sessionsMu.Lock()
s, ok := sessions[id]
sessionsMu.Unlock()
if !ok {
return nil, fmt.Errorf("会话已失效,请重新连接")
}
s.mu.Lock()
s.LastUsed = time.Now()
s.mu.Unlock()
return s, nil
}
func CloseSession(id string) {
sessionsMu.Lock()
s, ok := sessions[id]
if ok {
delete(sessions, id)
}
sessionsMu.Unlock()
if ok {
s.closeAll()
}
}
func expireSessions() {
cutoff := time.Now().Add(-sessionTTL)
sessionsMu.Lock()
var dead []*Session
for id, s := range sessions {
s.mu.Lock()
stale := s.LastUsed.Before(cutoff)
s.mu.Unlock()
if stale {
dead = append(dead, s)
delete(sessions, id)
}
}
sessionsMu.Unlock()
for _, s := range dead {
s.closeAll()
}
}
func (s *Session) closeAll() {
s.mu.Lock()
defer s.mu.Unlock()
for k, dbh := range s.sqlPool {
_ = dbh.Close()
delete(s.sqlPool, k)
}
s.poolUsed = map[string]time.Time{}
for k, cli := range s.redisClis {
_ = cli.Close()
delete(s.redisClis, k)
}
// agent 句柄随会话释放;JVM 进程归全局池,进程生命周期不挂钩会话
if len(s.agentHandles) > 0 {
if cli, err := getAgentClient(s.Conn.Driver); err == nil {
for _, h := range s.agentHandles {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
_ = cli.call(ctx, "close", map[string]any{"handle": h}, nil)
cancel()
}
}
s.agentHandles = map[string]string{}
}
if s.tunnel != nil {
s.tunnel.Close()
s.tunnel = nil
}
}
// agentConn 按 database 懒建/复用 agent JDBC 句柄(模式同 sqlDB)。
// 调用方不得持有 s.mu(本函数内部自行加锁;getAgentClient 用独立的全局池锁)。
func (s *Session) agentConn(c ConnReq) (*AgentClient, string, error) {
cli, err := getAgentClient(c.Driver)
if err != nil {
return nil, "", err
}
key := c.Database
s.mu.Lock()
if h, ok := s.agentHandles[key]; ok {
s.mu.Unlock()
return cli, h, nil
}
s.mu.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
h, err := agentOpen(ctx, cli, c)
if err != nil {
return nil, "", err
}
s.mu.Lock()
defer s.mu.Unlock()
if old, ok := s.agentHandles[key]; ok {
// 双检:并发下已有别的 goroutine 建好同键句柄,回收多余句柄
go func() {
ctx2, cancel2 := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel2()
_ = cli.call(ctx2, "close", map[string]any{"handle": h}, nil)
}()
return cli, old, nil
}
s.agentHandles[key] = h
return cli, h, nil
}
// touch 滚动续期(GetSession 已调用,池化路径同样经 GetSession 进入)。
func (s *Session) touch() {
s.mu.Lock()
s.LastUsed = time.Now()
s.mu.Unlock()
}
// sqlDB 按 database 懒加载缓存 *sql.DB,会话内复用连接池;
// 超过 maxPoolHandles 个句柄时关闭最久未用的池。
func (s *Session) sqlDB(database string) (*sql.DB, error) {
// SQLite 的 database 是文件路径,对象树的 "main" 不能覆盖它
if driverKey(s.Conn.Driver) == "sqlite" {
database = ""
}
s.mu.Lock()
if dbh, ok := s.sqlPool[database]; ok {
s.poolUsed[database] = time.Now()
s.mu.Unlock()
return dbh, nil
}
s.mu.Unlock()
conn := s.Conn
if database != "" {
conn.Database = database
}
dbh, err := Open(conn)
if err != nil {
return nil, err
}
s.mu.Lock()
defer s.mu.Unlock()
// 双检:并发下可能已有别的 goroutine 建好同键池
if old, ok := s.sqlPool[database]; ok {
s.poolUsed[database] = time.Now()
_ = dbh.Close()
return old, nil
}
// 超上限先淘汰最久未用
for len(s.sqlPool) >= maxPoolHandles {
var lruKey string
var lruTime time.Time
for k := range s.sqlPool {
t := s.poolUsed[k]
if lruKey == "" || t.Before(lruTime) {
lruKey, lruTime = k, t
}
}
_ = s.sqlPool[lruKey].Close()
delete(s.sqlPool, lruKey)
delete(s.poolUsed, lruKey)
}
s.sqlPool[database] = dbh
s.poolUsed[database] = time.Now()
return dbh, nil
}
// redisClient 懒建会话级 Redis 单例,避免每次命令新建连接。
func (s *Session) redisClient() (redis.UniversalClient, error) {
// 按连接配置缓存:DB 序号是 go-redis 的连接级参数(建连时钉死),
// 对象树里切换 db0/db1 必须各建一个客户端,否则全部打到首个 db。
return s.redisClientFor(s.Conn)
}
func (s *Session) redisClientFor(c ConnReq) (redis.UniversalClient, error) {
db := redisDBIndex(c)
s.mu.Lock()
defer s.mu.Unlock()
if s.redisClis == nil {
s.redisClis = map[int]redis.UniversalClient{}
}
if cli, ok := s.redisClis[db]; ok {
return cli, nil
}
// 以会话凭据 + 目标 DB 建连(Sentinel/Cluster 模式 DB 恒为 0,天然只有一份)
cli, err := openRedis(ConnReq{
Driver: s.Conn.Driver,
Host: s.Conn.Host,
Port: s.Conn.Port,
User: s.Conn.User,
Password: s.Conn.Password,
Database: strconv.Itoa(db),
RedisMode: s.Conn.RedisMode,
RedisNodes: s.Conn.RedisNodes,
MasterName: s.Conn.MasterName,
SentinelPassword: s.Conn.SentinelPassword,
SSL: s.Conn.SSL,
SSLMode: s.Conn.SSLMode,
SSLCA: s.Conn.SSLCA,
SSLCert: s.Conn.SSLCert,
SSLKey: s.Conn.SSLKey,
SSHEnabled: s.Conn.SSHEnabled,
SSHHost: s.Conn.SSHHost,
SSHPort: s.Conn.SSHPort,
SSHUser: s.Conn.SSHUser,
SSHAuth: s.Conn.SSHAuth,
SSHPassword: s.Conn.SSHPassword,
SSHPrivateKey: s.Conn.SSHPrivateKey,
SSHPassphrase: s.Conn.SSHPassphrase,
})
if err != nil {
return nil, err
}
s.redisClis[db] = cli
return cli, nil
}
func (s *Session) ConnFor(req ConnReq) ConnReq {
c := s.Conn
// SQLite 的 database 是文件路径,对象树里的 "main" 不能覆盖它
if req.Database != "" && driverKey(c.Driver) != "sqlite" {
c.Database = req.Database
}
c.SQL = req.SQL
c.Table = req.Table
c.Key = req.Key
c.Pattern = req.Pattern
c.Keyword = req.Keyword
c.Format = req.Format
c.Limit = req.Limit
c.Offset = req.Offset
c.ConfirmDanger = req.ConfirmDanger
// 逐请求字段:新增 ConnReq 字段时若属于"每次请求携带"语义,必须在此补拷贝,
// 否则会话路径会静默丢弃(前端全部走会话,等于功能失效)
c.TimeoutMs = req.TimeoutMs
c.Cursor = req.Cursor
c.RequestID = req.RequestID
c.SessionID = s.ID
return c
}
func ResolveConn(req ConnReq) (ConnReq, error) {
if req.SessionID == "" {
return req, nil
}
s, err := GetSession(req.SessionID)
if err != nil {
return req, err
}
return s.ConnFor(req), nil
}
+244
View File
@@ -0,0 +1,244 @@
package db
import (
"context"
"database/sql"
"os"
"path/filepath"
"testing"
"time"
)
// touchFile 先建空文件(零字节即合法 SQLite 库),绕开存在性检查。
func touchFile(t *testing.T, path string) {
t.Helper()
if err := os.WriteFile(path, []byte{}, 0600); err != nil {
t.Fatal(err)
}
}
func newTestSession(driver, database string) *Session {
return &Session{
ID: "test-pool",
Conn: ConnReq{Driver: driver, Database: database},
sqlPool: map[string]*sql.DB{},
poolUsed: map[string]time.Time{},
LastUsed: time.Now(),
}
}
// 池化:同 database 复用同一 *sql.DB 句柄。
func TestSessionSqlDBReuse(t *testing.T) {
dbFile := filepath.Join(t.TempDir(), "pool.db")
touchFile(t, dbFile)
s := newTestSession("sqlite", dbFile)
defer s.closeAll()
a, err := s.sqlDB("")
if err != nil {
t.Fatal(err)
}
b, err := s.sqlDB("")
if err != nil {
t.Fatal(err)
}
if a != b {
t.Fatal("同 database 应复用同一 *sql.DB 句柄")
}
// sqlite 下 database 参数不覆盖文件路径
c, err := s.sqlDB("main")
if err != nil {
t.Fatal(err)
}
if c != a {
t.Fatal("sqlite 会话应忽略 database 复用同一句柄")
}
}
// 池化上限:超过 8 个句柄时关闭最久未用(LRU 轮换)。
func TestSessionSqlDBLRU(t *testing.T) {
// mysql 驱动 sql.Open 不立即建连,可离线验证池管理逻辑
s := newTestSession("mysql", "")
defer s.closeAll()
first, err := s.sqlDB("db0")
if err != nil {
t.Fatal(err)
}
for i := 1; i <= maxPoolHandles; i++ {
time.Sleep(2 * time.Millisecond) // 保证 LRU 时间序确定
if _, err := s.sqlDB(string(rune('a' + i))); err != nil {
t.Fatal(err)
}
}
s.mu.Lock()
n := len(s.sqlPool)
_, firstAlive := s.sqlPool["db0"]
s.mu.Unlock()
if n > maxPoolHandles {
t.Fatalf("池句柄数 %d 超上限 %d", n, maxPoolHandles)
}
if firstAlive {
t.Fatal("最久未用句柄应被淘汰")
}
_ = first
}
// 会话路径端到端:ExecSQL 走池化,两次执行共享句柄且结果正确。
func TestExecSQLSessionPooling(t *testing.T) {
dbFile := filepath.Join(t.TempDir(), "exec.db")
touchFile(t, dbFile)
id, err := OpenSession(ConnReq{Driver: "sqlite", Database: dbFile})
if err != nil {
t.Fatal(err)
}
defer CloseSession(id)
mk := func(sqlText string) ConnReq {
c, err := ResolveConn(ConnReq{SessionID: id, SQL: sqlText})
if err != nil {
t.Fatal(err)
}
return c
}
if r := ExecSQL(context.Background(), mk("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")); r.Error != "" {
t.Fatal(r.Error)
}
if r := ExecSQL(context.Background(), mk("INSERT INTO t (v) VALUES ('a'),('b'),('c')")); r.Error != "" {
t.Fatal(r.Error)
}
r := ExecSQL(context.Background(), mk("SELECT * FROM t ORDER BY id"))
if r.Error != "" || r.RowCount != 3 {
t.Fatalf("查询结果错误: %#v", r)
}
s, err := GetSession(id)
if err != nil {
t.Fatal(err)
}
s.mu.Lock()
poolSize := len(s.sqlPool)
s.mu.Unlock()
if poolSize != 1 {
t.Fatalf("会话池应仅 1 个句柄,实际 %d", poolSize)
}
// offset 分页:跳过前 1 行取 1 行
p := mk("SELECT * FROM t ORDER BY id")
p.Offset = 1
p.Limit = 1
r2 := ExecSQL(context.Background(), p)
if r2.Error != "" || r2.RowCount != 1 {
t.Fatalf("offset 分页错误: %#v", r2)
}
if r2.Rows[0]["v"] != "b" {
t.Fatalf("offset=1 应返回第 2 行,实际 %#v", r2.Rows[0])
}
if r2.Offset != 1 || !r2.Truncated {
t.Fatalf("分页元数据错误: offset=%d truncated=%v", r2.Offset, r2.Truncated)
}
// 会话路径危险确认:ConfirmDanger 必须经 ConnFor 透传
d := mk("DROP TABLE t")
if r := ExecSQL(context.Background(), d); !r.NeedConfirm {
t.Fatalf("会话路径未确认应拦截: %#v", r)
}
d.ConfirmDanger = true
if r := ExecSQL(context.Background(), d); r.Error != "" || r.NeedConfirm {
t.Fatalf("会话路径确认后应放行: %#v", r)
}
// 逐请求字段透传:RequestID/TimeoutMs 丢失会让取消/超时在前端(全走会话路径)静默失效
s2, err := GetSession(id)
if err != nil {
t.Fatal(err)
}
c2 := s2.ConnFor(ConnReq{Driver: "sqlite", SessionID: id, SQL: "SELECT 1", RequestID: "req-1", TimeoutMs: 5000})
if c2.RequestID != "req-1" || c2.TimeoutMs != 5000 {
t.Fatalf("ConnFor 应透传 RequestID/TimeoutMs: %#v", c2)
}
}
// 危险语句二段流:未确认拦截,确认后放行。
func TestExecSQLDangerConfirm(t *testing.T) {
dbFile := filepath.Join(t.TempDir(), "danger.db")
touchFile(t, dbFile)
c := ConnReq{Driver: "sqlite", Database: dbFile}
r := ExecSQL(context.Background(), ConnReq{Driver: "sqlite", Database: dbFile, SQL: "CREATE TABLE t (id int)"})
if r.Error != "" {
t.Fatal(r.Error)
}
// 未确认:DROP 不执行
r = ExecSQL(context.Background(), ConnReq{Driver: "sqlite", Database: dbFile, SQL: "DROP TABLE t"})
if !r.NeedConfirm {
t.Fatalf("应 needConfirm: %#v", r)
}
if r.Error == "" {
t.Fatal("应返回命中原因")
}
// 表仍在
r = ExecSQL(context.Background(), ConnReq{Driver: "sqlite", Database: dbFile, SQL: "SELECT count(*) AS n FROM t"})
if r.Error != "" {
t.Fatalf("DROP 不应已执行: %s", r.Error)
}
// 确认后放行
c.SQL = "DROP TABLE t"
c.ConfirmDanger = true
r = ExecSQL(context.Background(), c)
if r.Error != "" || r.NeedConfirm {
t.Fatalf("确认后应放行: %#v", r)
}
// 表已删
r = ExecSQL(context.Background(), ConnReq{Driver: "sqlite", Database: dbFile, SQL: "SELECT count(*) FROM t"})
if r.Error == "" {
t.Fatal("确认后 DROP 应已生效")
}
}
// bigint 精度:超 2^53 的 int64 以字符串无损返回。
func TestExecSQLBigint(t *testing.T) {
dbFile := filepath.Join(t.TempDir(), "bigint.db")
touchFile(t, dbFile)
mk := func(sqlText string) ConnReq {
return ConnReq{Driver: "sqlite", Database: dbFile, SQL: sqlText}
}
if r := ExecSQL(context.Background(), mk("CREATE TABLE b (id INTEGER PRIMARY KEY)")); r.Error != "" {
t.Fatal(r.Error)
}
if r := ExecSQL(context.Background(), mk("INSERT INTO b (id) VALUES (9223372036854775807)")); r.Error != "" {
t.Fatal(r.Error)
}
r := ExecSQL(context.Background(), mk("SELECT id FROM b"))
if r.Error != "" {
t.Fatal(r.Error)
}
if r.Rows[0]["id"] != "9223372036854775807" {
t.Fatalf("bigint 应字符串保真: %#v", r.Rows[0]["id"])
}
}
// EXPLAIN:sqlite 行列计划非空。
func TestExplainSQLite(t *testing.T) {
dbFile := filepath.Join(t.TempDir(), "explain.db")
touchFile(t, dbFile)
c := ConnReq{Driver: "sqlite", Database: dbFile, SQL: "CREATE TABLE e (id int)"}
if r := ExecSQL(context.Background(), c); r.Error != "" {
t.Fatal(r.Error)
}
r := ExplainSQL(context.Background(), ConnReq{Driver: "sqlite", Database: dbFile, SQL: "SELECT * FROM e"})
if r.Error != "" {
t.Fatal(r.Error)
}
if len(r.Rows) == 0 {
t.Fatal("sqlite EXPLAIN QUERY PLAN 应返回行列计划")
}
}
// SQLite 打开不存在文件应报错而非静默建空库。
func TestSQLiteMissingFile(t *testing.T) {
_, err := Open(ConnReq{Driver: "sqlite", Database: filepath.Join(t.TempDir(), "not-exist", "x.db")})
if err == nil {
t.Fatal("不存在的 sqlite 路径应报错")
}
}
+97
View File
@@ -0,0 +1,97 @@
package db
import (
"context"
"database/sql"
"fmt"
"os"
"strings"
_ "modernc.org/sqlite"
)
// openSQLite 打开 SQLite 库文件;文件不存在时报错而非静默建空库。
// DSN 追加 busy_timeout / foreign_keys pragma。
func openSQLite(c ConnReq) (*sql.DB, error) {
path := c.Database
if path == "" {
path = c.Host
}
if path == "" {
return nil, fmt.Errorf("SQLite 需要 database 填文件路径")
}
dsn := path
switch {
case path == ":memory:":
// 统一走 file: URI 以便携带 _pragma 参数
dsn = "file::memory:"
case strings.HasPrefix(path, "file:"):
// 调用方自带 URI 参数,不再做存在性检查
default:
if _, err := os.Stat(path); err != nil {
return nil, fmt.Errorf("SQLite 文件不存在: %s", path)
}
dsn = "file:" + path
}
sep := "?"
if strings.Contains(dsn, "?") {
sep = "&"
}
dsn += sep + "_pragma=busy_timeout(5000)&_pragma=foreign_keys(1)"
return sql.Open("sqlite", dsn)
}
func quoteSQLiteIdent(s string) string {
return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
}
func listDatabasesSQLite(_ *sql.DB) ([]string, error) {
return []string{"main"}, nil
}
func listTablesSQLite(ctx context.Context, db *sql.DB, _ string) ([]string, error) {
rows, err := db.QueryContext(ctx, `SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY 1`)
if err != nil {
return nil, err
}
defer rows.Close()
out := []string{}
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
out = append(out, name)
}
return out, rows.Err()
}
func listColumnsSQLite(ctx context.Context, dbh *sql.DB, _, table string) ([]ColumnInfo, error) {
rows, err := dbh.QueryContext(ctx, `PRAGMA table_info(`+quoteSQLiteIdent(table)+`)`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ColumnInfo
for rows.Next() {
var cid int
var name, typ string
var notnull, pk int
var def sql.NullString
if err := rows.Scan(&cid, &name, &typ, &notnull, &def, &pk); err != nil {
return nil, err
}
col := ColumnInfo{Name: name, Type: typ, Nullable: notnull == 0}
if pk > 0 {
col.Key = "PRI"
}
if def.Valid {
col.Default = def.String
}
out = append(out, col)
}
if out == nil {
out = []ColumnInfo{}
}
return out, rows.Err()
}
+172
View File
@@ -0,0 +1,172 @@
package db
import "strings"
// SplitStatements 按分号切分 SQL 为多条语句,忽略单/双引号字符串与
// 行注释(--、#)、块注释(/* */)内的分号。语义对齐前端 sql.ts splitStatements。
// 返回的语句已去除首尾空白,空语句被丢弃。
func SplitStatements(sql string) []string {
var out []string
var cur strings.Builder
i := 0
n := len(sql)
inS, inD, inLine, inBlock := false, false, false, false
flush := func() {
if s := strings.TrimSpace(cur.String()); s != "" {
out = append(out, s)
}
cur.Reset()
}
for i < n {
ch := sql[i]
var nxt byte
if i+1 < n {
nxt = sql[i+1]
}
switch {
case inLine:
cur.WriteByte(ch)
if ch == '\n' {
inLine = false
}
i++
case inBlock:
cur.WriteByte(ch)
if ch == '*' && nxt == '/' {
cur.WriteByte(nxt)
inBlock = false
i += 2
} else {
i++
}
case !inS && !inD && ch == '-' && nxt == '-':
inLine = true
cur.WriteString("--")
i += 2
case !inS && !inD && ch == '#':
inLine = true
cur.WriteByte(ch)
i++
case !inS && !inD && ch == '/' && nxt == '*':
inBlock = true
cur.WriteString("/*")
i += 2
case ch == '\'' && !inD:
inS = !inS
cur.WriteByte(ch)
i++
case ch == '"' && !inS:
inD = !inD
cur.WriteByte(ch)
i++
case ch == ';' && !inS && !inD:
flush()
i++
default:
cur.WriteByte(ch)
i++
}
}
flush()
return out
}
// StripComments 剥除 --、#、/* */ 注释,并把字符串字面量替换为 ” 占位,
// 便于在纯代码文本上做危险关键字判定。
// MySQL 可执行注释 /*! ... */ 的内容会被保留参与检查(去掉包裹的 /*! 与 */)。
func StripComments(sql string) string {
var b strings.Builder
b.Grow(len(sql))
i := 0
n := len(sql)
for i < n {
ch := sql[i]
var nxt byte
if i+1 < n {
nxt = sql[i+1]
}
switch {
// 行注释 -- 与 #
case ch == '-' && nxt == '-':
for i < n && sql[i] != '\n' {
i++
}
b.WriteByte(' ')
case ch == '#':
for i < n && sql[i] != '\n' {
i++
}
b.WriteByte(' ')
// 块注释;/*! ... */ 为 MySQL 可执行注释,保留内容
case ch == '/' && nxt == '*':
exec := i+2 < n && sql[i+2] == '!'
start := i + 2
if exec {
start = i + 3
}
j := strings.Index(sql[start:], "*/")
if j < 0 {
if exec {
b.WriteString(sql[start:])
}
return b.String()
}
if exec {
b.WriteByte(' ')
b.WriteString(sql[start : start+j])
b.WriteByte(' ')
} else {
b.WriteByte(' ')
}
i = start + j + 2
// 单引号字面量(处理 '' 与 \' 转义)
case ch == '\'':
b.WriteString("''")
i++
for i < n {
if sql[i] == '\'' {
if i+1 < n && sql[i+1] == '\'' {
i += 2
continue
}
i++
break
}
if sql[i] == '\\' && i+1 < n {
i += 2
continue
}
i++
}
// 双引号标识符/字面量同样置空内容(防 "DROP" 作为标识符误报)
case ch == '"':
b.WriteString(`""`)
i++
for i < n {
if sql[i] == '"' {
if i+1 < n && sql[i+1] == '"' {
i += 2
continue
}
i++
break
}
i++
}
// 反引号标识符内容置空
case ch == '`':
b.WriteString("``")
i++
for i < n && sql[i] != '`' {
i++
}
if i < n {
i++
}
default:
b.WriteByte(ch)
i++
}
}
return b.String()
}
+57
View File
@@ -0,0 +1,57 @@
package db
import (
"strings"
"testing"
)
func TestSplitStatements(t *testing.T) {
cases := []struct {
in string
want []string
}{
{"SELECT 1", []string{"SELECT 1"}},
{"SELECT 1; SELECT 2;", []string{"SELECT 1", "SELECT 2"}},
{"SELECT ';'; SELECT 2", []string{"SELECT ';'", "SELECT 2"}},
{"SELECT 1; -- c;mt\nSELECT 2", []string{"SELECT 1", "-- c;mt\nSELECT 2"}},
{"/* a;b */ SELECT 1; SELECT 2", []string{"/* a;b */ SELECT 1", "SELECT 2"}},
{`SELECT "a;b"; SELECT 2`, []string{`SELECT "a;b"`, "SELECT 2"}},
{";;", nil},
{"", nil},
}
for _, tc := range cases {
got := SplitStatements(tc.in)
if len(got) != len(tc.want) {
t.Fatalf("SplitStatements(%q) = %#v, want %#v", tc.in, got, tc.want)
}
for i := range got {
if got[i] != tc.want[i] {
t.Fatalf("SplitStatements(%q)[%d] = %q, want %q", tc.in, i, got[i], tc.want[i])
}
}
}
}
func TestStripComments(t *testing.T) {
cases := []struct {
in string
has string // 结果应包含(大小写敏感原文)
missing string // 结果不应包含
}{
{"SELECT 1 -- WHERE x", "SELECT 1", "WHERE x"},
{"UPDATE t SET x=1 # note\nWHERE id=1", "WHERE id=1", "note"},
{"SELECT /* hint */ 1", "SELECT", "hint"},
{"SELECT /*!50000 DROP TABLE t*/", "DROP TABLE t", "/*!"},
{"SELECT 'DROP TABLE x'", "SELECT", "DROP TABLE x"},
{`SELECT "col;drop" FROM t`, "FROM t", "col;drop"},
}
for _, tc := range cases {
got := StripComments(tc.in)
if tc.has != "" && !strings.Contains(got, tc.has) {
t.Fatalf("StripComments(%q) = %q, 应包含 %q", tc.in, got, tc.has)
}
if tc.missing != "" && strings.Contains(got, tc.missing) {
t.Fatalf("StripComments(%q) = %q, 不应包含 %q", tc.in, got, tc.missing)
}
}
}
+226
View File
@@ -0,0 +1,226 @@
package db
import (
"fmt"
"io"
"net"
"os"
"path/filepath"
"strings"
"sync"
"time"
"golang.org/x/crypto/ssh"
)
// SSHTunnel holds a local TCP listener that forwards through an SSH session
// to the remote database host:port (local port forward semantics).
type SSHTunnel struct {
LocalPort int
listener net.Listener
client *ssh.Client
cancel chan struct{}
wg sync.WaitGroup
once sync.Once
}
func (t *SSHTunnel) Close() {
if t == nil {
return
}
t.once.Do(func() {
close(t.cancel)
if t.listener != nil {
_ = t.listener.Close()
}
if t.client != nil {
_ = t.client.Close()
}
t.wg.Wait()
})
}
// StartLocalForward dials SSH and listens on 127.0.0.1:0, forwarding to c.Host:c.Port.
func StartLocalForward(c ConnReq) (*SSHTunnel, error) {
sshHost := strings.TrimSpace(c.SSHHost)
if sshHost == "" {
return nil, fmt.Errorf("请填写 SSH 主机")
}
sshPort := c.SSHPort
if sshPort <= 0 {
sshPort = 22
}
sshUser := strings.TrimSpace(c.SSHUser)
if sshUser == "" {
return nil, fmt.Errorf("请填写 SSH 用户")
}
auth, err := sshAuthMethods(c)
if err != nil {
return nil, err
}
cfg := &ssh.ClientConfig{
User: sshUser,
Auth: auth,
HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // product tunnel; host-key pinning later
Timeout: 12 * time.Second,
}
addr := net.JoinHostPort(sshHost, fmt.Sprintf("%d", sshPort))
client, err := ssh.Dial("tcp", addr, cfg)
if err != nil {
return nil, fmt.Errorf("无法连接跳板机 %s: %w", addr, err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
_ = client.Close()
return nil, fmt.Errorf("本地监听失败: %w", err)
}
localPort := ln.Addr().(*net.TCPAddr).Port
remoteHost := c.Host
if remoteHost == "" {
remoteHost = "127.0.0.1"
}
remotePort := c.Port
if remotePort <= 0 {
remotePort = 3306
}
remoteAddr := net.JoinHostPort(remoteHost, fmt.Sprintf("%d", remotePort))
tun := &SSHTunnel{
LocalPort: localPort,
listener: ln,
client: client,
cancel: make(chan struct{}),
}
tun.wg.Add(1)
go func() {
defer tun.wg.Done()
for {
conn, err := ln.Accept()
if err != nil {
select {
case <-tun.cancel:
return
default:
return
}
}
tun.wg.Add(1)
go func(local net.Conn) {
defer tun.wg.Done()
defer local.Close()
remote, err := client.Dial("tcp", remoteAddr)
if err != nil {
return
}
defer remote.Close()
pipe(local, remote)
}(conn)
}
}()
return tun, nil
}
func sshAuthMethods(c ConnReq) ([]ssh.AuthMethod, error) {
authMode := strings.ToLower(strings.TrimSpace(c.SSHAuth))
if authMode == "" {
authMode = "password"
}
switch authMode {
case "password":
if c.SSHPassword == "" {
return nil, fmt.Errorf("请填写 SSH 密码")
}
return []ssh.AuthMethod{ssh.Password(c.SSHPassword)}, nil
case "privatekey", "key", "private_key":
keyMaterial, err := loadSSHPrivateKey(c.SSHPrivateKey)
if err != nil {
return nil, err
}
var signer ssh.Signer
if strings.TrimSpace(c.SSHPassphrase) != "" {
signer, err = ssh.ParsePrivateKeyWithPassphrase(keyMaterial, []byte(c.SSHPassphrase))
} else {
signer, err = ssh.ParsePrivateKey(keyMaterial)
}
if err != nil {
return nil, fmt.Errorf("解析私钥失败: %w", err)
}
return []ssh.AuthMethod{ssh.PublicKeys(signer)}, nil
default:
return nil, fmt.Errorf("不支持的 SSH 认证方式: %s", c.SSHAuth)
}
}
func loadSSHPrivateKey(raw string) ([]byte, error) {
s := strings.TrimSpace(raw)
if s == "" {
return nil, fmt.Errorf("请填写私钥 PEM 或路径")
}
if strings.Contains(s, "BEGIN") && strings.Contains(s, "PRIVATE KEY") {
return []byte(s), nil
}
path := s
if strings.HasPrefix(path, "~/") {
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("展开 ~ 失败: %w", err)
}
path = filepath.Join(home, path[2:])
}
b, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("读取私钥文件失败: %w", err)
}
return b, nil
}
func pipe(a, b net.Conn) {
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
_, _ = io.Copy(a, b)
_ = a.Close()
}()
go func() {
defer wg.Done()
_, _ = io.Copy(b, a)
_ = b.Close()
}()
wg.Wait()
}
// ApplySSHTunnel starts a local forward when sshEnabled and rewrites Host/Port to 127.0.0.1:<local>.
// Caller must Close the returned tunnel (may be nil).
func ApplySSHTunnel(c ConnReq) (ConnReq, *SSHTunnel, error) {
if !c.SSHEnabled {
return c, nil, nil
}
tun, err := StartLocalForward(c)
if err != nil {
return c, nil, fmt.Errorf("SSH: %w", err)
}
return rewriteHostForTunnel(c, tun.LocalPort), tun, nil
}
// rewriteHostForTunnel 把连接改写到隧道本地端口(127.0.0.1:<localPort>)。
// 改写会丢失真实 Host,导致隧道内 TLS ServerName 校验失败(127.0.0.1 不在证书 SAN 里),
// 故改写前把原 Host 记入 TLSServerName(用户显式指定时保留其值)。
func rewriteHostForTunnel(c ConnReq, localPort int) ConnReq {
out := c
if strings.TrimSpace(out.TLSServerName) == "" {
out.TLSServerName = out.Host
}
out.Host = "127.0.0.1"
out.Port = localPort
// Avoid nested tunnel if helpers re-enter
out.SSHEnabled = false
return out
}
+36
View File
@@ -0,0 +1,36 @@
package db
import "testing"
// R2-B1:SSH 隧道改写 Host 后,TLS ServerName 必须保留原 Host(隧道内 TLS 校验依赖)。
func TestRewriteHostForTunnelPreservesServerName(t *testing.T) {
c := ConnReq{Host: "db.internal", Port: 3306, SSHEnabled: true}
out := rewriteHostForTunnel(c, 45678)
if out.Host != "127.0.0.1" || out.Port != 45678 {
t.Fatalf("隧道地址改写错误: %s:%d", out.Host, out.Port)
}
if out.SSHEnabled {
t.Fatal("SSHEnabled 应置 false,防嵌套隧道")
}
if out.TLSServerName != "db.internal" {
t.Fatalf("TLSServerName 应保留原 Host, got %q", out.TLSServerName)
}
}
// 用户显式指定 tlsServerName 时,隧道改写不得覆盖。
func TestRewriteHostForTunnelKeepsExplicitServerName(t *testing.T) {
c := ConnReq{Host: "db.internal", TLSServerName: "db.example.com", SSHEnabled: true}
out := rewriteHostForTunnel(c, 40001)
if out.TLSServerName != "db.example.com" {
t.Fatalf("显式 TLSServerName 应保留, got %q", out.TLSServerName)
}
}
func TestTLSServerNameFor(t *testing.T) {
if got := tlsServerNameFor(ConnReq{Host: "db.example"}); got != "db.example" {
t.Fatalf("空 tlsserverName 应用 Host, got %q", got)
}
if got := tlsServerNameFor(ConnReq{Host: "127.0.0.1", TLSServerName: "db.example"}); got != "db.example" {
t.Fatalf("显式 tlsserverName 优先, got %q", got)
}
}
+148
View File
@@ -0,0 +1,148 @@
package db
import (
"crypto/tls"
"crypto/x509"
"fmt"
"os"
"strings"
)
// EffectiveSSLMode resolves ssl / sslMode for drivers.
// Checkbox off or mode disable → disable. Legacy ssl:true with empty mode → require.
func EffectiveSSLMode(c ConnReq) string {
mode := strings.ToLower(strings.TrimSpace(c.SSLMode))
switch mode {
case "disable", "require", "skip-verify", "verify-ca":
if mode == "disable" {
return "disable"
}
return mode
}
if c.SSL {
return "require"
}
return "disable"
}
// SSLEnabled is true when TLS should be negotiated.
func SSLEnabled(c ConnReq) bool {
return EffectiveSSLMode(c) != "disable"
}
// tlsServerNameFor 返回 TLS 校验用的 ServerName:显式 tlsServerName 优先,否则用 Host。
// SSH 隧道把 Host 改写成 127.0.0.1 之前,ApplySSHTunnel 会把原 Host 存入 TLSServerName。
func tlsServerNameFor(c ConnReq) string {
if v := strings.TrimSpace(c.TLSServerName); v != "" {
return v
}
return strings.TrimSpace(c.Host)
}
// tlsConfigFrom builds a *tls.Config from ConnReq paths/mode.
// Returns (nil, nil) when SSL is disabled.
func tlsConfigFrom(c ConnReq) (*tls.Config, error) {
mode := EffectiveSSLMode(c)
if mode == "disable" {
return nil, nil
}
cfg := &tls.Config{
MinVersion: tls.VersionTLS12,
}
if host := tlsServerNameFor(c); host != "" {
cfg.ServerName = host
}
switch mode {
case "skip-verify":
cfg.InsecureSkipVerify = true
case "require":
// Encrypt; verify with system / provided CAs when present.
case "verify-ca":
// Prefer explicit CA; still verify chain (hostname via ServerName when set).
default:
return nil, fmt.Errorf("未知 SSL 模式: %s", mode)
}
caContent := strings.TrimSpace(c.SSLCAContent)
caPath := strings.TrimSpace(c.SSLCA)
if caContent != "" || caPath != "" {
// 内容优先于路径(免在服务端落盘/共享文件)
var pem []byte
if caContent != "" {
pem = []byte(caContent)
} else {
var err error
pem, err = os.ReadFile(caPath)
if err != nil {
return nil, fmt.Errorf("读取 CA 证书失败 (%s): %w", caPath, err)
}
}
pool := x509.NewCertPool()
if !pool.AppendCertsFromPEM(pem) {
if caContent != "" {
return nil, fmt.Errorf("CA 证书内容不是有效 PEM")
}
return nil, fmt.Errorf("CA 证书不是有效 PEM: %s", caPath)
}
cfg.RootCAs = pool
if mode == "verify-ca" {
// verify-ca: trust given CA but allow hostname mismatch (common for IP hosts)
cfg.InsecureSkipVerify = true
roots := pool
cfg.VerifyPeerCertificate = func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
if len(rawCerts) == 0 {
return fmt.Errorf("服务端未提供证书")
}
certs := make([]*x509.Certificate, 0, len(rawCerts))
for i, raw := range rawCerts {
cert, err := x509.ParseCertificate(raw)
if err != nil {
return fmt.Errorf("解析服务端证书失败: %w", err)
}
if i == 0 {
// keep leaf for Intermediates below
}
certs = append(certs, cert)
}
leaf := certs[0]
inter := x509.NewCertPool()
for _, cert := range certs[1:] {
inter.AddCert(cert)
}
_, err := leaf.Verify(x509.VerifyOptions{Roots: roots, Intermediates: inter})
return err
}
}
} else if mode == "verify-ca" {
return nil, fmt.Errorf("verify-ca 模式需要填写 CA 证书路径")
}
certContent := strings.TrimSpace(c.SSLCertContent)
keyContent := strings.TrimSpace(c.SSLKeyContent)
certPath := strings.TrimSpace(c.SSLCert)
keyPath := strings.TrimSpace(c.SSLKey)
if certContent != "" || keyContent != "" {
// 内容与路径二选一,各自成对:证书与密钥必须同时给
if certContent == "" || keyContent == "" {
return nil, fmt.Errorf("客户端证书与密钥需同时填写")
}
cert, err := tls.X509KeyPair([]byte(certContent), []byte(keyContent))
if err != nil {
return nil, fmt.Errorf("加载客户端证书内容失败: %w", err)
}
cfg.Certificates = []tls.Certificate{cert}
} else if certPath != "" || keyPath != "" {
if certPath == "" || keyPath == "" {
return nil, fmt.Errorf("客户端证书与密钥需同时填写")
}
cert, err := tls.LoadX509KeyPair(certPath, keyPath)
if err != nil {
return nil, fmt.Errorf("加载客户端证书失败: %w", err)
}
cfg.Certificates = []tls.Certificate{cert}
}
return cfg, nil
}
+183
View File
@@ -0,0 +1,183 @@
package db
import (
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"math/big"
"os"
"path/filepath"
"testing"
"time"
)
func TestEffectiveSSLMode(t *testing.T) {
if got := EffectiveSSLMode(ConnReq{}); got != "disable" {
t.Fatalf("empty -> disable, got %s", got)
}
if got := EffectiveSSLMode(ConnReq{SSL: true}); got != "require" {
t.Fatalf("ssl true -> require, got %s", got)
}
if got := EffectiveSSLMode(ConnReq{SSL: true, SSLMode: "skip-verify"}); got != "skip-verify" {
t.Fatalf("mode wins, got %s", got)
}
if got := EffectiveSSLMode(ConnReq{SSL: true, SSLMode: "disable"}); got != "disable" {
t.Fatalf("explicit disable, got %s", got)
}
if !SSLEnabled(ConnReq{SSLMode: "require"}) {
t.Fatal("require should enable")
}
}
func TestTLSConfigFromDisable(t *testing.T) {
cfg, err := tlsConfigFrom(ConnReq{})
if err != nil || cfg != nil {
t.Fatalf("disable => nil,nil got %#v %v", cfg, err)
}
}
func TestTLSConfigFromSkipVerify(t *testing.T) {
cfg, err := tlsConfigFrom(ConnReq{SSL: true, SSLMode: "skip-verify", Host: "db.example"})
if err != nil {
t.Fatal(err)
}
if cfg == nil || !cfg.InsecureSkipVerify {
t.Fatalf("expected InsecureSkipVerify, got %#v", cfg)
}
if cfg.ServerName != "db.example" {
t.Fatalf("ServerName %q", cfg.ServerName)
}
}
func TestTLSConfigFromVerifyCARequiresPath(t *testing.T) {
_, err := tlsConfigFrom(ConnReq{SSLMode: "verify-ca"})
if err == nil {
t.Fatal("expected error without CA")
}
}
func TestTLSConfigFromVerifyCAWithPEM(t *testing.T) {
dir := t.TempDir()
caPath := filepath.Join(dir, "ca.pem")
if err := writeTestCA(caPath); err != nil {
t.Fatal(err)
}
cfg, err := tlsConfigFrom(ConnReq{SSLMode: "verify-ca", SSLCA: caPath, Host: "127.0.0.1"})
if err != nil {
t.Fatal(err)
}
if cfg == nil || cfg.RootCAs == nil {
t.Fatal("expected RootCAs")
}
if cfg.VerifyPeerCertificate == nil {
t.Fatal("expected custom verify for verify-ca")
}
}
func TestTLSConfigFromBadCAPath(t *testing.T) {
_, err := tlsConfigFrom(ConnReq{SSLMode: "require", SSLCA: "/no/such/ca.pem"})
if err == nil {
t.Fatal("expected read error")
}
}
// R2-B1:显式 tlsServerName 优先于 Host(tlsConfigFrom 层)。
func TestTLSConfigFromTLSServerName(t *testing.T) {
cfg, err := tlsConfigFrom(ConnReq{SSLMode: "require", Host: "127.0.0.1", TLSServerName: "db.example.com"})
if err != nil {
t.Fatal(err)
}
if cfg == nil || cfg.ServerName != "db.example.com" {
t.Fatalf("ServerName 应取 tlsServerName, got %#v", cfg)
}
}
// R2-B2:CA/客户端证书内容优先于路径;成对校验仍然生效。
func TestTLSConfigFromContentCerts(t *testing.T) {
caPEM, certPEM, keyPEM := genTestCertMaterials(t)
cfg, err := tlsConfigFrom(ConnReq{
SSLMode: "require", Host: "db.example.com",
SSLCAContent: caPEM, SSLCertContent: certPEM, SSLKeyContent: keyPEM,
})
if err != nil {
t.Fatal(err)
}
if cfg == nil || cfg.RootCAs == nil {
t.Fatal("内容 CA 应构建 RootCAs")
}
if len(cfg.Certificates) != 1 {
t.Fatal("内容客户端证书应加载")
}
if cfg.ServerName != "db.example.com" {
t.Fatalf("ServerName 错误: %q", cfg.ServerName)
}
// 内容 CA 非法 PEM 应报错(不回退路径)
if _, err := tlsConfigFrom(ConnReq{SSLMode: "require", SSLCAContent: "not-a-pem"}); err == nil {
t.Fatal("非法 CA 内容应报错")
}
// 证书与密钥仍需成对:只给内容证书不给内容密钥
if _, err := tlsConfigFrom(ConnReq{SSLMode: "require", SSLCertContent: certPEM}); err == nil {
t.Fatal("半对内容证书应报错")
}
}
// genTestCertMaterials 生成自签名 CA 与客户端证书 PEM 串。
func genTestCertMaterials(t *testing.T) (caPEM, certPEM, keyPEM string) {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
t.Fatal(err)
}
keyDER, err := x509.MarshalECPrivateKey(key)
if err != nil {
t.Fatal(err)
}
keyPEM = string(pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}))
tmpl := &x509.Certificate{
SerialNumber: big.NewInt(2),
Subject: pkix.Name{CommonName: "opsxdb-test-cert"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(24 * time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
BasicConstraintsValid: true,
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
if err != nil {
t.Fatal(err)
}
certPEM = string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}))
return certPEM, certPEM, keyPEM // 自签名,CA 与客户端证书同一张
}
func writeTestCA(path string) error {
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return err
}
tmpl := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "opsxdb-test-ca"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(24 * time.Hour),
IsCA: true,
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
BasicConstraintsValid: true,
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
if err != nil {
return err
}
f, err := os.Create(path)
if err != nil {
return err
}
defer f.Close()
return pem.Encode(f, &pem.Block{Type: "CERTIFICATE", Bytes: der})
}
+157
View File
@@ -0,0 +1,157 @@
package db
import (
"encoding/hex"
"strconv"
"time"
"unicode/utf8"
)
const DefaultRowLimit = 1000
const MaxRowLimit = 10000
// ConnReq 连接 / 执行请求。sessionId 有值时服务端用会话里的账密,请求里的 password 可空。
type ConnReq struct {
Driver string `json:"driver"`
Host string `json:"host"`
Port int `json:"port"`
User string `json:"user"`
Password string `json:"password"`
Database string `json:"database"`
SQL string `json:"sql"`
Service string `json:"service"`
SSL bool `json:"ssl"`
SSLMode string `json:"sslMode"` // disable|require|skip-verify|verify-ca
SSLCA string `json:"sslCA"`
SSLCert string `json:"sslCert"`
SSLKey string `json:"sslKey"`
// 证书 PEM 内容,优先于上面的文件路径(客户端与服务端分离部署时免共享文件)
SSLCAContent string `json:"sslCAContent"`
SSLCertContent string `json:"sslCertContent"`
SSLKeyContent string `json:"sslKeyContent"`
// TLSServerName 显式指定 TLS 校验用的 ServerName;SSH 隧道改写 Host 前
// 会把原 Host 自动填入,避免隧道内 TLS 校验失败。
TLSServerName string `json:"tlsServerName"`
SessionID string `json:"sessionId"`
Table string `json:"table"`
Limit int `json:"limit"`
Offset int `json:"offset"` // 结果分页:跳过前 offset 行再按 limit 收集
Key string `json:"key"` // Redis key (get/delete)
Pattern string `json:"pattern"` // Redis SCAN match
Cursor int64 `json:"cursor"` // Redis 值分页游标(hash/set=SCAN 游标;list/zset=偏移)
Keyword string `json:"keyword"` // 对象搜索关键词
Format string `json:"format"` // 导出格式: csv | json | insert
// TimeoutMs > 0 时 /db/exec、/db/explain 以此为语句超时
TimeoutMs int `json:"timeoutMs"`
// RequestID 非空时注册 in-flight 注册表,可通过 /api/db/cancel 取消
RequestID string `json:"requestId"`
// ConfirmDanger 为 true 表示前端已完成危险语句确认,服务端放行本次执行。
ConfirmDanger bool `json:"confirmDanger"`
SSHEnabled bool `json:"sshEnabled"`
SSHHost string `json:"sshHost"`
SSHPort int `json:"sshPort"`
SSHUser string `json:"sshUser"`
SSHAuth string `json:"sshAuth"` // password | privateKey
SSHPassword string `json:"sshPassword"`
SSHPrivateKey string `json:"sshPrivateKey"`
SSHPassphrase string `json:"sshPassphrase"`
// Redis 拓扑:standalone(默认)| sentinel | cluster
RedisMode string `json:"redisMode"`
// Sentinel/Cluster 节点列表,host:port 多行或逗号分隔
RedisNodes string `json:"redisNodes"`
// Sentinel 主节点名(必填)
MasterName string `json:"masterName"`
// Sentinel 自身密码(可空);Redis 密码仍用 Password
SentinelPassword string `json:"sentinelPassword"`
}
type ColumnInfo struct {
Name string `json:"name"`
Type string `json:"type"`
Nullable bool `json:"nullable"`
Key string `json:"key,omitempty"`
Default any `json:"default,omitempty"`
Comment string `json:"comment,omitempty"`
}
// ExecResult 与前端 QueryResult 字段对齐
type ExecResult struct {
Columns []string `json:"columns"`
Rows []map[string]interface{} `json:"rows"`
RowCount int `json:"rowCount"`
TimeMs int64 `json:"timeMs"`
Error string `json:"error,omitempty"`
Truncated bool `json:"truncated,omitempty"`
Limit int `json:"limit,omitempty"`
Offset int `json:"offset,omitempty"`
// NeedConfirm 为 true 表示命中危险语句,需前端确认后带 confirmDanger 重发。
NeedConfirm bool `json:"needConfirm,omitempty"`
// Redis 值游标分页:>0 还有更多(HGETALL/SMEMBERS/ZRANGE/LRANGE 单键读透传)
NextCursor int64 `json:"nextCursor,omitempty"`
}
// ExplainResult 执行计划结果:行列(计划表格)与 Text(XML/文本计划)二选一或兼有。
type ExplainResult struct {
Columns []string `json:"columns,omitempty"`
Rows []map[string]interface{} `json:"rows,omitempty"`
Text string `json:"text,omitempty"` // MSSQL SHOWPLAN_XML / Oracle 文本计划
TimeMs int64 `json:"timeMs"`
Error string `json:"error,omitempty"`
}
func normalize(v interface{}) interface{} {
switch t := v.(type) {
case nil:
return nil
case int64:
// bigint 超 JS 安全整数时转字符串保真,避免前端精度丢失错行
if t > 9007199254740991 || t < -9007199254740991 {
return strconv.FormatInt(t, 10)
}
return t
case []byte:
if utf8.Valid(t) {
return string(t)
}
// BLOB 等非 UTF-8 数据以十六进制标注返回,替代乱码
return `\x` + hex.EncodeToString(t)
case time.Time:
return t.Format(time.RFC3339Nano)
default:
return t
}
}
// Redis key browse (dedicated endpoints; do not reshape SQL tables []).
type RedisKeyInfo struct {
Name string `json:"name"`
Type string `json:"type"`
TTL int64 `json:"ttl"` // seconds; -1 no expire; -2 missing
}
type RedisKeysResult struct {
Keys []RedisKeyInfo `json:"keys"`
Truncated bool `json:"truncated,omitempty"`
TimeMs int64 `json:"timeMs"`
Error string `json:"error,omitempty"`
}
// RedisKeyValue feeds ResultPane (columns/rows) plus type/ttl meta.
type RedisKeyValue struct {
Name string `json:"name"`
Type string `json:"type"`
TTL int64 `json:"ttl"`
Columns []string `json:"columns"`
Rows []map[string]interface{} `json:"rows"`
RowCount int `json:"rowCount"`
TimeMs int64 `json:"timeMs"`
Truncated bool `json:"truncated,omitempty"`
Limit int `json:"limit,omitempty"`
NextCursor int64 `json:"nextCursor"` // >0 表示还有更多(续载时回传 cursor)
Error string `json:"error,omitempty"`
}
+61
View File
@@ -0,0 +1,61 @@
package db
import (
"strings"
"testing"
)
func TestNormalizeBigint(t *testing.T) {
// 超 JS 安全整数 → 字符串保真
if got := normalize(int64(9223372036854775807)); got != "9223372036854775807" {
t.Fatalf("bigint max: %#v", got)
}
if got := normalize(int64(-9223372036854775807)); got != "-9223372036854775807" {
t.Fatalf("bigint min: %#v", got)
}
// 安全范围内保持 number
if got := normalize(int64(9007199254740991)); got != int64(9007199254740991) {
t.Fatalf("safe int: %#v", got)
}
if got := normalize(int64(9007199254740992)); got != "9007199254740992" {
t.Fatalf("边界+1 应转字符串: %#v", got)
}
}
func TestNormalizeBytes(t *testing.T) {
if got := normalize([]byte("hello")); got != "hello" {
t.Fatalf("utf8 bytes: %#v", got)
}
if got := normalize([]byte("中文")); got != "中文" {
t.Fatalf("utf8 中文: %#v", got)
}
// 非 UTF-8 → \x 十六进制
if got := normalize([]byte{0xff, 0x00, 0x01}); got != `\xff0001` {
t.Fatalf("blob hex: %#v", got)
}
if got := normalize(nil); got != nil {
t.Fatalf("nil: %#v", got)
}
}
func TestMssqlDSNEncoding(t *testing.T) {
// 含空格 / '+' 的密码应经 userinfo 编码,不被 QueryEscape 破坏
dsn, err := mssqlDSN(ConnReq{Driver: "sqlserver", Host: "h1", Port: 1433, User: "sa", Password: "p@ ss+word", Database: "db1"})
if err != nil {
t.Fatal(err)
}
want := "p%40%20ss+word"
if !strings.Contains(dsn, want) {
t.Fatalf("DSN 密码编码不符合预期: %s(应含 %s)", dsn, want)
}
if !strings.Contains(dsn, "sqlserver://sa:") || !strings.Contains(dsn, "@h1:1433") {
t.Fatalf("DSN 结构错误: %s", dsn)
}
if !strings.Contains(dsn, "database=db1") {
t.Fatalf("DSN 缺 database: %s", dsn)
}
// verify-ca 无 CA 应报错
if _, err := mssqlDSN(ConnReq{Driver: "sqlserver", Host: "h1", SSLMode: "verify-ca"}); err == nil {
t.Fatal("verify-ca 缺 CA 应报错")
}
}