package user_migration import ( "fmt" "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "gorm.io/gorm" ) func (s *Service) ExecuteBatch(batchId int) error { items, err := model.ListExecutableUserMigrationItems(batchId) if err != nil { return err } providerSlugMap, err := buildProviderSlugMap() if err != nil { return err } for _, item := range items { if item.Status == model.UserMigrationItemStatusMigrated { continue } if err := s.executeItem(item, providerSlugMap); err != nil { _ = model.MarkUserMigrationItemFailed(item.Id, err.Error()) } } return model.RefreshUserMigrationBatchExecutionStats(batchId) } func (s *Service) executeItem(item *model.UserMigrationItem, providerSlugMap map[string]int) error { switch item.ResolutionStrategy { case model.UserMigrationStrategyCreateNew: return s.executeCreateNew(item, providerSlugMap) case model.UserMigrationStrategyMergeExisting: return s.executeMergeExisting(item, providerSlugMap) case model.UserMigrationStrategySkip: return model.MarkUserMigrationItemSkipped(item.Id) default: return fmt.Errorf("unsupported strategy: %q", item.ResolutionStrategy) } } func (s *Service) executeCreateNew(item *model.UserMigrationItem, providerSlugMap map[string]int) error { if item.TargetUserId != 0 { snapshot, err := s.client.GetUser(item.SourceUserId) if err != nil { return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err) } if err := s.revalidateCreateNew(item, snapshot, providerSlugMap); err != nil { return err } if err := markItemMigrating(item.Id); err != nil { return err } cnUser, err := model.GetUserById(item.TargetUserId, true) if err != nil { return err } return s.ensureRemoteConverted(item, cnUser.Quota) } snapshot, err := s.client.GetUser(item.SourceUserId) if err != nil { return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err) } if err := s.revalidateCreateNew(item, snapshot, providerSlugMap); err != nil { return err } if err := markItemMigrating(item.Id); err != nil { return err } affCode := snapshot.AffCode if affCode != "" { if _, err := model.GetUserIdByAffCode(affCode); err == nil { affCode = common.GetRandomString(4) } } cnUser := &model.User{ Username: snapshot.Username, Password: snapshot.Password, Email: snapshot.Email, DisplayName: snapshot.DisplayName, Status: snapshot.Status, Role: snapshot.Role, Group: snapshot.Group, Quota: snapshot.Quota, AffCode: affCode, CreatedAt: snapshot.CreatedAt, Setting: snapshot.Setting, GitHubId: snapshot.GitHubId, DiscordId: snapshot.DiscordId, OidcId: snapshot.OidcId, WeChatId: snapshot.WeChatId, TelegramId: snapshot.TelegramId, LinuxDOId: snapshot.LinuxDOId, } err = model.DB.Transaction(func(tx *gorm.DB) error { if err := model.InsertImportedUserWithHashedPasswordTx(tx, cnUser); err != nil { if isUniqueConstraintErr(err) { return fmt.Errorf("execution-time identity drift detected, rescan required: %w", err) } return fmt.Errorf("create CN user: %w", err) } if err := copyMappedBindingsTx(tx, cnUser.Id, snapshot.OAuthBindings); err != nil { return err } if err := tx.Model(&model.UserMigrationItem{}).Where("id = ?", item.Id). Update("target_user_id", cnUser.Id).Error; err != nil { return err } item.TargetUserId = cnUser.Id return nil }) if err != nil { return err } return s.ensureRemoteConverted(item, cnUser.Quota) } func (s *Service) executeMergeExisting(item *model.UserMigrationItem, providerSlugMap map[string]int) error { snapshot, err := s.client.GetUser(item.SourceUserId) if err != nil { return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err) } if err := s.revalidateMergeExisting(item, snapshot, providerSlugMap); err != nil { return err } if err := markItemMigrating(item.Id); err != nil { return err } finalQuota := 0 err = model.DB.Transaction(func(tx *gorm.DB) error { _, err := model.AtomicGrantQuotaIfNotExists(tx, item.TargetUserId, snapshot.Quota, item.MigrationToken) if err != nil { return err } var cnUser model.User if err := tx.Select("quota").First(&cnUser, item.TargetUserId).Error; err != nil { return err } finalQuota = cnUser.Quota return copyMappedBindingsTx(tx, item.TargetUserId, snapshot.OAuthBindings) }) if err != nil { return err } return s.ensureRemoteConverted(item, finalQuota) } func copyMappedBindingsTx(tx *gorm.DB, userId int, bindings []RemoteOAuthBinding) error { if len(bindings) == 0 { return nil } slugs := make([]string, 0, len(bindings)) for _, binding := range bindings { slugs = append(slugs, binding.ProviderSlug) } providers, err := model.GetCustomOAuthProvidersBySlugs(slugs) if err != nil { return err } slugToID := make(map[string]int, len(providers)) for _, provider := range providers { slugToID[provider.Slug] = provider.Id } localBindings := make([]*model.UserOAuthBinding, 0, len(bindings)) for _, binding := range bindings { localID, ok := slugToID[binding.ProviderSlug] if !ok { continue } localBindings = append(localBindings, &model.UserOAuthBinding{ ProviderId: localID, ProviderUserId: binding.ProviderUserId, }) } return model.CopyOAuthBindingsIfMissingTx(tx, userId, localBindings) } func (s *Service) ensureRemoteConverted(item *model.UserMigrationItem, syncedQuota int) error { remoteUser, err := s.client.GetUser(item.SourceUserId) if err != nil { return err } if remoteUser.Source != common.UserSourceSynced || remoteUser.RemoteUserId != item.TargetUserId { if err := s.client.ConvertToSynced(item.SourceUserId, &ConvertRemoteUserRequest{ RemoteUserId: item.TargetUserId, SyncedQuota: syncedQuota, }); err != nil { return err } } return markItemMigrated(item.Id) } func (s *Service) revalidateCreateNew(item *model.UserMigrationItem, snapshot *RemoteUserSnapshot, providerSlugMap map[string]int) error { if snapshot.Source == common.UserSourceSynced { if item.TargetUserId != 0 && snapshot.RemoteUserId == item.TargetUserId { return nil } return fmt.Errorf("execution-time identity drift detected, remote user %d is already synced, rescan required", item.SourceUserId) } result, err := s.analyzeExecutionSnapshot(snapshot, providerSlugMap) if err != nil { return err } if item.TargetUserId == 0 { if result.Status != model.UserMigrationItemStatusReady || result.Strategy != model.UserMigrationStrategyCreateNew || result.TargetCNUserID != 0 { return fmt.Errorf("execution-time identity drift detected, rescan required") } return nil } if result.Status != model.UserMigrationItemStatusReady { return fmt.Errorf("execution-time identity drift detected, rescan required") } if result.Strategy == model.UserMigrationStrategyCreateNew && result.TargetCNUserID == 0 { return nil } if result.Strategy == model.UserMigrationStrategyMergeExisting && result.TargetCNUserID == item.TargetUserId { return nil } return fmt.Errorf("execution-time identity drift detected, rescan required") } func (s *Service) revalidateMergeExisting(item *model.UserMigrationItem, snapshot *RemoteUserSnapshot, providerSlugMap map[string]int) error { if snapshot.Source == common.UserSourceSynced { if snapshot.RemoteUserId == item.TargetUserId { return nil } return fmt.Errorf("execution-time identity drift detected, remote user %d is already synced to CN user %d", item.SourceUserId, snapshot.RemoteUserId) } result, err := s.analyzeExecutionSnapshot(snapshot, providerSlugMap) if err != nil { return err } if result.Status != model.UserMigrationItemStatusReady || result.Strategy != model.UserMigrationStrategyMergeExisting || result.TargetCNUserID != item.TargetUserId { return fmt.Errorf("execution-time merge target drift detected, rescan required") } return nil } func (s *Service) analyzeExecutionSnapshot(snapshot *RemoteUserSnapshot, providerSlugMap map[string]int) (*conflictResult, error) { result, err := analyzeConflicts(snapshot, providerSlugMap, s.client) if err != nil { return nil, fmt.Errorf("revalidate conflicts for user %d: %w", snapshot.Id, err) } return result, nil } func markItemMigrating(itemID int) error { return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{ "status": model.UserMigrationItemStatusMigrating, "error_message": "", }).Error } func markItemMigrated(itemID int) error { return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{ "status": model.UserMigrationItemStatusMigrated, "migrated_at": timeNow(), "error_message": "", }).Error } func isUniqueConstraintErr(err error) bool { if err == nil { return false } message := err.Error() return strings.Contains(message, "UNIQUE constraint failed") || strings.Contains(message, "Duplicate entry") || strings.Contains(strings.ToLower(message), "duplicate key value") } func timeNow() int64 { return common.GetTimestamp() }