|
- package model
-
- import (
- "github.com/QuantumNous/new-api/common"
- "gorm.io/gorm"
- "gorm.io/gorm/clause"
- )
-
- const (
- UserMigrationBatchStatusDraft = "draft"
- UserMigrationBatchStatusScanning = "scanning"
- 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", UserMigrationBatchStatusScanning)
- 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
- }
|