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