model/option_primary_key_migration.go
Rama main · réplica de lectura del repositorio original.
package model
import (
"context"
"fmt"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
)
const (
optionPrimaryKeyTmpTable = "options_pk_tmp"
optionPrimaryKeyLockName = "new_api_options_pk"
optionPrimaryKeyLockID = 75820193
optionLegacyTablePrefix = "options_legacy_"
)
func migrateOptionPrimaryKey(db *gorm.DB) error {
if db == nil {
return fmt.Errorf("migrate options primary key: database is nil")
}
if !db.Migrator().HasTable(&Option{}) {
return nil
}
unique, err := optionsKeyIsUnique(db)
if err != nil {
return err
}
if unique {
return nil
}
return withOptionPrimaryKeyLock(db, func(locked *gorm.DB) error {
unique, err := optionsKeyIsUnique(locked)
if err != nil {
return err
}
if unique {
return nil
}
return repairOptionPrimaryKey(locked)
})
}
func optionsKeyIsUnique(db *gorm.DB) (bool, error) {
indexes, err := db.Migrator().GetIndexes(&Option{})
if err != nil {
return false, fmt.Errorf("inspect options indexes: %w", err)
}
for _, index := range indexes {
if len(index.Columns()) != 1 || index.Columns()[0] != "key" {
continue
}
if unique, ok := index.Unique(); ok && unique {
return true, nil
}
if primary, ok := index.PrimaryKey(); ok && primary {
return true, nil
}
}
if db.Dialector.Name() != "postgres" {
return false, nil
}
// PostgreSQL GetIndexes omits constraint-backed primary/unique keys.
var count int64
if err := db.Raw(`
SELECT count(*)
FROM pg_catalog.pg_constraint AS constraint_meta
JOIN pg_catalog.pg_attribute AS attribute_meta
ON attribute_meta.attrelid = constraint_meta.conrelid
AND attribute_meta.attnum = constraint_meta.conkey[1]
WHERE constraint_meta.conrelid = to_regclass('options')
AND constraint_meta.contype IN ('p', 'u')
AND cardinality(constraint_meta.conkey) = 1
AND attribute_meta.attname = 'key'`).Scan(&count).Error; err != nil {
return false, fmt.Errorf("inspect options constraints: %w", err)
}
return count > 0, nil
}
func withOptionPrimaryKeyLock(db *gorm.DB, fn func(*gorm.DB) error) error {
switch db.Dialector.Name() {
case "mysql":
sqlDB, err := db.DB()
if err != nil {
return fmt.Errorf("lock options table: %w", err)
}
ctx := context.Background()
conn, err := sqlDB.Conn(ctx)
if err != nil {
return fmt.Errorf("lock options table: %w", err)
}
defer conn.Close()
var acquired int
if err := conn.QueryRowContext(ctx, "SELECT GET_LOCK(?, 60)", optionPrimaryKeyLockName).Scan(&acquired); err != nil {
return fmt.Errorf("lock options table: %w", err)
}
if acquired != 1 {
return fmt.Errorf("lock options table: timeout")
}
defer conn.ExecContext(ctx, "SELECT RELEASE_LOCK(?)", optionPrimaryKeyLockName)
return fn(db)
case "postgres":
return db.Transaction(func(tx *gorm.DB) error {
if err := tx.Exec("SELECT pg_advisory_xact_lock(?)", optionPrimaryKeyLockID).Error; err != nil {
return fmt.Errorf("lock options table: %w", err)
}
return fn(tx)
})
default:
return db.Transaction(func(tx *gorm.DB) error {
return fn(tx)
})
}
}
func repairOptionPrimaryKey(db *gorm.DB) error {
var rows []Option
if err := db.Find(&rows).Error; err != nil {
return fmt.Errorf("read options rows: %w", err)
}
deduped, skippedEmpty := dedupeOptionRows(rows)
if db.Migrator().HasTable(optionPrimaryKeyTmpTable) {
if err := db.Migrator().DropTable(optionPrimaryKeyTmpTable); err != nil {
return fmt.Errorf("drop leftover %s: %w", optionPrimaryKeyTmpTable, err)
}
}
if err := db.Table(optionPrimaryKeyTmpTable).Migrator().CreateTable(&Option{}); err != nil {
return fmt.Errorf("create %s: %w", optionPrimaryKeyTmpTable, err)
}
tmpCreated := true
defer func() {
if tmpCreated && db.Migrator().HasTable(optionPrimaryKeyTmpTable) {
_ = db.Migrator().DropTable(optionPrimaryKeyTmpTable)
}
}()
if len(deduped) > 0 {
if err := db.Table(optionPrimaryKeyTmpTable).CreateInBatches(deduped, 100).Error; err != nil {
return fmt.Errorf("insert rebuilt options: %w", err)
}
}
var written int64
if err := db.Table(optionPrimaryKeyTmpTable).Count(&written).Error; err != nil {
return fmt.Errorf("count rebuilt options: %w", err)
}
if int(written) != len(deduped) {
return fmt.Errorf("options rebuild wrote %d rows, want %d", written, len(deduped))
}
backup := fmt.Sprintf("%s%d", optionLegacyTablePrefix, time.Now().UnixNano())
if err := swapOptionTables(db, optionPrimaryKeyTmpTable, backup); err != nil {
return err
}
tmpCreated = false
unique, err := optionsKeyIsUnique(db)
if err != nil {
return err
}
if !unique {
return fmt.Errorf("options table still has no unique key after rebuild")
}
common.SysLog(common.LogText("rebuilt options table with primary key from %d rows into %d keys; previous rows kept in %s", len(rows)-skippedEmpty, len(deduped), backup))
return nil
}
func dedupeOptionRows(rows []Option) ([]Option, int) {
seen := make(map[string]int)
deduped := make([]Option, 0, len(rows))
skippedEmpty := 0
for _, row := range rows {
if row.Key == "" {
skippedEmpty++
continue
}
if i, ok := seen[row.Key]; ok {
if deduped[i].Value != row.Value {
common.SysError(common.LogText("options key %q had conflicting values; keeping the last row", row.Key))
}
deduped[i].Value = row.Value
continue
}
seen[row.Key] = len(deduped)
deduped = append(deduped, Option{Key: row.Key, Value: row.Value})
}
if skippedEmpty > 0 {
common.SysError(common.LogText("skipped %d options rows with empty keys while rebuilding the primary key", skippedEmpty))
}
return deduped, skippedEmpty
}
func swapOptionTables(db *gorm.DB, tmp, backup string) error {
if !optionSafeIdent(tmp) || !optionSafeIdent(backup) {
return fmt.Errorf("unsafe options table name")
}
if db.Dialector.Name() == "mysql" {
sql := "RENAME TABLE `options` TO `" + backup + "`, `" + tmp + "` TO `options`"
if err := db.Exec(sql).Error; err != nil {
return fmt.Errorf("rename options tables: %w", err)
}
return nil
}
if err := db.Migrator().RenameTable("options", backup); err != nil {
return fmt.Errorf("rename options to %s: %w", backup, err)
}
if err := db.Migrator().RenameTable(tmp, "options"); err != nil {
return fmt.Errorf("rename %s to options: %w", tmp, err)
}
return nil
}
func optionSafeIdent(name string) bool {
if name == "" {
return false
}
return !strings.ContainsFunc(name, func(r rune) bool {
return r != '_' && (r < '0' || r > '9') && (r < 'a' || r > 'z')
})
}