TsDB2/DB/driver/driver.go

package vrtxsql

// database/sql 驱动壳: 连接注册表 (同路径共享一个引擎)、
// Conn/Tx/Stmt/Rows 适配。引擎是一写多读模型:
//   - shared.sem 容量 16: 允许多读者并行进入引擎; 写操作也走同一条通道,
//     但 core.accMu.Lock() 会把并发写排他化 (Begin/Commit/DDL 持写锁)
//   - 锁获取尊重 ctx —— 排队等锁阶段可被超时/取消中断
//   - SELECT 结果集流式返回 (worker 回调 + channel), Close 前持续持锁
//
// 调用方建议 SetMaxOpenConns(16) (≥ semaphore 容量), 否则并发读会在
// sql.DB 连接池层被串行化, 引擎内部的 RLock 优势发挥不出来。

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{} // 容量16: 允许多个读者并行; Begin/Commit 的写锁由 core.accMu 排他保证
	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{}, 16)}, 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{}, 16)}
	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 }