TsDB2/gormvrtx/dialector.go

// Package gormvrtx 把 TsDB2 (vrtxsql database/sql 驱动) 接入 GORM:
// CRUD 走 GORM 生成 SQL → vrtx 驱动执行; schema 自省 (AutoMigrate) 不经过 SQL,
// 由 Migrator 直接调引擎 API。
//
// 引擎是嵌入式单写者存储, 与 mysql/pgx 这类网络库有两个本质差别, 本包据此裁剪:
//   - 无版本协商、无嵌套事务 (SavePoint/RollbackTo 不声明, 让 GORM 默认行为生效);
//   - 行存储即主键索引, 表必须有整型主键列承载。
package gormvrtx

import (
	"context"
	"database/sql"
	"fmt"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/callbacks"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/schema"

	"git.ycl.pub/yl32940/TsDB2/DB/core"
	"git.ycl.pub/yl32940/TsDB2/DB/core/types"
	vrtxsql "git.ycl.pub/yl32940/TsDB2/DB/driver"
)

// DriverName 是 vrtxsql 注册的 database/sql 驱动名。
const DriverName = "vrtx"

// Config Dialector 配置。
type Config struct {
	DriverName string // database/sql 驱动名, 默认 "vrtx"
	DSN        string // 数据库文件路径 (可带 ?mode=ro / ?cache=private)
}

// Dialector 实现 gorm.Dialector。
type Dialector struct {
	*Config
	db    *core.DB // Initialize 时经 SharedDB 取得的引擎实例, Migrator 直接用
	sqlDB *sql.DB  // 连接池本体 (供调用方 Close 时正确释放)
}

var _ gorm.Dialector = (*Dialector)(nil)

// Open 以默认配置打开: DriverName="vrtx", DSN=文件路径。
func Open(dsn string) gorm.Dialector {
	return &Dialector{Config: &Config{DriverName: DriverName, DSN: dsn}}
}

// New 按配置构造 Dialector; DriverName 为空时用 "vrtx"。
func New(config Config) gorm.Dialector {
	if config.DriverName == "" {
		config.DriverName = DriverName
	}
	return &Dialector{Config: &config}
}

// Name 方言名。
func (d *Dialector) Name() string { return "vrtx" }

// Initialize 打开驱动连接池与引擎实例。不做任何 SELECT VERSION() 探测 ——
// 引擎无版本概念, 打开时的 CRC/WAL 恢复即完整性校验。
//
// 两个引擎相关的适配点:
//   - SetMaxOpenConns(1): 引擎是单写者模型, 连接池压成 1 让所有 SQL 串行;
//   - ConnPool 包一层 rewritePool: 内联 LIMIT/OFFSET 整数参数 (驱动解析器支持限定名)
//     (旧解析器只认裸列名/字面量 LIMIT, 该包装让 GORM 的 SQL 形态稳定可用)。
func (d *Dialector) Initialize(db *gorm.DB) error {
	if d.DriverName == "" {
		d.DriverName = DriverName
	}
	sqlDB, err := sql.Open(d.DriverName, d.DSN)
	if err != nil {
		return err
	}
	sqlDB.SetMaxOpenConns(1)

	eng, err := vrtxsql.SharedDB(d.DSN)
	if err != nil {
		_ = sqlDB.Close()
		return fmt.Errorf("vrtx: 打开引擎 %s: %w", d.DSN, err)
	}
	d.db = eng
	d.sqlDB = sqlDB

	db.ConnPool = rewritePool{DB: sqlDB}

	// GORM v1.25 的 Open 不注册默认回调 —— 由 Dialector 负责 (与 mysql/sqlite 驱动一致)。
	callbacks.RegisterDefaultCallbacks(db, &callbacks.Config{})

	// vrtx 的 SQL INSERT 语义要求显式提供整型主键 (引擎行存储定位键), 不像
	// AUTO_INCREMENT 那样省略列由库自增回填; 而 gorm 只有 tag 显式 autoIncrement
	// 才认为主键有库侧默认值 (LastInsertId 回填路径只对这类字段生效)。
	// 折衷: 建 create 前回调, 用引擎游标预分配主键写进对象 —— INSERT 带上该值走
	// 显式主键路径 (预留号免查重), 行为对齐"自增主键 + 回填"的直觉。
	db.Callback().Create().Before("gorm:create").Register("vrtx:assign_pk", d.assignPKCallback)

	return nil
}

