openteam/server/internal/store/db.go

package store

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
	gormlogger "gorm.io/gorm/logger"
)

// schemaVersion 当前 schema 版本。struct 变更（加列/改列/删列）时递增，
// 触发一次 AutoMigrate 并把新版本写入库（SQLite 用 PRAGMA user_version）。
// AutoMigrate 对已有表的列判定不收敛（每次都重建表：CREATE __temp + INSERT SELECT + DROP），
// 大表上一次重建数十秒且每次重启重演，所以之后版本未变就直接跳过。
//
// 注意：c39c3e7 给 Model 加了 cache_write_price/context_length/supports_tools/
// supports_vision/supports_reasoning 等列并删掉了 display_name，但当时漏改这里，
// 导致已建库的实例永远跳过重迁移、新增模型时 Create 因缺列而失败。这里补到 3 触发一次迁移。
// v4：users 加 token_version（会话撤销）。
// v5：usage_logs 补 key_id/channel_id 索引与 (user,created)/(channel,created) 复合索引。
// v6：修正 v5 复合索引列序——AutoMigrate 按名跳过已存在的索引，旧库上的
// idx_user_created 是 user_id 单列，光改 struct tag 不会重建，需删掉重迁一次。
// v7：usage_logs 加 reasoning_tokens / cache_write_1h_tokens / served_model 三列。
// v8：去掉 usage_logs 上没人查询的 key_id / trace_id 索引（每行两次多余的 B-tree 插入）。
// v9：同样去掉没人查询的 idx_channel_created。
const schemaVersion = 9

// Open 打开数据库连接并自动迁移。
// 开发默认 SQLite（自动追加 WAL 与 busy_timeout，见 withSQLitePragmas），生产可切 postgres。
func Open(driver, dsn string) (*gorm.DB, error) {
	var dialector gorm.Dialector
	switch driver {
	case "postgres":
		dialector = postgresDialector(dsn)
	default:
		dsn = withSQLitePragmas(dsn)
		// 确保 SQLite 文件所在目录存在
		if dir := sqliteDir(dsn); dir != "" {
			_ = os.MkdirAll(dir, 0o755)
		}
		dialector = sqlite.Open(dsn)
	}

	// PrepareStmt 让 GORM 缓存预编译语句：实测请求路径上 sqlite3Prepare（SQL 解析）
	// 占 CPU 的 15%，每个请求十来个查询全在重复解析同几条 SQL。
	db, err := gorm.Open(dialector, &gorm.Config{
		Logger:      gormlogger.Default.LogMode(gormlogger.Warn),
		PrepareStmt: true,
	})
	if err != nil {
		return nil, err
	}
	configurePool(db, driver)

	if driver != "postgres" && currentSQLiteVersion(db) >= schemaVersion {
		log.Printf("store: connected driver=%s (schema up-to-date v%d, skip migrate)", driver, schemaVersion)
		return db, nil
	}

	if err := db.AutoMigrate(AllModels()...); err != nil {
		return nil, err
	}
	reshapeUsageLogIndexes(db)
	dropUnusedUsageLogIndexes(db)
	if driver != "postgres" {
		setSQLiteVersion(db, schemaVersion)
	}
	log.Printf("store: connected driver=%s (migrated, schema v%d)", driver, schemaVersion)
	// 将已有渠道的超时时间从 120000ms 更新为 300000ms（幂等操作）
	db.Model(&Channel{}).Where("timeout_ms = ?", 120000).Update("timeout_ms", 300000)
	return db, nil
}

// reshapeUsageLogIndexes 一次性重建 usage_logs 的复合索引：AutoMigrate 按名跳过已存在的
// 索引，历史库上的 idx_user_created 是 user_id 单列、idx_channel_created 形状同样不对，
// struct tag 改了也不会生效。删掉后紧接的 AutoMigrate 按新定义重建。
// 用 system_configs 打标记：PG 每次启动都走迁移路径，没有标记会反复重建索引。
func reshapeUsageLogIndexes(db *gorm.DB) {
	var done string
	db.Model(&SystemConfig{}).Where("key = ?", "index_reshape_v6").Pluck("value", &done)
	if done != "" {
		return
	}
	for _, name := range []string{"idx_user_created", "idx_channel_created"} {
		if db.Migrator().HasIndex(&UsageLog{}, name) {
			if err := db.Migrator().DropIndex(&UsageLog{}, name); err != nil {
				log.Printf("store: drop legacy index %s: %v", name, err)
				return
			}
		}
	}
	if err := db.AutoMigrate(&UsageLog{}); err != nil {
		log.Printf("store: recreate usage_logs indexes: %v", err)
		return
	}
	db.Create(&SystemConfig{Key: "index_reshape_v6", Value: `"done"`})
	log.Printf("store: usage_logs indexes reshaped (user+created / channel+created)")
}

