diff --git a/model/main.go b/model/main.go index 371702c..64b87d3 100644 --- a/model/main.go +++ b/model/main.go @@ -286,6 +286,9 @@ func migrateDB() error { &QuotaSyncLog{}, &EmailQuotaRule{}, &UserModelRateLimit{}, + &UserMigrationBatch{}, + &UserMigrationItem{}, + &MigrationQuotaGrant{}, ) if err != nil { return err @@ -350,6 +353,9 @@ func migrateDBFast() error { {&PendingSyncRecord{}, "PendingSyncRecord"}, {&QuotaSyncLog{}, "QuotaSyncLog"}, {&EmailQuotaRule{}, "EmailQuotaRule"}, + {&UserMigrationBatch{}, "UserMigrationBatch"}, + {&UserMigrationItem{}, "UserMigrationItem"}, + {&MigrationQuotaGrant{}, "MigrationQuotaGrant"}, } // 动态计算migration数量,确保errChan缓冲区足够大 errChan := make(chan error, len(migrations)) diff --git a/model/user_migration.go b/model/user_migration.go new file mode 100644 index 0000000..de70222 --- /dev/null +++ b/model/user_migration.go @@ -0,0 +1,239 @@ +package model + +import ( + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const ( + UserMigrationBatchStatusDraft = "draft" + UserMigrationBatchStatusScanned = "scanned" + UserMigrationBatchStatusReady = "ready" + UserMigrationBatchStatusRunning = "running" + UserMigrationBatchStatusCompleted = "completed" + UserMigrationBatchStatusFailed = "failed" + UserMigrationBatchStatusCancelled = "cancelled" + + UserMigrationItemStatusPendingScan = "pending_scan" + UserMigrationItemStatusConflict = "conflict" + UserMigrationItemStatusReady = "ready" + UserMigrationItemStatusMigrating = "migrating" + UserMigrationItemStatusMigrated = "migrated" + UserMigrationItemStatusSkipped = "skipped" + UserMigrationItemStatusFailed = "failed" + + UserMigrationStrategyCreateNew = "create_new" + UserMigrationStrategyMergeExisting = "merge_into_existing" + UserMigrationStrategySkip = "skip" +) + +type UserMigrationBatch struct { + Id int `json:"id" gorm:"primaryKey"` + Name string `json:"name" gorm:"type:varchar(128);not null"` + SourceRegion string `json:"source_region" gorm:"type:varchar(32);not null"` + TargetRegion string `json:"target_region" gorm:"type:varchar(32);not null"` + Status string `json:"status" gorm:"type:varchar(32);not null;default:'draft';index"` + TotalUsers int `json:"total_users" gorm:"default:0"` + ConflictUsers int `json:"conflict_users" gorm:"default:0"` + MigratedUsers int `json:"migrated_users" gorm:"default:0"` + FailedUsers int `json:"failed_users" gorm:"default:0"` + OperatorId int `json:"operator_id" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"bigint"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +type UserMigrationItem struct { + Id int `json:"id" gorm:"primaryKey"` + BatchId int `json:"batch_id" gorm:"index:idx_batch_source_user,unique"` + SourceUserId int `json:"source_user_id" gorm:"index:idx_batch_source_user,unique"` + TargetUserId int `json:"target_user_id" gorm:"default:0;index"` + SourceUsername string `json:"source_username" gorm:"type:varchar(128);not null"` + SourceEmail string `json:"source_email" gorm:"type:varchar(255);default:''"` + SourceQuota int `json:"source_quota" gorm:"default:0"` + ConflictFlags string `json:"conflict_flags" gorm:"type:text"` + ResolutionStrategy string `json:"resolution_strategy" gorm:"type:varchar(32);default:''"` + Status string `json:"status" gorm:"type:varchar(32);not null;default:'pending_scan';index"` + ErrorMessage string `json:"error_message" gorm:"type:text"` + Snapshot string `json:"snapshot" gorm:"type:text"` + MigrationToken string `json:"migration_token" gorm:"type:varchar(64);default:'';index"` + MigratedAt int64 `json:"migrated_at" gorm:"bigint;default:0"` + CreatedAt int64 `json:"created_at" gorm:"bigint"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +type MigrationQuotaGrant struct { + Id int `json:"id" gorm:"primaryKey"` + MigrationToken string `json:"migration_token" gorm:"type:varchar(64);not null;uniqueIndex"` + CreatedAt int64 `json:"created_at" gorm:"bigint"` +} + +func (b *UserMigrationBatch) BeforeCreate(tx *gorm.DB) error { + now := common.GetTimestamp() + if b.Status == "" { + b.Status = UserMigrationBatchStatusDraft + } + b.CreatedAt = now + b.UpdatedAt = now + return nil +} + +func (b *UserMigrationBatch) BeforeUpdate(tx *gorm.DB) error { + b.UpdatedAt = common.GetTimestamp() + return nil +} + +func (i *UserMigrationItem) BeforeCreate(tx *gorm.DB) error { + now := common.GetTimestamp() + if i.Status == "" { + i.Status = UserMigrationItemStatusPendingScan + } + i.CreatedAt = now + i.UpdatedAt = now + return nil +} + +func (i *UserMigrationItem) BeforeUpdate(tx *gorm.DB) error { + i.UpdatedAt = common.GetTimestamp() + return nil +} + +func (g *MigrationQuotaGrant) BeforeCreate(tx *gorm.DB) error { + g.CreatedAt = common.GetTimestamp() + return nil +} + +func CreateUserMigrationBatch(batch *UserMigrationBatch) error { + return DB.Create(batch).Error +} + +func UpsertUserMigrationItem(item *UserMigrationItem) error { + var existing UserMigrationItem + result := DB.Where("batch_id = ? AND source_user_id = ?", item.BatchId, item.SourceUserId).Limit(1).Find(&existing) + if result.Error != nil { + return result.Error + } + if result.RowsAffected > 0 { + updates := map[string]any{ + "source_username": item.SourceUsername, + "source_email": item.SourceEmail, + "source_quota": item.SourceQuota, + "conflict_flags": item.ConflictFlags, + "status": item.Status, + "snapshot": item.Snapshot, + } + if existing.ResolutionStrategy == "" && item.ResolutionStrategy != "" { + updates["resolution_strategy"] = item.ResolutionStrategy + } + if existing.TargetUserId == 0 && item.TargetUserId != 0 { + updates["target_user_id"] = item.TargetUserId + } + item.Id = existing.Id + return DB.Model(&existing).Updates(updates).Error + } + return DB.Create(item).Error +} + +func ListUserMigrationItemsByBatch(batchId int) ([]*UserMigrationItem, error) { + var items []*UserMigrationItem + err := DB.Where("batch_id = ?", batchId).Order("id asc").Find(&items).Error + return items, err +} + +func ListExecutableUserMigrationItems(batchId int) ([]*UserMigrationItem, error) { + var items []*UserMigrationItem + err := DB.Where("batch_id = ? AND status IN ?", batchId, []string{ + UserMigrationItemStatusReady, + UserMigrationItemStatusMigrating, + UserMigrationItemStatusFailed, + }).Order("id asc").Find(&items).Error + return items, err +} + +func MarkUserMigrationItemFailed(id int, errMsg string) error { + return DB.Model(&UserMigrationItem{}).Where("id = ?", id).Updates(map[string]any{ + "status": UserMigrationItemStatusFailed, + "error_message": errMsg, + }).Error +} + +func MarkUserMigrationItemSkipped(id int) error { + return DB.Model(&UserMigrationItem{}).Where("id = ?", id). + Update("status", UserMigrationItemStatusSkipped).Error +} + +func RefreshUserMigrationBatchStats(batchId int) error { + var total, conflicts, pendingScan int64 + DB.Model(&UserMigrationItem{}).Where("batch_id = ?", batchId).Count(&total) + DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusConflict).Count(&conflicts) + DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusPendingScan).Count(&pendingScan) + + status := UserMigrationBatchStatusScanned + if conflicts == 0 && pendingScan == 0 && total > 0 { + status = UserMigrationBatchStatusReady + } + + return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{ + "total_users": int(total), + "conflict_users": int(conflicts), + "status": status, + }).Error +} + +func RefreshUserMigrationBatchExecutionStats(batchId int) error { + var migrated, failed, pending int64 + DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrated).Count(&migrated) + DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusFailed).Count(&failed) + DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status IN ?", batchId, []string{ + UserMigrationItemStatusReady, + UserMigrationItemStatusMigrating, + }).Count(&pending) + + status := UserMigrationBatchStatusCompleted + if failed > 0 { + status = UserMigrationBatchStatusFailed + } + if pending > 0 { + status = UserMigrationBatchStatusRunning + } + + return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{ + "migrated_users": int(migrated), + "failed_users": int(failed), + "status": status, + }).Error +} + +func TryLockBatchForExecution(batchId int) (bool, error) { + result := DB.Model(&UserMigrationBatch{}). + Where("id = ? AND status IN ?", batchId, []string{ + UserMigrationBatchStatusReady, + UserMigrationBatchStatusFailed, + }). + Update("status", UserMigrationBatchStatusRunning) + return result.RowsAffected == 1, result.Error +} + +func TryLockBatchForScan(batchId int) (bool, error) { + result := DB.Model(&UserMigrationBatch{}). + Where("id = ? AND status IN ?", batchId, []string{ + UserMigrationBatchStatusDraft, + UserMigrationBatchStatusScanned, + UserMigrationBatchStatusReady, + }). + Update("status", UserMigrationBatchStatusScanned) + return result.RowsAffected == 1, result.Error +} + +func AtomicGrantQuotaIfNotExists(tx *gorm.DB, userId, delta int, token string) (bool, error) { + result := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&MigrationQuotaGrant{MigrationToken: token}) + if result.Error != nil { + return false, result.Error + } + if result.RowsAffected == 0 { + return true, nil + } + + return false, tx.Model(&User{}).Where("id = ?", userId). + Update("quota", gorm.Expr("quota + ?", delta)).Error +} diff --git a/model/user_migration_test.go b/model/user_migration_test.go new file mode 100644 index 0000000..65c3a2d --- /dev/null +++ b/model/user_migration_test.go @@ -0,0 +1,150 @@ +package model + +import ( + "testing" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupUserMigrationTestDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + require.NoError(t, db.AutoMigrate( + &User{}, + &CustomOAuthProvider{}, + &UserOAuthBinding{}, + &UserMigrationBatch{}, + &UserMigrationItem{}, + &MigrationQuotaGrant{}, + )) + + return db +} + +func TestUserMigrationModels_AutoMigrate(t *testing.T) { + db := setupUserMigrationTestDB(t) + + require.True(t, db.Migrator().HasTable(&UserMigrationBatch{})) + require.True(t, db.Migrator().HasTable(&UserMigrationItem{})) + require.True(t, db.Migrator().HasTable(&MigrationQuotaGrant{})) +} + +func TestCreateUserMigrationBatch_DefaultStatus(t *testing.T) { + db := setupUserMigrationTestDB(t) + orig := DB + DB = db + defer func() { + DB = orig + }() + + batch := &UserMigrationBatch{ + Name: "ov-to-cn-wave-1", + SourceRegion: "overseas", + TargetRegion: "cn", + OperatorId: 1, + } + + require.NoError(t, CreateUserMigrationBatch(batch)) + require.Equal(t, UserMigrationBatchStatusDraft, batch.Status) + require.Greater(t, batch.Id, 0) +} + +func TestUpsertMigrationItem_PreservesHumanSetStrategy(t *testing.T) { + db := setupUserMigrationTestDB(t) + orig := DB + DB = db + defer func() { + DB = orig + }() + + batch := &UserMigrationBatch{ + Name: "b1", + SourceRegion: "ov", + TargetRegion: "cn", + OperatorId: 1, + } + require.NoError(t, CreateUserMigrationBatch(batch)) + + first := &UserMigrationItem{ + BatchId: batch.Id, + SourceUserId: 10000001, + SourceUsername: "alice", + SourceEmail: "alice@example.com", + SourceQuota: 12345, + ResolutionStrategy: "", + Status: UserMigrationItemStatusConflict, + } + require.NoError(t, UpsertUserMigrationItem(first)) + + require.NoError(t, DB.Model(&UserMigrationItem{}). + Where("batch_id = ? AND source_user_id = ?", batch.Id, 10000001). + Update("resolution_strategy", UserMigrationStrategyMergeExisting).Error) + + second := &UserMigrationItem{ + BatchId: batch.Id, + SourceUserId: 10000001, + SourceUsername: "alice", + SourceEmail: "alice@example.com", + SourceQuota: 99999, + ResolutionStrategy: UserMigrationStrategyCreateNew, + Status: UserMigrationItemStatusReady, + } + require.NoError(t, UpsertUserMigrationItem(second)) + + items, err := ListUserMigrationItemsByBatch(batch.Id) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy) + require.Equal(t, UserMigrationItemStatusReady, items[0].Status) + require.Equal(t, 99999, items[0].SourceQuota) +} + +func TestAtomicGrantQuotaIfNotExists_Idempotent(t *testing.T) { + db := setupUserMigrationTestDB(t) + orig := DB + DB = db + defer func() { + DB = orig + }() + + require.NoError(t, db.Create(&User{ + Id: 1, + Username: "cn-user", + Password: "hash", + AffCode: "AFF1", + Quota: 500, + }).Error) + + var alreadyDone bool + err := db.Transaction(func(tx *gorm.DB) error { + var innerErr error + alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc") + return innerErr + }) + require.NoError(t, err) + require.False(t, alreadyDone) + + var user User + require.NoError(t, db.First(&user, 1).Error) + require.Equal(t, 700, user.Quota) + + err = db.Transaction(func(tx *gorm.DB) error { + var innerErr error + alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc") + return innerErr + }) + require.NoError(t, err) + require.True(t, alreadyDone) + + require.NoError(t, db.First(&user, 1).Error) + require.Equal(t, 700, user.Quota) +}