Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

789 строки
23 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "log"
  5. "os"
  6. "strings"
  7. "sync"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/constant"
  11. "github.com/glebarez/sqlite"
  12. "gorm.io/driver/mysql"
  13. "gorm.io/driver/postgres"
  14. "gorm.io/gorm"
  15. )
  16. var commonGroupCol string
  17. var commonKeyCol string
  18. var commonTrueVal string
  19. var commonFalseVal string
  20. var logKeyCol string
  21. var logGroupCol string
  22. func initCol() {
  23. // init common column names
  24. if common.UsingPostgreSQL {
  25. commonGroupCol = `"group"`
  26. commonKeyCol = `"key"`
  27. commonTrueVal = "true"
  28. commonFalseVal = "false"
  29. } else {
  30. commonGroupCol = "`group`"
  31. commonKeyCol = "`key`"
  32. commonTrueVal = "1"
  33. commonFalseVal = "0"
  34. }
  35. if os.Getenv("LOG_SQL_DSN") != "" {
  36. switch common.LogSqlType {
  37. case common.DatabaseTypePostgreSQL:
  38. logGroupCol = `"group"`
  39. logKeyCol = `"key"`
  40. default:
  41. logGroupCol = commonGroupCol
  42. logKeyCol = commonKeyCol
  43. }
  44. } else {
  45. // LOG_SQL_DSN 为空时,日志数据库与主数据库相同
  46. if common.UsingPostgreSQL {
  47. logGroupCol = `"group"`
  48. logKeyCol = `"key"`
  49. } else {
  50. logGroupCol = commonGroupCol
  51. logKeyCol = commonKeyCol
  52. }
  53. }
  54. // log sql type and database type
  55. //common.SysLog("Using Log SQL Type: " + common.LogSqlType)
  56. }
  57. var DB *gorm.DB
  58. var LOG_DB *gorm.DB
  59. func createRootAccountIfNeed() error {
  60. var user User
  61. //if user.Status != common.UserStatusEnabled {
  62. if err := DB.First(&user).Error; err != nil {
  63. common.SysLog("no user exists, create a root user for you: username is root, password is 123456")
  64. hashedPassword, err := common.Password2Hash("123456")
  65. if err != nil {
  66. return err
  67. }
  68. rootUser := User{
  69. Username: "root",
  70. Password: hashedPassword,
  71. Role: common.RoleRootUser,
  72. Status: common.UserStatusEnabled,
  73. DisplayName: "Root User",
  74. AccessToken: nil,
  75. Quota: 100000000,
  76. }
  77. DB.Create(&rootUser)
  78. }
  79. return nil
  80. }
  81. func CheckSetup() {
  82. setup := GetSetup()
  83. if setup == nil {
  84. // No setup record exists, check if we have a root user
  85. if RootUserExists() {
  86. common.SysLog("system is not initialized, but root user exists")
  87. // Create setup record
  88. newSetup := Setup{
  89. Version: common.Version,
  90. InitializedAt: time.Now().Unix(),
  91. }
  92. err := DB.Create(&newSetup).Error
  93. if err != nil {
  94. common.SysLog("failed to create setup record: " + err.Error())
  95. }
  96. constant.Setup = true
  97. } else {
  98. common.SysLog("system is not initialized and no root user exists")
  99. constant.Setup = false
  100. }
  101. } else {
  102. // Setup record exists, system is initialized
  103. common.SysLog("system is already initialized at: " + time.Unix(setup.InitializedAt, 0).String())
  104. constant.Setup = true
  105. }
  106. }
  107. func chooseDB(envName string, isLog bool) (*gorm.DB, error) {
  108. defer func() {
  109. initCol()
  110. }()
  111. dsn := os.Getenv(envName)
  112. if dsn != "" {
  113. if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
  114. // Use PostgreSQL
  115. common.SysLog("using PostgreSQL as database")
  116. if !isLog {
  117. common.UsingPostgreSQL = true
  118. } else {
  119. common.LogSqlType = common.DatabaseTypePostgreSQL
  120. }
  121. return gorm.Open(postgres.New(postgres.Config{
  122. DSN: dsn,
  123. PreferSimpleProtocol: true, // disables implicit prepared statement usage
  124. }), &gorm.Config{
  125. PrepareStmt: true, // precompile SQL
  126. })
  127. }
  128. if strings.HasPrefix(dsn, "local") {
  129. common.SysLog("SQL_DSN not set, using SQLite as database")
  130. if !isLog {
  131. common.UsingSQLite = true
  132. } else {
  133. common.LogSqlType = common.DatabaseTypeSQLite
  134. }
  135. return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
  136. PrepareStmt: true, // precompile SQL
  137. })
  138. }
  139. // Use MySQL
  140. common.SysLog("using MySQL as database")
  141. // check parseTime
  142. if !strings.Contains(dsn, "parseTime") {
  143. if strings.Contains(dsn, "?") {
  144. dsn += "&parseTime=true"
  145. } else {
  146. dsn += "?parseTime=true"
  147. }
  148. }
  149. if !isLog {
  150. common.UsingMySQL = true
  151. } else {
  152. common.LogSqlType = common.DatabaseTypeMySQL
  153. }
  154. return gorm.Open(mysql.Open(dsn), &gorm.Config{
  155. PrepareStmt: true, // precompile SQL
  156. })
  157. }
  158. // Use SQLite
  159. common.SysLog("SQL_DSN not set, using SQLite as database")
  160. common.UsingSQLite = true
  161. return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
  162. PrepareStmt: true, // precompile SQL
  163. })
  164. }
  165. func InitDB() (err error) {
  166. db, err := chooseDB("SQL_DSN", false)
  167. if err == nil {
  168. if common.DebugEnabled {
  169. db = db.Debug()
  170. }
  171. DB = db
  172. // MySQL charset/collation startup check: ensure Chinese-capable charset
  173. if common.UsingMySQL {
  174. if err := checkMySQLChineseSupport(DB); err != nil {
  175. panic(err)
  176. }
  177. }
  178. sqlDB, err := DB.DB()
  179. if err != nil {
  180. return err
  181. }
  182. sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100))
  183. sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000))
  184. sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60)))
  185. if !common.IsMasterNode {
  186. return nil
  187. }
  188. if common.UsingMySQL {
  189. //_, _ = sqlDB.Exec("ALTER TABLE channels MODIFY model_mapping TEXT;") // TODO: delete this line when most users have upgraded
  190. }
  191. common.SysLog("database migration started")
  192. err = migrateDB()
  193. if err != nil {
  194. return err
  195. }
  196. LoadEmailQuotaCache()
  197. return nil
  198. } else {
  199. common.FatalLog(err)
  200. }
  201. return err
  202. }
  203. func InitLogDB() (err error) {
  204. if os.Getenv("LOG_SQL_DSN") == "" {
  205. LOG_DB = DB
  206. return
  207. }
  208. db, err := chooseDB("LOG_SQL_DSN", true)
  209. if err == nil {
  210. if common.DebugEnabled {
  211. db = db.Debug()
  212. }
  213. LOG_DB = db
  214. // If log DB is MySQL, also ensure Chinese-capable charset
  215. if common.LogSqlType == common.DatabaseTypeMySQL {
  216. if err := checkMySQLChineseSupport(LOG_DB); err != nil {
  217. panic(err)
  218. }
  219. }
  220. sqlDB, err := LOG_DB.DB()
  221. if err != nil {
  222. return err
  223. }
  224. sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100))
  225. sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000))
  226. sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60)))
  227. if !common.IsMasterNode {
  228. return nil
  229. }
  230. common.SysLog("database migration started")
  231. err = migrateLOGDB()
  232. return err
  233. } else {
  234. common.FatalLog(err)
  235. }
  236. return err
  237. }
  238. func migrateDB() error {
  239. // Migrate price_amount column from float/double to decimal for existing tables
  240. migrateSubscriptionPlanPriceAmount()
  241. // Drop bound_channel_id column from tokens table (deprecated field)
  242. migrateTokenDropBoundChannelId()
  243. err := DB.AutoMigrate(
  244. &Channel{},
  245. &ChannelAssetCredential{},
  246. &Token{},
  247. &User{},
  248. &PasskeyCredential{},
  249. &Option{},
  250. &Redemption{},
  251. &Ability{},
  252. // &Log{} is intentionally omitted here. On PostgreSQL the logs table is a
  253. // partitioned table managed by the pg_partman extension; GORM's AutoMigrate
  254. // would try to alter its composite primary key (id, created_at) and create
  255. // redundant per-partition indexes, breaking startup. The logs table is
  256. // handled separately by ensureLogTable() below. On SQLite/MySQL it falls
  257. // back to the original AutoMigrate behavior.
  258. &Midjourney{},
  259. &TopUp{},
  260. &QuotaData{},
  261. &Task{},
  262. &Model{},
  263. &Vendor{},
  264. &PrefillGroup{},
  265. &Setup{},
  266. &TwoFA{},
  267. &TwoFABackupCode{},
  268. &Checkin{},
  269. &SubscriptionOrder{},
  270. &UserSubscription{},
  271. &SubscriptionPreConsumeRecord{},
  272. &CustomOAuthProvider{},
  273. &UserOAuthBinding{},
  274. &PendingSyncRecord{},
  275. &QuotaSyncLog{},
  276. &EmailQuotaRule{},
  277. &UserModelRateLimit{},
  278. &UserAssetChannel{},
  279. &UserAssetGroup{},
  280. &UserMigrationBatch{},
  281. &UserMigrationItem{},
  282. &MigrationQuotaGrant{},
  283. )
  284. if err != nil {
  285. return err
  286. }
  287. if err := ensureLogTable(); err != nil {
  288. return err
  289. }
  290. if err := DB.Exec("DROP TABLE IF EXISTS channel_pricings").Error; err != nil {
  291. return err
  292. }
  293. if err := DB.Exec("DROP TABLE IF EXISTS pricing_tags").Error; err != nil {
  294. return err
  295. }
  296. if common.UsingSQLite {
  297. if err := ensureSubscriptionPlanTableSQLite(); err != nil {
  298. return err
  299. }
  300. } else {
  301. if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil {
  302. return err
  303. }
  304. }
  305. // 将现有 sort_order=0 的模型和供应商更新为默认大数
  306. DB.Model(&Model{}).Where("sort_order = 0").Update("sort_order", 999999)
  307. DB.Model(&Vendor{}).Where("sort_order = 0").Update("sort_order", 999999)
  308. migrateChannelPublicName()
  309. return nil
  310. }
  311. func migrateDBFast() error {
  312. // Drop bound_channel_id column from tokens table (deprecated field)
  313. migrateTokenDropBoundChannelId()
  314. var wg sync.WaitGroup
  315. migrations := []struct {
  316. model interface{}
  317. name string
  318. }{
  319. {&Channel{}, "Channel"},
  320. {&ChannelAssetCredential{}, "ChannelAssetCredential"},
  321. {&Token{}, "Token"},
  322. {&User{}, "User"},
  323. {&PasskeyCredential{}, "PasskeyCredential"},
  324. {&Option{}, "Option"},
  325. {&Redemption{}, "Redemption"},
  326. {&Ability{}, "Ability"},
  327. // &Log{} omitted: see ensureLogTable() / migrateDB() for rationale
  328. // (pg_partman-managed partitioned table on PostgreSQL).
  329. {&Midjourney{}, "Midjourney"},
  330. {&TopUp{}, "TopUp"},
  331. {&QuotaData{}, "QuotaData"},
  332. {&Task{}, "Task"},
  333. {&Model{}, "Model"},
  334. {&Vendor{}, "Vendor"},
  335. {&PrefillGroup{}, "PrefillGroup"},
  336. {&Setup{}, "Setup"},
  337. {&TwoFA{}, "TwoFA"},
  338. {&TwoFABackupCode{}, "TwoFABackupCode"},
  339. {&Checkin{}, "Checkin"},
  340. {&SubscriptionOrder{}, "SubscriptionOrder"},
  341. {&UserSubscription{}, "UserSubscription"},
  342. {&SubscriptionPreConsumeRecord{}, "SubscriptionPreConsumeRecord"},
  343. {&CustomOAuthProvider{}, "CustomOAuthProvider"},
  344. {&UserOAuthBinding{}, "UserOAuthBinding"},
  345. {&PendingSyncRecord{}, "PendingSyncRecord"},
  346. {&QuotaSyncLog{}, "QuotaSyncLog"},
  347. {&EmailQuotaRule{}, "EmailQuotaRule"},
  348. {&UserAssetChannel{}, "UserAssetChannel"},
  349. {&UserAssetGroup{}, "UserAssetGroup"},
  350. {&UserMigrationBatch{}, "UserMigrationBatch"},
  351. {&UserMigrationItem{}, "UserMigrationItem"},
  352. {&MigrationQuotaGrant{}, "MigrationQuotaGrant"},
  353. }
  354. // 动态计算migration数量,确保errChan缓冲区足够大
  355. errChan := make(chan error, len(migrations))
  356. for _, m := range migrations {
  357. wg.Add(1)
  358. go func(model interface{}, name string) {
  359. defer wg.Done()
  360. if err := DB.AutoMigrate(model); err != nil {
  361. errChan <- fmt.Errorf("failed to migrate %s: %v", name, err)
  362. }
  363. }(m.model, m.name)
  364. }
  365. // Wait for all migrations to complete
  366. wg.Wait()
  367. close(errChan)
  368. // Check for any errors
  369. for err := range errChan {
  370. if err != nil {
  371. return err
  372. }
  373. }
  374. if err := ensureLogTable(); err != nil {
  375. return err
  376. }
  377. if common.UsingSQLite {
  378. if err := ensureSubscriptionPlanTableSQLite(); err != nil {
  379. return err
  380. }
  381. } else {
  382. if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil {
  383. return err
  384. }
  385. }
  386. common.SysLog("database migrated")
  387. return nil
  388. }
  389. // ensureLogTable handles the logs table, which behaves differently depending on
  390. // the database backend:
  391. //
  392. // - SQLite / MySQL: behaves as before, GORM AutoMigrate creates/maintains it.
  393. // - PostgreSQL: the logs table is a RANGE-partitioned table on created_at,
  394. // created and maintained by the pg_partman extension (initdb script). GORM's
  395. // AutoMigrate must NOT touch it, otherwise it would try to alter the composite
  396. // primary key (id, created_at) - required by the partition key - and create
  397. // redundant per-partition indexes from the model's index tags, breaking
  398. // startup. Here we only verify the table exists and warn if it does not.
  399. //
  400. // No partition creation/maintenance logic lives in application code; pg_partman
  401. // plus pg_cron own that responsibility entirely on the database side.
  402. func ensureLogTable() error {
  403. if !common.UsingPostgreSQL {
  404. return DB.AutoMigrate(&Log{})
  405. }
  406. var exists bool
  407. if err := DB.Raw(`SELECT EXISTS (
  408. SELECT 1 FROM information_schema.tables
  409. WHERE table_schema = current_schema() AND table_name = 'logs'
  410. )`).Scan(&exists).Error; err != nil {
  411. return err
  412. }
  413. if exists {
  414. common.SysLog("logs table is a pg_partman-managed partitioned table, skipping GORM AutoMigrate")
  415. return nil
  416. }
  417. common.SysLog("WARNING: PostgreSQL detected but 'logs' table not found. " +
  418. "Ensure the pg_partman init script created the partitioned logs table before starting new-api.")
  419. return nil
  420. }
  421. func migrateLOGDB() error {
  422. // When LOG_SQL_DSN is empty, LOG_DB == DB and InitLogDB returns early without
  423. // calling this function, so the logs table is already handled by ensureLogTable()
  424. // during migrateDB(). This branch only runs for a dedicated PostgreSQL log
  425. // database: keep pg_partman-managed behavior (skip AutoMigrate) consistent with
  426. // the main database.
  427. if common.LogSqlType == common.DatabaseTypePostgreSQL {
  428. var exists bool
  429. if err := LOG_DB.Raw(`SELECT EXISTS (
  430. SELECT 1 FROM information_schema.tables
  431. WHERE table_schema = current_schema() AND table_name = 'logs'
  432. )`).Scan(&exists).Error; err != nil {
  433. return err
  434. }
  435. if exists {
  436. common.SysLog("logs table is a pg_partman-managed partitioned table, skipping GORM AutoMigrate")
  437. return nil
  438. }
  439. common.SysLog("WARNING: PostgreSQL log database detected but 'logs' table not found. " +
  440. "Ensure the pg_partman init script created the partitioned logs table before starting new-api.")
  441. return nil
  442. }
  443. var err error
  444. if err = LOG_DB.AutoMigrate(&Log{}); err != nil {
  445. return err
  446. }
  447. return nil
  448. }
  449. type sqliteColumnDef struct {
  450. Name string
  451. DDL string
  452. }
  453. func ensureSubscriptionPlanTableSQLite() error {
  454. if !common.UsingSQLite {
  455. return nil
  456. }
  457. tableName := "subscription_plans"
  458. if !DB.Migrator().HasTable(tableName) {
  459. createSQL := `CREATE TABLE ` + "`" + tableName + "`" + ` (
  460. ` + "`id`" + ` integer,
  461. ` + "`title`" + ` varchar(128) NOT NULL,
  462. ` + "`subtitle`" + ` varchar(255) DEFAULT '',
  463. ` + "`price_amount`" + ` decimal(10,6) NOT NULL,
  464. ` + "`currency`" + ` varchar(8) NOT NULL DEFAULT 'USD',
  465. ` + "`duration_unit`" + ` varchar(16) NOT NULL DEFAULT 'month',
  466. ` + "`duration_value`" + ` integer NOT NULL DEFAULT 1,
  467. ` + "`custom_seconds`" + ` bigint NOT NULL DEFAULT 0,
  468. ` + "`enabled`" + ` numeric DEFAULT 1,
  469. ` + "`sort_order`" + ` integer DEFAULT 0,
  470. ` + "`stripe_price_id`" + ` varchar(128) DEFAULT '',
  471. ` + "`creem_product_id`" + ` varchar(128) DEFAULT '',
  472. ` + "`max_purchase_per_user`" + ` integer DEFAULT 0,
  473. ` + "`upgrade_group`" + ` varchar(64) DEFAULT '',
  474. ` + "`total_amount`" + ` bigint NOT NULL DEFAULT 0,
  475. ` + "`quota_reset_period`" + ` varchar(16) DEFAULT 'never',
  476. ` + "`quota_reset_custom_seconds`" + ` bigint DEFAULT 0,
  477. ` + "`created_at`" + ` bigint,
  478. ` + "`updated_at`" + ` bigint,
  479. PRIMARY KEY (` + "`id`" + `)
  480. )`
  481. return DB.Exec(createSQL).Error
  482. }
  483. var cols []struct {
  484. Name string `gorm:"column:name"`
  485. }
  486. if err := DB.Raw("PRAGMA table_info(`" + tableName + "`)").Scan(&cols).Error; err != nil {
  487. return err
  488. }
  489. existing := make(map[string]struct{}, len(cols))
  490. for _, c := range cols {
  491. existing[c.Name] = struct{}{}
  492. }
  493. required := []sqliteColumnDef{
  494. {Name: "title", DDL: "`title` varchar(128) NOT NULL"},
  495. {Name: "subtitle", DDL: "`subtitle` varchar(255) DEFAULT ''"},
  496. {Name: "price_amount", DDL: "`price_amount` decimal(10,6) NOT NULL"},
  497. {Name: "currency", DDL: "`currency` varchar(8) NOT NULL DEFAULT 'USD'"},
  498. {Name: "duration_unit", DDL: "`duration_unit` varchar(16) NOT NULL DEFAULT 'month'"},
  499. {Name: "duration_value", DDL: "`duration_value` integer NOT NULL DEFAULT 1"},
  500. {Name: "custom_seconds", DDL: "`custom_seconds` bigint NOT NULL DEFAULT 0"},
  501. {Name: "enabled", DDL: "`enabled` numeric DEFAULT 1"},
  502. {Name: "sort_order", DDL: "`sort_order` integer DEFAULT 0"},
  503. {Name: "stripe_price_id", DDL: "`stripe_price_id` varchar(128) DEFAULT ''"},
  504. {Name: "creem_product_id", DDL: "`creem_product_id` varchar(128) DEFAULT ''"},
  505. {Name: "max_purchase_per_user", DDL: "`max_purchase_per_user` integer DEFAULT 0"},
  506. {Name: "upgrade_group", DDL: "`upgrade_group` varchar(64) DEFAULT ''"},
  507. {Name: "total_amount", DDL: "`total_amount` bigint NOT NULL DEFAULT 0"},
  508. {Name: "quota_reset_period", DDL: "`quota_reset_period` varchar(16) DEFAULT 'never'"},
  509. {Name: "quota_reset_custom_seconds", DDL: "`quota_reset_custom_seconds` bigint DEFAULT 0"},
  510. {Name: "created_at", DDL: "`created_at` bigint"},
  511. {Name: "updated_at", DDL: "`updated_at` bigint"},
  512. }
  513. for _, col := range required {
  514. if _, ok := existing[col.Name]; ok {
  515. continue
  516. }
  517. if err := DB.Exec("ALTER TABLE `" + tableName + "` ADD COLUMN " + col.DDL).Error; err != nil {
  518. return err
  519. }
  520. }
  521. return nil
  522. }
  523. // migrateSubscriptionPlanPriceAmount migrates price_amount column from float/double to decimal(10,6)
  524. // This is safe to run multiple times - it checks the column type first
  525. func migrateSubscriptionPlanPriceAmount() {
  526. // SQLite doesn't support ALTER COLUMN, and its type affinity handles this automatically
  527. // Skip early to avoid GORM parsing the existing table DDL which may cause issues
  528. if common.UsingSQLite {
  529. return
  530. }
  531. tableName := "subscription_plans"
  532. columnName := "price_amount"
  533. // Check if table exists first
  534. if !DB.Migrator().HasTable(tableName) {
  535. return
  536. }
  537. // Check if column exists
  538. if !DB.Migrator().HasColumn(&SubscriptionPlan{}, columnName) {
  539. return
  540. }
  541. var alterSQL string
  542. if common.UsingPostgreSQL {
  543. // PostgreSQL: Check if already decimal/numeric
  544. var dataType string
  545. DB.Raw(`SELECT data_type FROM information_schema.columns
  546. WHERE table_name = ? AND column_name = ?`, tableName, columnName).Scan(&dataType)
  547. if dataType == "numeric" {
  548. return // Already decimal/numeric
  549. }
  550. alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE decimal(10,6) USING %s::decimal(10,6)`,
  551. tableName, columnName, columnName)
  552. } else if common.UsingMySQL {
  553. // MySQL: Check if already decimal
  554. var columnType string
  555. DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
  556. WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
  557. tableName, columnName).Scan(&columnType)
  558. if strings.HasPrefix(strings.ToLower(columnType), "decimal") {
  559. return // Already decimal
  560. }
  561. alterSQL = fmt.Sprintf("ALTER TABLE %s MODIFY COLUMN %s decimal(10,6) NOT NULL DEFAULT 0",
  562. tableName, columnName)
  563. } else {
  564. return
  565. }
  566. if alterSQL != "" {
  567. if err := DB.Exec(alterSQL).Error; err != nil {
  568. common.SysLog(fmt.Sprintf("Warning: failed to migrate %s.%s to decimal: %v", tableName, columnName, err))
  569. } else {
  570. common.SysLog(fmt.Sprintf("Successfully migrated %s.%s to decimal(10,6)", tableName, columnName))
  571. }
  572. }
  573. }
  574. // migrateTokenDropBoundChannelId 删除 tokens 表中的 bound_channel_id 列(已废弃的字段)
  575. func migrateTokenDropBoundChannelId() {
  576. // SQLite 不支持 DROP COLUMN,跳过迁移
  577. if common.UsingSQLite {
  578. return
  579. }
  580. tableName := "tokens"
  581. columnName := "bound_channel_id"
  582. // 检查表是否存在
  583. if !DB.Migrator().HasTable(tableName) {
  584. return
  585. }
  586. // 检查列是否存在
  587. if !DB.Migrator().HasColumn(&Token{}, columnName) {
  588. return
  589. }
  590. // MySQL 和 PostgreSQL 都支持此语法
  591. dropSQL := fmt.Sprintf("ALTER TABLE %s DROP COLUMN %s", tableName, columnName)
  592. if err := DB.Exec(dropSQL).Error; err != nil {
  593. common.SysLog(fmt.Sprintf("Warning: failed to drop %s.%s: %v", tableName, columnName, err))
  594. } else {
  595. common.SysLog(fmt.Sprintf("Successfully dropped deprecated column %s.%s", tableName, columnName))
  596. }
  597. }
  598. func closeDB(db *gorm.DB) error {
  599. sqlDB, err := db.DB()
  600. if err != nil {
  601. return err
  602. }
  603. err = sqlDB.Close()
  604. return err
  605. }
  606. func CloseDB() error {
  607. if LOG_DB != DB {
  608. err := closeDB(LOG_DB)
  609. if err != nil {
  610. return err
  611. }
  612. }
  613. return closeDB(DB)
  614. }
  615. // checkMySQLChineseSupport ensures the MySQL connection and current schema
  616. // default charset/collation can store Chinese characters. It allows common
  617. // Chinese-capable charsets (utf8mb4, utf8, gbk, big5, gb18030) and panics otherwise.
  618. func checkMySQLChineseSupport(db *gorm.DB) error {
  619. // 仅检测:当前库默认字符集/排序规则 + 各表的排序规则(隐含字符集)
  620. // Read current schema defaults
  621. var schemaCharset, schemaCollation string
  622. err := db.Raw("SELECT DEFAULT_CHARACTER_SET_NAME, DEFAULT_COLLATION_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = DATABASE()").Row().Scan(&schemaCharset, &schemaCollation)
  623. if err != nil {
  624. return fmt.Errorf("读取当前库默认字符集/排序规则失败 / Failed to read schema default charset/collation: %v", err)
  625. }
  626. toLower := func(s string) string { return strings.ToLower(s) }
  627. // Allowed charsets that can store Chinese text
  628. allowedCharsets := map[string]string{
  629. "utf8mb4": "utf8mb4_",
  630. "utf8": "utf8_",
  631. "gbk": "gbk_",
  632. "big5": "big5_",
  633. "gb18030": "gb18030_",
  634. }
  635. isChineseCapable := func(cs, cl string) bool {
  636. csLower := toLower(cs)
  637. clLower := toLower(cl)
  638. if prefix, ok := allowedCharsets[csLower]; ok {
  639. if clLower == "" {
  640. return true
  641. }
  642. return strings.HasPrefix(clLower, prefix)
  643. }
  644. // 如果仅提供了排序规则,尝试按排序规则前缀判断
  645. for _, prefix := range allowedCharsets {
  646. if strings.HasPrefix(clLower, prefix) {
  647. return true
  648. }
  649. }
  650. return false
  651. }
  652. // 1) 当前库默认值必须支持中文
  653. if !isChineseCapable(schemaCharset, schemaCollation) {
  654. return fmt.Errorf("当前库默认字符集/排序规则不支持中文:schema(%s/%s)。请将库设置为 utf8mb4/utf8/gbk/big5/gb18030 / Schema default charset/collation is not Chinese-capable: schema(%s/%s). Please set to utf8mb4/utf8/gbk/big5/gb18030",
  655. schemaCharset, schemaCollation, schemaCharset, schemaCollation)
  656. }
  657. // 2) 所有物理表的排序规则(隐含字符集)必须支持中文
  658. type tableInfo struct {
  659. Name string
  660. Collation *string
  661. }
  662. var tables []tableInfo
  663. if err := db.Raw("SELECT TABLE_NAME, TABLE_COLLATION FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() AND TABLE_TYPE = 'BASE TABLE'").Scan(&tables).Error; err != nil {
  664. return fmt.Errorf("读取表排序规则失败 / Failed to read table collations: %v", err)
  665. }
  666. var badTables []string
  667. for _, t := range tables {
  668. // NULL 或空表示继承库默认设置,已在上面校验库默认,视为通过
  669. if t.Collation == nil || *t.Collation == "" {
  670. continue
  671. }
  672. cl := *t.Collation
  673. // 仅凭排序规则判断是否中文可用
  674. ok := false
  675. lower := strings.ToLower(cl)
  676. for _, prefix := range allowedCharsets {
  677. if strings.HasPrefix(lower, prefix) {
  678. ok = true
  679. break
  680. }
  681. }
  682. if !ok {
  683. badTables = append(badTables, fmt.Sprintf("%s(%s)", t.Name, cl))
  684. }
  685. }
  686. if len(badTables) > 0 {
  687. // 限制输出数量以避免日志过长
  688. maxShow := 20
  689. shown := badTables
  690. if len(shown) > maxShow {
  691. shown = shown[:maxShow]
  692. }
  693. return fmt.Errorf(
  694. "存在不支持中文的表,请修复其排序规则/字符集。示例(最多展示 %d 项):%v / Found tables not Chinese-capable. Please fix their collation/charset. Examples (showing up to %d): %v",
  695. maxShow, shown, maxShow, shown,
  696. )
  697. }
  698. return nil
  699. }
  700. var (
  701. lastPingTime time.Time
  702. pingMutex sync.Mutex
  703. )
  704. func PingDB() error {
  705. pingMutex.Lock()
  706. defer pingMutex.Unlock()
  707. if time.Since(lastPingTime) < time.Second*10 {
  708. return nil
  709. }
  710. sqlDB, err := DB.DB()
  711. if err != nil {
  712. log.Printf("Error getting sql.DB from GORM: %v", err)
  713. return err
  714. }
  715. err = sqlDB.Ping()
  716. if err != nil {
  717. log.Printf("Error pinging DB: %v", err)
  718. return err
  719. }
  720. lastPingTime = time.Now()
  721. common.SysLog("Database pinged successfully")
  722. return nil
  723. }
  724. func migrateChannelPublicName() {
  725. result := DB.Model(&Channel{}).
  726. Where("public_name = '' OR public_name IS NULL").
  727. Update("public_name", gorm.Expr("name"))
  728. if result.Error != nil {
  729. common.SysError("[Migration] migrateChannelPublicName failed: " + result.Error.Error())
  730. } else if result.RowsAffected > 0 {
  731. common.SysLog(fmt.Sprintf("[Migration] migrateChannelPublicName: backfilled %d channels", result.RowsAffected))
  732. }
  733. }