// withSQLitePragmas 给 DSN 补上 busy_timeout 与 WAL。
// glebarez（modernc）驱动只解析 _pragma/_txlock/_time_format/vfs 四种参数——
// 旧注释宣称的 `?_journal_mode=WAL` 实际不生效，缺 WAL 时并发写会频繁 database is locked。
// 用户 DSN 里已显式声明的项不覆盖。
// SQLite 连接参数。实测（单表 2000 条记账、暖机后）：
//   - 默认 deferred 事务：并发写升级锁失败会触发 SQLITE_BUSY 重试，吞吐 11–20k rows/s
//     且抖动极大；
//   - _txlock=immediate：写事务在 BEGIN 时就拿写锁，不再有升级冲突，28k rows/s；
//   - 再加 synchronous(NORMAL)：WAL 下的标准配置（提交不 fsync），稳定 30k rows/s，
//     波动 <1%。
//
// 这两项都是 SQLite 官方推荐的并发写法，代价是断电时可能丢最后一个事务。
func withSQLitePragmas(dsn string) string {
	for _, p := range []string{
		"_pragma=busy_timeout(10000)",
		"_pragma=journal_mode(WAL)",
		"_pragma=synchronous(NORMAL)",
		"_txlock=immediate",
	} {
		if strings.Contains(dsn, p) {
			continue
		}
		if strings.Contains(dsn, "?") {
			dsn += "&" + p
		} else {
			dsn += "?" + p
		}
	}
	return dsn
}

// configurePool 连接池上限。此前全仓无任何配置：PG 连接数无上限，
// SQLite 多连接并发写叠加无 WAL 加剧锁竞争。
// dropUnusedUsageLogIndexes 删除历史库里已经没用的索引：AutoMigrate 不会自动删索引，
// 去掉 struct tag 只是不再新建，旧库上它们还会继续拖慢每次插入。
func dropUnusedUsageLogIndexes(db *gorm.DB) {
	for _, name := range []string{"idx_usage_logs_key_id", "idx_usage_logs_trace_id", "idx_channel_created"} {
		if err := db.Exec("DROP INDEX IF EXISTS " + name).Error; err != nil {
			log.Printf("store: drop index %s: %v", name, err)
		}
	}
}

func configurePool(db *gorm.DB, driver string) {
	sqlDB, err := db.DB()
	if err != nil {
		return
	}
	if driver == "postgres" {
		sqlDB.SetMaxOpenConns(30)
		sqlDB.SetMaxIdleConns(10)
		sqlDB.SetConnMaxLifetime(30 * time.Minute)
		sqlDB.SetConnMaxIdleTime(5 * time.Minute)
		return
	}
	// SQLite：WAL 下读写可并发，写仍串行——限连接数让写竞争在 busy_timeout 内排队
	sqlDB.SetMaxOpenConns(8)
	sqlDB.SetConnMaxIdleTime(10 * time.Minute)
}

// currentSQLiteVersion 读取 PRAGMA user_version。
func currentSQLiteVersion(db *gorm.DB) int {
	var v int
	db.Raw("PRAGMA user_version").Scan(&v)
	return v
}

// setSQLiteVersion 写入 PRAGMA user_version。
func setSQLiteVersion(db *gorm.DB, v int) {
	db.Exec(fmt.Sprintf("PRAGMA user_version = %d", v))
}

// sqliteDir 提取 SQLite DSN 中的目录部分（忽略 file: 前缀与查询参数）。
func sqliteDir(dsn string) string {
	d := dsn
	if i := strings.IndexByte(d, '?'); i >= 0 {
		d = d[:i]
	}
	if strings.HasPrefix(d, "file:") {
		d = d[len("file:"):]
	}
	if d == "" || d == ":memory:" || strings.Contains(d, "::") {
		return ""
	}
	return filepath.Dir(d)
}