|
|
@@ -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 |
|
|
|
|
|
} |