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) }