|
- 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()
- }
|