|
- package model
-
- import (
- "fmt"
- "slices"
-
- "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"`
- SelectionMode string `json:"selection_mode" gorm:"type:varchar(32);not null;default:'all'"`
- SelectionPayload string `json:"selection_payload" gorm:"type:text"`
- RequestedUserCount int `json:"requested_user_count" gorm:"default:0"`
- MatchedUserCount int `json:"matched_user_count" gorm:"default:0"`
- ExcludedUserCount int `json:"excluded_user_count" gorm:"default:0"`
- SelectionSummary string `json:"selection_summary" gorm:"type:text"`
- 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"`
- }
-
- type ActiveMigrationUserConflict struct {
- SourceUserID int `json:"source_user_id"`
- BatchID int `json:"batch_id"`
- BatchStatus string `json:"batch_status"`
- }
-
- func (b *UserMigrationBatch) BeforeCreate(tx *gorm.DB) error {
- now := common.GetTimestamp()
- if b.Status == "" {
- b.Status = UserMigrationBatchStatusDraft
- }
- if b.SelectionMode == "" {
- b.SelectionMode = "all"
- }
- 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 CancelUserMigrationBatch(batchId int) (bool, error) {
- var batch UserMigrationBatch
- if err := DB.First(&batch, batchId).Error; err != nil {
- return false, err
- }
-
- switch batch.Status {
- case UserMigrationBatchStatusCancelled:
- return true, nil
- case UserMigrationBatchStatusDraft,
- UserMigrationBatchStatusScanned,
- UserMigrationBatchStatusReady,
- UserMigrationBatchStatusFailed:
- return true, DB.Model(&UserMigrationBatch{}).
- Where("id = ?", batchId).
- Update("status", UserMigrationBatchStatusCancelled).Error
- default:
- return false, nil
- }
- }
-
- func FindActiveMigrationUserConflicts(sourceUserIDs []int) ([]ActiveMigrationUserConflict, error) {
- if len(sourceUserIDs) == 0 {
- return nil, nil
- }
-
- activeStatuses := []string{
- UserMigrationBatchStatusDraft,
- UserMigrationBatchStatusScanning,
- UserMigrationBatchStatusScanned,
- UserMigrationBatchStatusReady,
- UserMigrationBatchStatusRunning,
- UserMigrationBatchStatusFailed,
- }
-
- type activeBatchSelection struct {
- ID int
- Status string
- SelectionMode string
- SelectionPayload string
- }
-
- conflictMap := make(map[string]ActiveMigrationUserConflict)
- addConflict := func(conflict ActiveMigrationUserConflict) {
- key := fmt.Sprintf("%d:%d", conflict.SourceUserID, conflict.BatchID)
- if _, exists := conflictMap[key]; exists {
- return
- }
- conflictMap[key] = conflict
- }
-
- var itemConflicts []ActiveMigrationUserConflict
- err := DB.Table("user_migration_items AS umi").
- Select("umi.source_user_id, umb.id AS batch_id, umb.status AS batch_status").
- Joins("JOIN user_migration_batches AS umb ON umb.id = umi.batch_id").
- Where("umi.source_user_id IN ?", sourceUserIDs).
- Where("umb.status IN ?", activeStatuses).
- Order("umi.source_user_id ASC, umb.id ASC").
- Scan(&itemConflicts).Error
- if err != nil {
- return nil, err
- }
- for _, conflict := range itemConflicts {
- addConflict(conflict)
- }
-
- var activeBatches []activeBatchSelection
- if err := DB.Model(&UserMigrationBatch{}).
- Select("id, status, selection_mode, selection_payload").
- Where("status IN ?", activeStatuses).
- Where("selection_mode = ?", "explicit_ids").
- Find(&activeBatches).Error; err != nil {
- return nil, err
- }
-
- sourceIDSet := make(map[int]struct{}, len(sourceUserIDs))
- for _, id := range sourceUserIDs {
- sourceIDSet[id] = struct{}{}
- }
-
- for _, batch := range activeBatches {
- if batch.SelectionPayload == "" {
- continue
- }
- var payload struct {
- SourceUserIDs []int `json:"source_user_ids"`
- }
- if err := common.UnmarshalJsonStr(batch.SelectionPayload, &payload); err != nil {
- return nil, err
- }
- for _, sourceUserID := range payload.SourceUserIDs {
- if _, ok := sourceIDSet[sourceUserID]; !ok {
- continue
- }
- addConflict(ActiveMigrationUserConflict{
- SourceUserID: sourceUserID,
- BatchID: batch.ID,
- BatchStatus: batch.Status,
- })
- }
- }
-
- conflicts := make([]ActiveMigrationUserConflict, 0, len(conflictMap))
- for _, conflict := range conflictMap {
- conflicts = append(conflicts, conflict)
- }
- slices.SortFunc(conflicts, func(a, b ActiveMigrationUserConflict) int {
- if a.SourceUserID != b.SourceUserID {
- return a.SourceUserID - b.SourceUserID
- }
- return a.BatchID - b.BatchID
- })
- return conflicts, nil
- }
-
- 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, migrating 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)
- DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrating).Count(&migrating)
-
- status := UserMigrationBatchStatusCompleted
- if failed > 0 {
- status = UserMigrationBatchStatusFailed
- }
- if pending > 0 {
- status = UserMigrationBatchStatusRunning
- }
- if pending == migrating && migrating > 0 {
- status = UserMigrationBatchStatusFailed
- }
-
- 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
- }
|