db/dialect/sqlite3/vtab/vtab.go

//go:build sqlite_vtable || vtable

// Package vtab 在 mattn/go-sqlite3 上提供干净的虚拟表框架。
//
// 编译要求：需加 build tag `-tags sqlite_vtable`
//
// 基本用法：
//
//	vtab.Register("my_module", &MyModule{})
//	db, _ := sql.Open(vtab.DriverName, ":memory:")
//	db.Exec(`CREATE VIRTUAL TABLE t USING my_module(arg1, arg2)`)
//	db.Query(`SELECT * FROM t WHERE col = ?`, value)
package vtab

import (
	"database/sql"
	"fmt"
	"sync"
	"time"

	"git.fsdpf.net/go/db"
	"git.fsdpf.net/go/db/exp"
	gosqlite3 "github.com/mattn/go-sqlite3"
)

// DriverName 是支持虚拟表的 SQLite3 驱动名称，同时内置 IF() 函数支持。
// 构建时需指定 tag：-tags vtable
const DriverName = "sqlite3_vtab"

// Op 是 WHERE 约束的操作符类型，直接复用 go-sqlite3 的定义。
type Op = gosqlite3.Op

// 操作符常量，与 SQLite C API 值一致。
const (
	OpEQ   Op = gosqlite3.OpEQ   // =
	OpGT   Op = gosqlite3.OpGT   // >
	OpLE   Op = gosqlite3.OpLE   // <=
	OpLT   Op = gosqlite3.OpLT   // <
	OpGE   Op = gosqlite3.OpGE   // >=
	OpLIKE Op = gosqlite3.OpLIKE // LIKE

	// OpLIMIT / OpOFFSET：go-sqlite3 尚未导出这两个常量，直接使用 SQLite C API 原始值。
	// BestIndex 中将它们标记为 Used=true 后，Filter 可收到 LIMIT / OFFSET 的实际值。
	// 注意：这两个约束的 Column 字段为 -1（不对应任何列）。
	OpLIMIT  Op = 73 // SQLITE_INDEX_CONSTRAINT_LIMIT
	OpOFFSET Op = 74 // SQLITE_INDEX_CONSTRAINT_OFFSET
)

// ConstraintInfo 是约束信息，BestIndex 和 Filter 阶段共用。
//   - BestIndex 阶段：Value 为 nil，Usable 表示 SQLite 是否允许使用该约束。
//   - Filter 阶段：Value 为实际约束值（已绑定），Usable 始终为 true。
type ConstraintInfo struct {
	Column int  // 列索引，对应 DeclareVTab 中列的顺序（从 0 开始）；OpLIMIT/OpOFFSET 时为 -1
	Op     Op   // 操作符
	Usable bool // BestIndex 阶段有效；Filter 阶段忽略
	Value  any  // Filter 阶段由适配层填充；BestIndex 阶段为 nil
}

// OrderByInfo 是查询中的 ORDER BY 信息。
type OrderByInfo struct {
	Column int  // 列索引
	Desc   bool // true 表示降序
}

// IndexOutput 是 BestIndex 的返回值，告知 SQLite 本表能处理哪些约束。
type IndexOutput struct {
	// Used[i]=true 表示第 i 个约束由本表自行处理。
	// 对应约束的值会按原顺序在 Filter.constraintValues 中传入。
	// len(Used) 必须等于传入 BestIndex 的 constraints 长度。
	Used []bool

	// IdxNum 和 IdxStr 是传给 Filter 的不透明标识，用于区分不同查询计划。
	IdxNum int
	IdxStr string

	// AlreadyOrdered 为 true 时 SQLite 不再对结果二次排序。
	AlreadyOrdered bool

	// EstimatedCost 扫描代价估算（越小越优先），默认 0 表示交给 SQLite 决定。
	EstimatedCost float64
	// EstimatedRows 预估返回行数，默认 0。
	EstimatedRows float64
}