// assignPKCallback create 前回调: 零值整型主键 → 引擎 AllocPKID 预分配并写回对象。
func (d *Dialector) assignPKCallback(tx *gorm.DB) {
	stmt := tx.Statement
	if stmt == nil || stmt.Schema == nil || stmt.DryRun {
		return
	}
	pk := stmt.Schema.PrioritizedPrimaryField
	if pk == nil {
		return
	}
	kind := pk.IndirectFieldType.Kind()
	isUint := kind == reflect.Uint || kind == reflect.Uint8 || kind == reflect.Uint16 ||
		kind == reflect.Uint32 || kind == reflect.Uint64
	isInt := kind == reflect.Int || kind == reflect.Int8 || kind == reflect.Int16 ||
		kind == reflect.Int32 || kind == reflect.Int64
	if !isUint && !isInt {
		return
	}
	// 列被 Select/Omit 排除时跳过 (INSERT 不会带它, 预分配反而写错对象)。
	selectColumns, restricted := stmt.SelectAndOmitColumns(true, false)
	if v, present := selectColumns[pk.DBName]; present {
		if !v {
			return
		}
	} else if restricted {
		return
	}
	tbl, ok := d.db.FindTable(stmt.Table)
	if !ok {
		return // 表不存在: 让 INSERT 自己报错
	}
	id := tbl.AllocPKID()
	assign := func(rv reflect.Value) {
		_, isZero := pk.ValueOf(stmt.Context, rv)
		if !isZero {
			return
		}
		if isUint {
			tx.AddError(pk.Set(stmt.Context, rv, id))
		} else {
			tx.AddError(pk.Set(stmt.Context, rv, int64(id)))
		}
	}
	rv := stmt.ReflectValue
	for rv.Kind() == reflect.Ptr {
		if rv.IsNil() {
			return
		}
		rv = rv.Elem()
	}
	switch rv.Kind() {
	case reflect.Struct:
		assign(rv)
	case reflect.Slice, reflect.Array:
		for i := 0; i < rv.Len(); i++ {
			assign(rv.Index(i))
		}
	}
}

// Migrator 返回自定义 Migrator: schema 自省直接调引擎 API (不经 SQL)。
func (d *Dialector) Migrator(db *gorm.DB) gorm.Migrator {
	return Migrator{db: db, eng: d.db, dsn: d.DSN, dialector: d}
}

// explicitTypeMapping 小型 SQL 类型映射: 把 gorm:"type:..." 标签里用户显式声明的
// SQL 类型归一化成引擎侧类型名 (前缀 = 空白或 '(' 之前的部分, 故 "bigint unsigned"
// 归 BIGINT、"varchar(500)" 归 VARCHAR 且宽度取括号值)。独立实现, 不依赖 DB/driver
// 的 SQL 文本解析。第三个返回值为 false 表示不认识 (decimal/json/...), 调用方落回
// Go 类型的默认映射并注释说明。
func explicitTypeMapping(tag string) (string, int, bool) {
	t := strings.ToLower(strings.TrimSpace(tag))
	if t == "" {
		return "", 0, false
	}
	name := t
	if i := strings.IndexAny(name, "( "); i >= 0 {
		name = name[:i]
	}
	size := 0
	if i := strings.IndexByte(t, '('); i >= 0 {
		if j := strings.IndexByte(t[i:], ')'); j > 0 {
			size, _ = strconv.Atoi(strings.TrimSpace(t[i+1 : i+j]))
		}
	}
	switch name {
	case "bigint":
		return "BIGINT", 0, true
	case "int":
		return "INT", 0, true
	case "varchar", "char":
		return "VARCHAR", size, true
	case "text":
		return "TEXT", 0, true
	case "double":
		return "DOUBLE", 0, true
	case "bool", "boolean":
		return "BOOLEAN", 0, true
	case "datetime":
		return "DATETIME", 0, true
	}
	return "", 0, false
}

