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 } for _, item := range items { if item.Status == model.UserMigrationItemStatusMigrated { continue } if err := s.executeItem(item); err != nil { _ = model.MarkUserMigrationItemFailed(item.Id, err.Error()) } } return model.RefreshUserMigrationBatchExecutionStats(batchId) } func (s *Service) executeItem(item *model.UserMigrationItem) error { switch item.ResolutionStrategy { case model.UserMigrationStrategyCreateNew: return s.executeCreateNew(item) case model.UserMigrationStrategyMergeExisting: return s.executeMergeExisting(item) case model.UserMigrationStrategySkip: return model.MarkUserMigrationItemSkipped(item.Id) default: return fmt.Errorf("unsupported strategy: %q", item.ResolutionStrategy) } } func (s *Service) executeCreateNew(item *model.UserMigrationItem) error { if item.TargetUserId != 0 { 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) } 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) error { snapshot, err := s.client.GetUser(item.SourceUserId) if err != nil { return fmt.Errorf("get remote user %d: %w", item.SourceUserId, 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 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() }