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 }