// DataTypeOf GORM 字段 → 引擎理解的 SQL 类型名 (供 FullDataTypeOf / 高级用法;
// 建表本身走引擎 API, 不经 SQL)。
//
// 显式 gorm:"type:..." 标签优先 (报告 #10 的逃生门: string 字段默认无 size 会落成
// 动态列, 带索引时引擎拒绝建索引 —— 用 type:varchar(500) 精确声明即可绕过默认映射):
// 按 explicitTypeMapping 解析, 认识 → 按其返回; 不认识 → 保持默认映射并附注释说明。
// 默认映射: string 带尺寸 → VARCHAR(size), 超过引擎行内定长上限 2048 时截断声明并
// 注释说明; 无尺寸 → TEXT (动态列)。
func (d *Dialector) DataTypeOf(field *schema.Field) string {
	if tag, ok := field.TagSettings["TYPE"]; ok {
		if name, size, ok := explicitTypeMapping(tag); ok {
			if name != "VARCHAR" {
				return name
			}
			switch {
			case size <= 0:
				return "VARCHAR /* 声明无宽度, 建列按动态列处理 */"
			case size > types.MaxInlineColumnWidth:
				return fmt.Sprintf("VARCHAR(2048) /* 声明 %d 超引擎行内上限 2048, 截断 */", size)
			default:
				return fmt.Sprintf("VARCHAR(%d)", size)
			}
		}
		return fmt.Sprintf("%s /* 未知 type:%s, 按 Go 类型默认映射 */", d.defaultDataTypeOf(field), tag)
	}
	return d.defaultDataTypeOf(field)
}

// defaultDataTypeOf 按 Go 类型给默认 SQL 类型名 (DataTypeOf 的回落分支)。
func (d *Dialector) defaultDataTypeOf(field *schema.Field) string {
	switch field.DataType {
	case schema.Uint:
		return "BIGINT UNSIGNED"
	case schema.Int:
		return "BIGINT"
	case schema.Float:
		return "DOUBLE"
	case schema.String:
		if field.Size > 0 {
			if field.Size > types.MaxInlineColumnWidth {
				return fmt.Sprintf("VARCHAR(2048) /* 声明 %d 超引擎行内上限 2048, 截断 */", field.Size)
			}
			return fmt.Sprintf("VARCHAR(%d)", field.Size)
		}
		return "TEXT"
	case schema.Bool:
		return "BOOLEAN"
	case schema.Time:
		return "DATETIME" // 引擎约定: 行内 8B Unix 毫秒
	case schema.Bytes:
		return "BLOB"
	default:
		// struct / JSON 等: 序列化为文本存动态列
		return "TEXT"
	}
}

// DefaultValueOf 引擎不支持列级 DEFAULT (未提供的列按 NULL/零值裁决), 给出占位表达式。
func (d *Dialector) DefaultValueOf(field *schema.Field) clause.Expression {
	return clause.Expr{SQL: "DEFAULT"}
}

// BindVarTo 占位符: 与 mysql 一致, 一律 "?" (驱动只认 ? 占位符)。
func (d *Dialector) BindVarTo(w clause.Writer, stmt *gorm.Statement, v interface{}) {
	w.WriteByte('?')
}

// QuoteTo 反引号包裹; 含 "." 时按段分别引用: t.col → `t`.`col`。
func (d *Dialector) QuoteTo(w clause.Writer, s string) {
	for i, part := range strings.Split(s, ".") {
		if i > 0 {
			w.WriteByte('.')
		}
		w.WriteByte('`')
		w.WriteString(strings.ReplaceAll(part, "`", "``"))
		w.WriteByte('`')
	}
}

// Explain 把 "?" 按序替换为字面值 (字符串单引号包裹并转义, 数字直写,
// bool → 1/0, nil → NULL)。仅用于 SQL 日志展示。
func (d *Dialector) Explain(sql string, vars ...interface{}) string {
	var sb strings.Builder
	idx := 0
	for i := 0; i < len(sql); i++ {
		c := sql[i]
		if c != '?' || idx >= len(vars) {
			sb.WriteByte(c)
			continue
		}
		sb.WriteString(literalOf(vars[idx]))
		idx++
	}
	return sb.String()
}

// literalOf 把一个参数值格式化为 SQL 字面量。
func literalOf(v interface{}) string {
	switch x := v.(type) {
	case nil:
		return "NULL"
	case string:
		return "'" + strings.ReplaceAll(x, "'", "''") + "'"
	case []byte:
		return "'" + strings.ReplaceAll(string(x), "'", "''") + "'"
	case bool:
		if x {
			return "1"
		}
		return "0"
	case time.Time:
		return "'" + x.Format("2006-01-02 15:04:05.999") + "'"
	}
	rv := reflect.ValueOf(v)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return strconv.FormatInt(rv.Int(), 10)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return strconv.FormatUint(rv.Uint(), 10)
	case reflect.Float32, reflect.Float64:
		return strconv.FormatFloat(rv.Float(), 'g', -1, 64)
	}
	return "'" + strings.ReplaceAll(fmt.Sprint(v), "'", "''") + "'"
}

