package vrtxsql // database/sql 驱动壳: 连接注册表 (同路径共享一个引擎 + 全局串行锁)、 // Conn/Tx/Stmt/Rows 适配。引擎是单写者模型, 所有操作 (含事务全程) 由 // shared 的容量1信号量串行化。锁获取尊重 ctx —— 排队等锁阶段可被超时/ // 取消中断; 但引擎一旦开始执行便不可中途打断 (单写者, 调用方应串行化, // 建议 SetMaxOpenConns(1))。SELECT 结果集流式返回 (worker 回调 + channel), // 结果集 Close 前持续持锁。 import ( "context" "database/sql" "database/sql/driver" "errors" "fmt" "io" "net/url" "reflect" "strconv" "strings" "sync" "time" "git.ycl.pub/yl32940/TsDB2/DB/core" "git.ycl.pub/yl32940/TsDB2/DB/core/types" ) func init() { sql.Register("vrtx", &vrtxDriver{}) } // vrtxDriver 实现 driver.Driver。 type vrtxDriver struct{} // registry 同一路径 (DSN) 的所有连接共享同一个 *core.DB —— 引擎的文件锁与 // 单写者语义不允许一个库被打开成多个实例。 var registry = struct { sync.Mutex m map[string]*shared }{m: make(map[string]*shared)} type shared struct { db *core.DB dsn string sem chan struct{} // 容量1: 全局串行 (引擎单写者); 事务持有至 Commit/Rollback refs int } // lock 获取串行锁; 排队期间尊重 ctx (超时/取消返回 ctx.Err)。 func (s *shared) lock(ctx context.Context) error { if ctx == nil { ctx = context.Background() } select { case s.sem <- struct{}{}: return nil case <-ctx.Done(): return ctx.Err() } } // unlock 释放串行锁 (非严格: 未持有时不阻塞、不 panic)。 func (s *shared) unlock() { select { case <-s.sem: default: } } // dsnConfig 是 DSN 解析结果: 文件路径 + 打开选项。 type dsnConfig struct { path string readOnly bool // mode=ro private bool // cache=private: 不进共享注册表, 独立引擎实例 } // parseDSN 解析 "文件路径?mode=ro&cache=private" 形式的 DSN。 func parseDSN(dsn string) (*dsnConfig, error) { cfg := &dsnConfig{path: dsn} query := "" if i := strings.IndexByte(dsn, '?'); i >= 0 { cfg.path = dsn[:i] query = dsn[i+1:] } if cfg.path == "" { return nil, errors.New("vrtx: DSN 为空 (应为数据库文件路径)") } values, err := url.ParseQuery(query) if err != nil { return nil, fmt.Errorf("vrtx: DSN 参数解析失败: %w", err) } switch strings.ToLower(values.Get("mode")) { case "", "rw", "readwrite": case "ro", "readonly": cfg.readOnly = true default: return nil, fmt.Errorf("vrtx: 不支持的 mode=%q (应为 ro 或 rw)", values.Get("mode")) } if strings.EqualFold(values.Get("cache"), "private") { cfg.private = true } return cfg, nil } // openEngine 按 DSN 选项打开引擎。 func openEngine(cfg *dsnConfig) (*core.DB, error) { if cfg.readOnly { return core.OpenWithOptions(cfg.path, core.OpenOptions{ReadOnly: true}) } return core.Open(cfg.path) } func sharedFor(dsn string) (*shared, error) { cfg, err := parseDSN(dsn) if err != nil { return nil, err } // cache=private: 每次独立引擎实例, 不进共享注册表 (release 时直接 Close)。 if cfg.private { db, err := openEngine(cfg) if err != nil { return nil, err } return &shared{db: db, dsn: dsn, refs: 1, sem: make(chan struct{}, 1)}, nil } registry.Lock() defer registry.Unlock() if s, ok := registry.m[dsn]; ok { s.refs++ return s, nil } db, err := openEngine(cfg) if err != nil { return nil, err } s := &shared{db: db, dsn: dsn, refs: 1, sem: make(chan struct{}, 1)} registry.m[dsn] = s return s, nil } func (s *shared) release() error { registry.Lock() defer registry.Unlock() s.refs-- if s.refs > 0 { return nil } delete(registry.m, s.dsn) return s.db.Close() } func (d *vrtxDriver) Open(dsn string) (driver.Conn, error) { if dsn == "" { return nil, errors.New("vrtx: DSN 为空 (应为数据库文件路径)") } s, err := sharedFor(dsn) if err != nil { return nil, fmt.Errorf("vrtx: 打开 %s: %w", dsn, err) } return &conn{s: s}, nil } // OpenConnector 实现 driver.DriverContext: database/sql 打开时优先走 Connector // (支持 sql.OpenDB、可固定 DSN); 此处提前校验 DSN, 参数错误在 Open 阶段就暴露。 func (d *vrtxDriver) OpenConnector(dsn string) (driver.Connector, error) { if dsn == "" { return nil, errors.New("vrtx: DSN 为空 (应为数据库文件路径)") } if _, err := parseDSN(dsn); err != nil { return nil, err } return &vrtxConnector{drv: d, dsn: dsn}, nil } // vrtxConnector 实现 driver.Connector。 type vrtxConnector struct { drv *vrtxDriver dsn string } // Connect 打开一个连接; 引擎打开/WAL 恢复不可中断, 但在已取消 ctx 下快速失败。 func (c *vrtxConnector) Connect(ctx context.Context) (driver.Conn, error) { if err := ctx.Err(); err != nil { return nil, err } s, err := sharedFor(c.dsn) if err != nil { return nil, fmt.Errorf("vrtx: 打开 %s: %w", c.dsn, err) } return &conn{s: s}, nil } func (c *vrtxConnector) Driver() driver.Driver { return c.drv } var _ driver.DriverContext = (*vrtxDriver)(nil) // conn 实现 driver.Conn / driver.ConnBeginTx / driver.ExecerContext / // driver.QueryerContext / driver.Pinger / driver.NamedValueChecker。 type conn struct { s *shared // activeTx 本连接上开着的事务: 事务期间持有串行锁, 拥有者连接上的 // Exec/Query 不得重复加锁 (否则自死锁); 其它连接的操作照常抢锁, // 在事务提交/回滚前一律阻塞 —— 这就是全局串行化的实现方式。 activeTx *core.Transaction } var ( _ driver.Conn = (*conn)(nil) _ driver.ConnBeginTx = (*conn)(nil) _ driver.ExecerContext = (*conn)(nil) _ driver.QueryerContext = (*conn)(nil) _ driver.Pinger = (*conn)(nil) _ driver.NamedValueChecker = (*conn)(nil) ) func (c *conn) Prepare(query string) (driver.Stmt, error) { st, err := parse(query) if err != nil { return nil, err } return &stmt{c: c, st: st}, nil } func (c *conn) Close() error { if c.s == nil { return nil } if c.activeTx != nil { // 防御: 连接带着未决事务被回收 —— 回滚并放锁 (正常流程不会发生, // database/sql 在事务结束前不会回收连接) _ = c.activeTx.Rollback() c.activeTx = nil c.s.unlock() } s := c.s c.s = nil return s.release() } func (c *conn) Begin() (driver.Tx, error) { return c.BeginTx(context.Background(), driver.TxOptions{}) } func (c *conn) BeginTx(ctx context.Context, _ driver.TxOptions) (driver.Tx, error) { if err := c.s.lock(ctx); err != nil { return nil, err } if c.s.db.HasActiveTransaction() { c.s.unlock() return nil, errors.New("vrtx: 已有活跃事务 (单写者模型禁止并发事务)") } c.activeTx = c.s.db.BeginTransaction() return &tx{txn: c.activeTx, s: c.s, c: c}, nil } func (c *conn) Ping(_ context.Context) error { // 引擎无活性探针: 句柄在 = 可服务 (打开时的 CRC/WAL 恢复已完成)。 return nil } // CheckNamedValue 实现 driver.NamedValueChecker: 经默认转换器把基础参数规整为 // driver 基础类型 (int→int64、float32→float64 等; time.Time/[]byte 原样保留), // 并明确拒绝命名参数 (解析器只认 ? 占位符)。 func (c *conn) CheckNamedValue(nv *driver.NamedValue) error { if nv.Name != "" { return fmt.Errorf("vrtx: 不支持命名参数 %q, 请使用 ? 占位符", nv.Name) } v, err := driver.DefaultParameterConverter.ConvertValue(nv.Value) if err != nil { return err } nv.Value = v return nil } func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) { st, err := parse(query) if err != nil { return nil, err } if st.kind == kSelect || st.kind == kShow { return nil, errors.New("vrtx: 查询语句 (SELECT/SHOW) 请使用 QueryContext / Query") } if c.activeTx == nil { if err := c.s.lock(ctx); err != nil { return nil, err } defer c.s.unlock() } res, err := execStatement(c.s.db, st, args) return res, classifyError(err) } func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { st, err := parse(query) if err != nil { return nil, err } if st.kind != kSelect && st.kind != kShow { return nil, errors.New("vrtx: 非查询语句请使用 ExecContext / Exec") } var release func() if c.activeTx == nil { if err := c.s.lock(ctx); err != nil { return nil, err } release = c.s.unlock } if st.kind == kShow { rr, e := execShow(c.s.db, st.show, release) return rr, classifyError(e) } rr, e := queryStatement(ctx, c.s.db, st.sel, args, release) return rr, classifyError(e) } // tx 实现 driver.Tx。Commit/Rollback 释放串行锁 (Begin 时持有)。 type tx struct { txn *core.Transaction s *shared c *conn } func (t *tx) Commit() error { err := t.txn.Commit() t.c.activeTx = nil t.s.unlock() return err } func (t *tx) Rollback() error { err := t.txn.Rollback() t.c.activeTx = nil t.s.unlock() return err } // stmt 实现 driver.Stmt (Prepare 路径: 解析已完成, 执行时只绑定参数)。 type stmt struct { c *conn st *statement } var ( _ driver.Stmt = (*stmt)(nil) _ driver.StmtExecContext = (*stmt)(nil) _ driver.StmtQueryContext = (*stmt)(nil) _ driver.NamedValueChecker = (*stmt)(nil) ) func (s *stmt) Close() error { return nil } func (s *stmt) NumInput() int { return s.st.argCount() } func (s *stmt) Exec(args []driver.Value) (driver.Result, error) { return s.execContext(context.Background(), valuesToNamed(args)) } func (s *stmt) Query(args []driver.Value) (driver.Rows, error) { return s.queryContext(context.Background(), valuesToNamed(args)) } // ExecContext 实现 driver.StmtExecContext。 func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) { return s.execContext(ctx, args) } // QueryContext 实现 driver.StmtQueryContext。 func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) { return s.queryContext(ctx, args) } func (s *stmt) CheckNamedValue(nv *driver.NamedValue) error { return s.c.CheckNamedValue(nv) } func (s *stmt) execContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) { if s.st.kind == kSelect || s.st.kind == kShow { return nil, errors.New("vrtx: 查询语句请使用 Query") } if s.c.activeTx == nil { if err := s.c.s.lock(ctx); err != nil { return nil, err } defer s.c.s.unlock() } res, err := execStatement(s.c.s.db, s.st, args) return res, classifyError(err) } func (s *stmt) queryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) { if s.st.kind != kSelect && s.st.kind != kShow { return nil, errors.New("vrtx: 非查询语句请使用 Exec") } var release func() if s.c.activeTx == nil { if err := s.c.s.lock(ctx); err != nil { return nil, err } release = s.c.s.unlock } if s.st.kind == kShow { rr, e := execShow(s.c.s.db, s.st.show, release) return rr, classifyError(e) } rr, e := queryStatement(ctx, s.c.s.db, s.st.sel, args, release) return rr, classifyError(e) } func valuesToNamed(args []driver.Value) []driver.NamedValue { out := make([]driver.NamedValue, len(args)) for i, v := range args { out[i] = driver.NamedValue{Ordinal: i + 1, Value: v} } return out } // streamBuf 流式结果集 channel 缓冲行数。 const streamBuf = 64 // rows 同时支持两种模式: // - 流式 (ch != nil): 普通 SELECT, worker 逐行回调送入; Close 前持锁。 // - 物化 (ch == nil): 单行聚合 / JOIN / SHOW, 执行已完成、锁已释放。 // // blobReader 物化 Blob 列的读取入口 (blobReaderOf 绑定具体表): Blob 列行内只有 // 8B 链引用, 结果集要的是内容字节 —— 流式路径在 worker 里 (持锁期间) 解引用。 // 物化模式的两条路径不经过它: JOIN 在 scanJoinTable 里当场解引用 (join.go), // 聚合/SHOW/字面量行根本不产出 BlobRef。 type rows struct { cols []string metas []colMeta // 流式 ch chan []driver.Value done chan struct{} ended chan struct{} once sync.Once err error // 物化 data [][]driver.Value pos int blobReader func(types.BlobRef) ([]byte, error) release func() } var ( _ driver.Rows = (*rows)(nil) _ driver.RowsColumnTypeDatabaseTypeName = (*rows)(nil) _ driver.RowsColumnTypeNullable = (*rows)(nil) _ driver.RowsColumnTypeScanType = (*rows)(nil) _ driver.RowsColumnTypeLength = (*rows)(nil) ) func newStreamRows(cols []string, metas []colMeta, blobReader func(types.BlobRef) ([]byte, error), release func()) *rows { return &rows{ cols: cols, metas: metas, ch: make(chan []driver.Value, streamBuf), done: make(chan struct{}), ended: make(chan struct{}), blobReader: blobReader, release: release, } } func materialRows(cols []string, metas []colMeta, data [][]driver.Value) *rows { return &rows{cols: cols, metas: metas, data: data} } // start 启动 worker: 在已持锁状态下逐行扫描, 按 colNames 投影后送入 ch。 // q.Scan 的回调给的是未投影完整行, 故这里按 FieldName 自行挑列。 func (r *rows) start(ctx context.Context, colNames []string, distinct bool, q *core.Query) { go func() { defer close(r.ended) defer close(r.ch) var seen map[string]bool if distinct { seen = make(map[string]bool) } scanErr := q.Scan(func(full []types.FieldValue) bool { // 列名→下标映射【每行重建】: 引擎的行是稀疏编码 (NULL 列物理缺席), // 每行字段数不同、同一下标在不同行对应不同列 — 按第一行建表复用会让 // 第二行起映射错位, 行更短直接索引越界 panic (迁移数据必读崩)。 pos := make(map[string]int, len(full)) for i, fv := range full { pos[strings.ToLower(string(fv.FieldName))] = i } vals := make([]driver.Value, len(colNames)) for i, cn := range colNames { idx, ok := pos[strings.ToLower(cn)] if !ok { // 旧版本行不含后加的列 (schema 演进): 新列对旧行 = NULL。 vals[i] = nil continue } raw := full[idx].Value if ref, isRef := raw.(types.BlobRef); isRef && r.blobReader != nil { // Blob 列行内只有 8B 链引用: 在持锁的 worker 里解引用成内容字节。 data, err := r.blobReader(ref) if err != nil { r.err = fmt.Errorf("读取 BLOB 列 %q: %w", cn, err) return false } vals[i] = data continue } dv, err := fieldValueToDriverTyped(r.metas[i].ft, raw) if err != nil { r.err = err return false } vals[i] = dv } if seen != nil { key := distinctKey(vals) if seen[key] { return true // 重复组合: 跳过本行, 继续扫描 } seen[key] = true } select { case r.ch <- vals: return true case <-r.done: return false case <-ctx.Done(): r.err = ctx.Err() return false } }) if scanErr != nil && r.err == nil { r.err = scanErr } }() } // distinctKey 把投影后的一行 driver.Value 编成无歧义字符串 (类型标记 + 长度前缀), // 作为 SELECT DISTINCT 的去重键。 func distinctKey(vals []driver.Value) string { var sb strings.Builder for _, v := range vals { switch x := v.(type) { case nil: sb.WriteString("N:0:") case string: sb.WriteString("S:") sb.WriteString(strconv.Itoa(len(x))) sb.WriteByte(':') sb.WriteString(x) case []byte: sb.WriteString("B:") sb.WriteString(strconv.Itoa(len(x))) sb.WriteByte(':') sb.Write(x) case time.Time: sb.WriteString("T:") sb.WriteString(strconv.FormatInt(x.UnixNano(), 10)) sb.WriteByte(':') case bool: sb.WriteString("L:") if x { sb.WriteByte('1') } else { sb.WriteByte('0') } sb.WriteByte(':') case int64: sb.WriteString("I:") sb.WriteString(strconv.FormatInt(x, 10)) sb.WriteByte(':') case float64: sb.WriteString("F:") sb.WriteString(strconv.FormatFloat(x, 'g', -1, 64)) sb.WriteByte(':') default: sb.WriteString("?:") sb.WriteString(fmt.Sprintf("%T:%v", x, x)) } } return sb.String() } func (r *rows) Columns() []string { return r.cols } func (r *rows) Close() error { r.once.Do(func() { if r.done != nil { close(r.done) } if r.ended != nil { <-r.ended // 等 worker 完全退出、不再访问引擎, 然后才放锁 } if r.release != nil { r.release() r.release = nil } }) return nil } func (r *rows) Next(dest []driver.Value) error { if r.ch != nil { row, ok := <-r.ch if !ok { <-r.ended if r.err != nil { return classifyError(r.err) } return io.EOF } copy(dest, row) return nil } if r.pos >= len(r.data) { return io.EOF } copy(dest, r.data[r.pos]) r.pos++ return nil } // ColumnType* 系列: 供 sqlx/GORM 等做列类型反射。 func (r *rows) ColumnTypeDatabaseTypeName(i int) string { if i >= 0 && i < len(r.metas) { return r.metas[i].dbName } return "" } func (r *rows) ColumnTypeNullable(i int) (bool, bool) { if i >= 0 && i < len(r.metas) { return r.metas[i].nullable, true } return true, false } func (r *rows) ColumnTypeScanType(i int) reflect.Type { if i >= 0 && i < len(r.metas) { return r.metas[i].scanType } return reflect.TypeOf("") } func (r *rows) ColumnTypeLength(i int) (int64, bool) { if i >= 0 && i < len(r.metas) && r.metas[i].hasLen { return r.metas[i].length, true } return 0, false } // result 实现 driver.Result。 type result struct { lastID int64 affected int64 } func (r *result) LastInsertId() (int64, error) { return r.lastID, nil } func (r *result) RowsAffected() (int64, error) { return r.affected, nil }