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

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 路径应报错")
	}
}