// =============================================================================
// rewritePool: SQL 进驱动前的轻量改写 (内联 LIMIT/OFFSET 整数参数)
// =============================================================================

// 注: 早期驱动解析器不支持表限定列名, 这里曾剥掉 "`t`." 前缀; 驱动现已原生
// 支持限定名 (单表剥限定 + JOIN 按别名归属), 剥前缀反而会制造歧义 (JOIN 下
// 两表同名列), 故只保留 LIMIT/OFFSET 参数内联。

// limitClauseRe 匹配 LIMIT/OFFSET 占位符形态 (GORM 把 Limit(1) 编成 "LIMIT ?",
// 而 vrtx 解析器只接受字面量整数, 需要按参数内联)。
var limitClauseRe = regexp.MustCompile(`(?i)\b(LIMIT|OFFSET)[ \t\r\n]+\?`)

// rewriteSQL 进驱动前的统一改写: 内联 LIMIT/OFFSET 整数参数 (限定名保持原样)。
func rewriteSQL(query string, args []interface{}) (string, []interface{}) {
	if !limitClauseRe.MatchString(query) {
		return query, args
	}
	// 找出 LIMIT/OFFSET 子句里 '?' 是整条 SQL 的第几个 '?' (即参数下标),
	// 把它替换成字面量并把对应参数从 args 剔除 (后面的参数依次前移)。
	mark := map[int]string{} // 第k个 '?' → 字面量
	for _, loc := range limitClauseRe.FindAllStringSubmatchIndex(query, -1) {
		qpos := strings.IndexByte(query[loc[0]:loc[1]], '?')
		if qpos < 0 {
			continue
		}
		abs := loc[0] + qpos
		k := 0
		for i := 0; i < abs; i++ {
			if query[i] == '?' {
				k++
			}
		}
		if k >= len(args) {
			continue
		}
		mark[k] = intLiteral(args[k])
	}
	if len(mark) == 0 {
		return query, args
	}
	var sb strings.Builder
	argIdx := 0
	for i := 0; i < len(query); i++ {
		c := query[i]
		if c != '?' {
			sb.WriteByte(c)
			continue
		}
		if lit, ok := mark[argIdx]; ok {
			sb.WriteString(lit)
		} else {
			sb.WriteByte('?')
		}
		argIdx++
	}
	kept := make([]interface{}, 0, len(args)-len(mark))
	for i, a := range args {
		if _, ok := mark[i]; !ok {
			kept = append(kept, a)
		}
	}
	return sb.String(), kept
}

// intLiteral 把 LIMIT/OFFSET 参数格式化为整数字面量; 不认识的类型原样占位返回 "0"。
func intLiteral(v interface{}) string {
	switch x := v.(type) {
	case int:
		return strconv.Itoa(x)
	case int8:
		return strconv.FormatInt(int64(x), 10)
	case int16:
		return strconv.FormatInt(int64(x), 10)
	case int32:
		return strconv.FormatInt(int64(x), 10)
	case int64:
		return strconv.FormatInt(x, 10)
	case uint:
		return strconv.FormatUint(uint64(x), 10)
	case uint8:
		return strconv.FormatUint(uint64(x), 10)
	case uint16:
		return strconv.FormatUint(uint64(x), 10)
	case uint32:
		return strconv.FormatUint(uint64(x), 10)
	case uint64:
		return strconv.FormatUint(x, 10)
	case float64:
		return strconv.FormatInt(int64(x), 10)
	case string:
		return x
	}
	return "0"
}

// rewritePool 包装 *sql.DB, SQL 进驱动前统一过 rewriteSQL (剥 “ `t`. “ 限定前缀、
// 内联 LIMIT/OFFSET)。嵌入 *sql.DB 顺带获得 BeginTx (TxBeginner) 等全部池能力。
type rewritePool struct {
	*sql.DB
}

func (p rewritePool) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
	q, _ := rewriteSQL(query, nil)
	return p.DB.PrepareContext(ctx, q)
}

func (p rewritePool) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	q, a := rewriteSQL(query, args)
	return p.DB.ExecContext(ctx, q, a...)
}

func (p rewritePool) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	q, a := rewriteSQL(query, args)
	return p.DB.QueryContext(ctx, q, a...)
}

func (p rewritePool) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
	q, a := rewriteSQL(query, args)
	return p.DB.QueryRowContext(ctx, q, a...)
}