func DialectOptions() *db.SQLDialectOptions {
	opts := db.DefaultDialectOptions()

	opts.SupportsReturn = true
	opts.SupportsOrderByOnUpdate = true
	opts.SupportsLimitOnUpdate = true
	opts.SupportsOrderByOnDelete = true
	opts.SupportsLimitOnDelete = true
	opts.SupportsConflictUpdateWhere = false
	opts.SupportsInsertIgnoreSyntax = true
	opts.SupportsConflictTarget = true
	opts.SupportsMultipleUpdateTables = false
	opts.WrapCompoundsInParens = false
	opts.SupportsDistinct = false // 设为 false 可全局禁止生成 DISTINCT 关键字
	opts.SupportsDistinctOn = false
	opts.SupportsWindowFunction = false
	opts.SupportsLateral = false

	opts.PlaceHolderFragment = []byte("?")
	opts.IncludePlaceholderNum = false
	opts.QuoteRune = '`'
	opts.DefaultValuesFragment = []byte("")
	opts.True = []byte("1")
	opts.False = []byte("0")
	opts.TimeFormat = time.RFC3339Nano
	opts.BooleanOperatorLookup = map[exp.BooleanOperation][]byte{
		exp.EqOp:             []byte("="),
		exp.NeqOp:            []byte("!="),
		exp.GtOp:             []byte(">"),
		exp.GteOp:            []byte(">="),
		exp.LtOp:             []byte("<"),
		exp.LteOp:            []byte("<="),
		exp.InOp:             []byte("IN"),
		exp.NotInOp:          []byte("NOT IN"),
		exp.IsOp:             []byte("IS"),
		exp.IsNotOp:          []byte("IS NOT"),
		exp.LikeOp:           []byte("LIKE"),
		exp.NotLikeOp:        []byte("NOT LIKE"),
		exp.ILikeOp:          []byte("LIKE"),
		exp.NotILikeOp:       []byte("NOT LIKE"),
		exp.RegexpLikeOp:     []byte("REGEXP"),
		exp.RegexpNotLikeOp:  []byte("NOT REGEXP"),
		exp.RegexpILikeOp:    []byte("REGEXP"),
		exp.RegexpNotILikeOp: []byte("NOT REGEXP"),
	}
	opts.UseLiteralIsBools = false
	opts.BitwiseOperatorLookup = map[exp.BitwiseOperation][]byte{
		exp.BitwiseOrOp:         []byte("|"),
		exp.BitwiseAndOp:        []byte("&"),
		exp.BitwiseLeftShiftOp:  []byte("<<"),
		exp.BitwiseRightShiftOp: []byte(">>"),
	}
	opts.EscapedRunes = map[rune][]byte{
		'\'': []byte("''"),
	}
	opts.InsertIgnoreClause = []byte("INSERT OR IGNORE INTO ")
	opts.ConflictFragment = []byte(" ON CONFLICT ")
	opts.ConflictDoUpdateFragment = []byte(" DO UPDATE SET ")
	opts.ConflictDoNothingFragment = []byte(" DO NOTHING ")
	opts.ForUpdateFragment = []byte("")
	opts.OfFragment = []byte("")
	opts.NowaitFragment = []byte("")
	return opts
}

var (
	registryMu sync.RWMutex
	registry   = map[string]Module{}
)

// Register 注册一个虚拟表模块。
// 必须在第一次打开数据库连接之前调用。
// 注册后，可在 SQL 中使用：CREATE VIRTUAL TABLE t USING moduleName(args...)
func Register(moduleName string, m Module) {
	registryMu.Lock()
	registry[moduleName] = m
	registryMu.Unlock()
}

func init() {
	sql.Register(DriverName, &gosqlite3.SQLiteDriver{
		ConnectHook: func(conn *gosqlite3.SQLiteConn) error {
			// 内置 IF(cond, trueVal, falseVal) 函数
			if err := conn.RegisterFunc("IF", func(cond int64, trueVal, falseVal interface{}) interface{} {
				if cond != 0 {
					return trueVal
				}
				return falseVal
			}, true); err != nil {
				return err
			}
			// 注册所有已登记的虚拟表模块
			registryMu.RLock()
			defer registryMu.RUnlock()

			for name, m := range registry {
				if err := conn.CreateModule(name, &moduleAdapter{mod: m}); err != nil {
					return fmt.Errorf("vtab: 注册模块 %q 失败: %w", name, err)
				}
			}
			return nil
		},
	})

	db.RegisterDialect("vtable", DialectOptions())
}