diff --git a/model/main.go b/model/main.go index fa61ea94f0..3a70622b21 100644 --- a/model/main.go +++ b/model/main.go @@ -315,6 +315,9 @@ func is64BitIntegerType(dbType common.DatabaseType, dataType string) bool { } func migrateDB() error { + if err := migrateTokenKeyUniqueness(DB); err != nil { + return err + } if err := migratePrefillGroupUniqueness(DB); err != nil { return err } diff --git a/model/token_migration.go b/model/token_migration.go new file mode 100644 index 0000000000..c289652ef3 --- /dev/null +++ b/model/token_migration.go @@ -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 + }) +} diff --git a/model/token_migration_test.go b/model/token_migration_test.go new file mode 100644 index 0000000000..f777dec8d0 --- /dev/null +++ b/model/token_migration_test.go @@ -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) + }) + } +}