fix(model): migrate legacy token key constraints

This commit is contained in:
CaIon
2026-08-31 11:51:59 +08:00
parent 8c8c4153d4
commit 27ff6a8767
3 changed files with 515 additions and 0 deletions
+3
View File
@@ -315,6 +315,9 @@ func is64BitIntegerType(dbType common.DatabaseType, dataType string) bool {
} }
func migrateDB() error { func migrateDB() error {
if err := migrateTokenKeyUniqueness(DB); err != nil {
return err
}
if err := migratePrefillGroupUniqueness(DB); err != nil { if err := migratePrefillGroupUniqueness(DB); err != nil {
return err return err
} }
+217
View File
@@ -0,0 +1,217 @@
package model
import (
"fmt"
"strings"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
const (
tokenKeyIndex = "idx_tokens_key"
postgresTokenKeyConstraint = "tokens_key_key"
gormTokenKeyConstraint = "uni_tokens_key"
)
type tokenKeyUniqueConstraint struct {
Name string `gorm:"column:constraint_name"`
Definition string `gorm:"column:constraint_definition"`
Deferrable bool `gorm:"column:is_deferrable"`
Validated bool `gorm:"column:is_validated"`
}
type tokenKeyIndexState struct {
exists bool
definitionValid bool
standaloneValid bool
}
func inspectTokenKeyUniqueConstraints(db *gorm.DB, tableName string) ([]tokenKeyUniqueConstraint, error) {
var constraints []tokenKeyUniqueConstraint
if err := db.Raw(`
SELECT constraint_meta.conname AS constraint_name,
pg_get_constraintdef(constraint_meta.oid) AS constraint_definition,
constraint_meta.condeferrable AS is_deferrable,
constraint_meta.convalidated AS is_validated
FROM pg_catalog.pg_constraint AS constraint_meta
WHERE constraint_meta.conrelid = to_regclass(?)
AND constraint_meta.contype = 'u'
AND cardinality(constraint_meta.conkey) = 1
AND EXISTS (
SELECT 1
FROM pg_catalog.pg_attribute AS attribute_meta
WHERE attribute_meta.attrelid = constraint_meta.conrelid
AND attribute_meta.attnum = constraint_meta.conkey[1]
AND attribute_meta.attname = ?
)
ORDER BY constraint_meta.conname`, tableName, "key").Scan(&constraints).Error; err != nil {
return nil, fmt.Errorf("inspect token key unique constraints: %w", err)
}
return constraints, nil
}
func validateTokenKeyUniqueConstraints(constraints []tokenKeyUniqueConstraint) error {
for _, constraint := range constraints {
switch constraint.Name {
case tokenKeyIndex, postgresTokenKeyConstraint, gormTokenKeyConstraint:
default:
return fmt.Errorf(
"tokens.key has unsupported unique constraint %q with definition %q",
constraint.Name,
constraint.Definition,
)
}
if constraint.Deferrable || !constraint.Validated || strings.Contains(strings.ToUpper(constraint.Definition), "NULLS NOT DISTINCT") {
return fmt.Errorf(
"tokens.key unique constraint %q has unsupported definition %q",
constraint.Name,
constraint.Definition,
)
}
}
return nil
}
func inspectTokenKeyIndex(db *gorm.DB, tableName string) (tokenKeyIndexState, error) {
var state struct {
Exists bool `gorm:"column:index_exists"`
DefinitionValid bool `gorm:"column:definition_valid"`
StandaloneValid bool `gorm:"column:standalone_valid"`
}
if err := db.Raw(`
SELECT count(*) > 0 AS index_exists,
COALESCE(bool_or(
index_meta.indisunique
AND index_meta.indisvalid
AND index_meta.indisready
AND NOT index_meta.indisprimary
AND index_meta.indpred IS NULL
AND index_meta.indexprs IS NULL
AND index_meta.indnatts = 1
AND attribute_meta.attname = ?
), false) AS definition_valid,
COALESCE(bool_or(
index_meta.indisunique
AND index_meta.indisvalid
AND index_meta.indisready
AND NOT index_meta.indisprimary
AND index_meta.indpred IS NULL
AND index_meta.indexprs IS NULL
AND index_meta.indnatts = 1
AND attribute_meta.attname = ?
AND NOT EXISTS (
SELECT 1
FROM pg_catalog.pg_constraint AS constraint_meta
WHERE constraint_meta.conindid = index_meta.indexrelid
)
), false) AS standalone_valid
FROM pg_catalog.pg_index AS index_meta
JOIN pg_catalog.pg_class AS index_class
ON index_class.oid = index_meta.indexrelid
LEFT JOIN pg_catalog.pg_attribute AS attribute_meta
ON attribute_meta.attrelid = index_meta.indrelid
AND attribute_meta.attnum = index_meta.indkey[0]
WHERE index_meta.indrelid = to_regclass(?)
AND index_class.relname = ?`, "key", "key", tableName, tokenKeyIndex).Scan(&state).Error; err != nil {
return tokenKeyIndexState{}, fmt.Errorf("inspect token key unique index: %w", err)
}
return tokenKeyIndexState{
exists: state.Exists,
definitionValid: state.DefinitionValid,
standaloneValid: state.StandaloneValid,
}, nil
}
// migrateTokenKeyUniqueness converts known PostgreSQL UNIQUE constraints left
// on tokens.key into the standalone uniqueIndex represented by the current
// model. Unknown constraint names are reported without modifying the schema.
func migrateTokenKeyUniqueness(db *gorm.DB) error {
if db == nil {
return fmt.Errorf("migrate token key uniqueness: database is nil")
}
if db.Dialector.Name() != "postgres" {
return nil
}
statement := &gorm.Statement{DB: db}
if err := statement.Parse(&Token{}); err != nil {
return fmt.Errorf("parse token schema: %w", err)
}
tableName := statement.Schema.Table
constraints, err := inspectTokenKeyUniqueConstraints(db, tableName)
if err != nil {
return err
}
if len(constraints) == 0 {
return nil
}
if err := validateTokenKeyUniqueConstraints(constraints); err != nil {
return err
}
return db.Transaction(func(tx *gorm.DB) error {
migrator := tx.Migrator()
if !migrator.HasTable(&Token{}) {
return nil
}
if err := tx.Exec(
"LOCK TABLE ? IN ACCESS EXCLUSIVE MODE",
clause.Table{Name: tableName},
).Error; err != nil {
return fmt.Errorf("lock tokens for key uniqueness migration: %w", err)
}
constraints, err := inspectTokenKeyUniqueConstraints(tx, tableName)
if err != nil {
return err
}
if len(constraints) == 0 {
return nil
}
if err := validateTokenKeyUniqueConstraints(constraints); err != nil {
return err
}
targetIndex, err := inspectTokenKeyIndex(tx, tableName)
if err != nil {
return err
}
if targetIndex.exists && !targetIndex.definitionValid {
return fmt.Errorf("token key index %q has an unexpected definition", tokenKeyIndex)
}
for _, constraint := range constraints {
if err := migrator.DropConstraint(&Token{}, constraint.Name); err != nil {
return fmt.Errorf("drop token key unique constraint %q: %w", constraint.Name, err)
}
}
targetIndex, err = inspectTokenKeyIndex(tx, tableName)
if err != nil {
return err
}
if !targetIndex.exists {
if err := migrator.CreateIndex(&Token{}, tokenKeyIndex); err != nil {
return fmt.Errorf("create token key unique index: %w", err)
}
targetIndex, err = inspectTokenKeyIndex(tx, tableName)
if err != nil {
return err
}
}
if !targetIndex.standaloneValid {
return fmt.Errorf("token key index %q has an unexpected definition", tokenKeyIndex)
}
remainingConstraints, err := inspectTokenKeyUniqueConstraints(tx, tableName)
if err != nil {
return err
}
if len(remainingConstraints) != 0 {
return fmt.Errorf("tokens.key still has unique constraints after migration")
}
return nil
})
}
+295
View File
@@ -0,0 +1,295 @@
package model
import (
"fmt"
"os"
"strings"
"testing"
"time"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
func requireTokenConstraintExists(t *testing.T, db *gorm.DB, constraintName string) {
t.Helper()
var count int64
require.NoError(t, db.Raw(`
SELECT count(*)
FROM pg_catalog.pg_constraint
WHERE conrelid = to_regclass(?)
AND conname = ?`, "tokens", constraintName).Scan(&count).Error)
require.EqualValues(t, 1, count)
}
func requireTokenIndexExists(t *testing.T, db *gorm.DB, indexName string) {
t.Helper()
var count int64
require.NoError(t, db.Raw(`
SELECT count(*)
FROM pg_catalog.pg_index AS index_meta
JOIN pg_catalog.pg_class AS index_class
ON index_class.oid = index_meta.indexrelid
WHERE index_meta.indrelid = to_regclass(?)
AND index_class.relname = ?`, "tokens", indexName).Scan(&count).Error)
require.EqualValues(t, 1, count)
}
func testTokenKeyMigrationNonPostgreSQL(t *testing.T, db *gorm.DB) {
t.Helper()
tableName := fmt.Sprintf("token_migration_%d", time.Now().UnixNano())
t.Cleanup(func() { _ = db.Migrator().DropTable(tableName) })
tableDB := db.Table(tableName)
require.NoError(t, tableDB.AutoMigrate(&Token{}))
require.NoError(t, tableDB.Create(&Token{UserId: 1, Key: "preserved-key"}).Error)
for range 2 {
require.NoError(t, migrateTokenKeyUniqueness(db))
require.NoError(t, tableDB.AutoMigrate(&Token{}))
}
var preserved Token
require.NoError(t, tableDB.Where(&Token{Key: "preserved-key"}).First(&preserved).Error)
assert.Equal(t, 1, preserved.UserId)
expectedIndex := db.NamingStrategy.IndexName(tableName, "key")
assert.True(t, db.Migrator().HasIndex(tableName, expectedIndex))
}
func TestMigrateTokenKeyUniquenessSQLite(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
testTokenKeyMigrationNonPostgreSQL(t, db)
}
func TestMigrateTokenKeyUniquenessMySQL(t *testing.T) {
dsn := strings.TrimSpace(os.Getenv("TEST_MYSQL_DSN"))
if dsn == "" {
t.Skip("TEST_MYSQL_DSN is not configured")
}
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
testTokenKeyMigrationNonPostgreSQL(t, db)
}
func TestMigrateTokenKeyUniquenessPostgreSQL(t *testing.T) {
dsn := strings.TrimSpace(os.Getenv("TEST_POSTGRES_DSN"))
if dsn == "" {
t.Skip("TEST_POSTGRES_DSN is not configured")
}
db, err := gorm.Open(postgres.New(postgres.Config{
DSN: dsn,
PreferSimpleProtocol: true,
}), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
tests := []struct {
name string
prepareOld func(*testing.T, *gorm.DB)
expectedError string
preservedConstraints []string
preservedIndexes []string
}{
{name: "fresh"},
{
name: "legacy_idx_constraint",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Migrator().DropIndex(&Token{}, tokenKeyIndex))
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: tokenKeyIndex},
clause.Column{Name: "key"},
).Error)
},
},
{
name: "gorm_generated_constraint",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: gormTokenKeyConstraint},
clause.Column{Name: "key"},
).Error)
},
},
{
name: "postgres_default_constraint_without_target_index",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Migrator().DropIndex(&Token{}, tokenKeyIndex))
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: postgresTokenKeyConstraint},
clause.Column{Name: "key"},
).Error)
},
},
{
name: "non_conflicting_uniqueness_is_preserved",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: postgresTokenKeyConstraint},
clause.Column{Name: "key"},
).Error)
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?, ?)",
clause.Table{Name: "tokens"},
clause.Column{Name: "keep_tokens_key_user_id"},
clause.Column{Name: "key"},
clause.Column{Name: "user_id"},
).Error)
require.NoError(t, tx.Exec(
"CREATE UNIQUE INDEX ? ON ? (?) WHERE user_id > 0",
clause.Column{Name: "keep_tokens_partial_key"},
clause.Table{Name: "tokens"},
clause.Column{Name: "key"},
).Error)
},
preservedConstraints: []string{"keep_tokens_key_user_id"},
preservedIndexes: []string{"keep_tokens_partial_key"},
},
{
name: "arbitrary_constraint_is_rejected",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: "keep_tokens_key_unique"},
clause.Column{Name: "key"},
).Error)
},
expectedError: "unsupported unique constraint",
preservedConstraints: []string{"keep_tokens_key_unique"},
},
{
name: "deferrable_constraint_is_rejected",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Migrator().DropIndex(&Token{}, tokenKeyIndex))
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?) DEFERRABLE INITIALLY DEFERRED",
clause.Table{Name: "tokens"},
clause.Column{Name: postgresTokenKeyConstraint},
clause.Column{Name: "key"},
).Error)
},
expectedError: "unsupported definition",
preservedConstraints: []string{postgresTokenKeyConstraint},
},
{
name: "invalid_target_index_is_rejected",
prepareOld: func(t *testing.T, tx *gorm.DB) {
t.Helper()
require.NoError(t, tx.Migrator().DropIndex(&Token{}, tokenKeyIndex))
require.NoError(t, tx.Exec(
"CREATE INDEX ? ON ? (?)",
clause.Column{Name: tokenKeyIndex},
clause.Table{Name: "tokens"},
clause.Column{Name: "key"},
).Error)
require.NoError(t, tx.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
clause.Table{Name: "tokens"},
clause.Column{Name: postgresTokenKeyConstraint},
clause.Column{Name: "key"},
).Error)
},
expectedError: "unexpected definition",
preservedConstraints: []string{postgresTokenKeyConstraint},
preservedIndexes: []string{tokenKeyIndex},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
tx := db.Begin()
require.NoError(t, tx.Error)
t.Cleanup(func() { _ = tx.Rollback().Error })
schemaName := fmt.Sprintf("token_migration_%d", time.Now().UnixNano())
require.NoError(t, tx.Exec(
"CREATE SCHEMA ?",
clause.Table{Name: schemaName},
).Error)
require.NoError(t, tx.Exec(
"SET LOCAL search_path TO ?",
clause.Table{Name: schemaName},
).Error)
require.NoError(t, migrateTokenKeyUniqueness(tx))
require.NoError(t, tx.AutoMigrate(&Token{}))
original := Token{UserId: 1, Key: "preserved-key", Name: "preserve me"}
require.NoError(t, tx.Create(&original).Error)
if test.prepareOld != nil {
test.prepareOld(t, tx)
}
if test.expectedError != "" {
err := migrateTokenKeyUniqueness(tx)
require.Error(t, err)
assert.Contains(t, err.Error(), test.expectedError)
for _, constraintName := range test.preservedConstraints {
requireTokenConstraintExists(t, tx, constraintName)
}
for _, indexName := range test.preservedIndexes {
requireTokenIndexExists(t, tx, indexName)
}
return
}
for range 2 {
require.NoError(t, migrateTokenKeyUniqueness(tx))
require.NoError(t, tx.AutoMigrate(&Token{}))
}
var preserved Token
require.NoError(t, tx.First(&preserved, original.Id).Error)
assert.Equal(t, original.Key, preserved.Key)
assert.Equal(t, original.Name, preserved.Name)
constraints, err := inspectTokenKeyUniqueConstraints(tx, "tokens")
require.NoError(t, err)
assert.Empty(t, constraints)
targetIndex, err := inspectTokenKeyIndex(tx, "tokens")
require.NoError(t, err)
assert.True(t, targetIndex.standaloneValid)
for _, constraintName := range test.preservedConstraints {
requireTokenConstraintExists(t, tx, constraintName)
}
for _, indexName := range test.preservedIndexes {
requireTokenIndexExists(t, tx, indexName)
}
duplicateError := tx.Transaction(func(duplicateTx *gorm.DB) error {
return duplicateTx.Create(&Token{UserId: 2, Key: original.Key}).Error
})
require.Error(t, duplicateError)
var totalRows int64
require.NoError(t, tx.Model(&Token{}).Count(&totalRows).Error)
assert.EqualValues(t, 1, totalRows)
})
}
}