|
- package user_migration
-
- import (
- "errors"
- "fmt"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "gorm.io/gorm"
- )
-
- const scanPageSize = 100
-
- type Service struct {
- client migrationClient
- }
-
- func NewService(client migrationClient) *Service {
- return &Service{client: client}
- }
-
- func (s *Service) ScanBatch(batchId int) error {
- providerSlugMap, err := buildProviderSlugMap()
- if err != nil {
- return err
- }
-
- var batch model.UserMigrationBatch
- if err := model.DB.First(&batch, batchId).Error; err != nil {
- return err
- }
-
- var selectionPayload SelectionPayload
- if batch.SelectionPayload != "" {
- if err := common.UnmarshalJsonStr(batch.SelectionPayload, &selectionPayload); err != nil {
- return fmt.Errorf("parse selection payload: %w", err)
- }
- }
-
- page := 0
- var latestSummary *SelectionSummary
- for {
- var respData []RemoteUserSnapshot
- var total int
- if batch.SelectionMode == SelectionModeExplicitIDs {
- resp, err := s.client.QueryUsers(page, scanPageSize, selectionPayload.SourceUserIDs)
- if err != nil {
- return fmt.Errorf("query users page %d: %w", page, err)
- }
- respData = resp.Data
- total = resp.Total
- latestSummary = resp.SelectionSummary
- } else {
- resp, err := s.client.ListUsers(page, scanPageSize, "")
- if err != nil {
- return fmt.Errorf("list users page %d: %w", page, err)
- }
- respData = resp.Data
- total = resp.Total
- }
- for _, remoteUser := range respData {
- result, err := analyzeConflicts(&remoteUser, providerSlugMap, s.client)
- if err != nil {
- return fmt.Errorf("analyze conflicts for user %d: %w", remoteUser.Id, err)
- }
-
- flagsJSON, err := common.Marshal(result.Flags)
- if err != nil {
- return err
- }
- snapshotJSON, err := common.Marshal(remoteUser)
- if err != nil {
- return err
- }
-
- item := &model.UserMigrationItem{
- BatchId: batchId,
- SourceUserId: remoteUser.Id,
- SourceUsername: remoteUser.Username,
- SourceEmail: remoteUser.Email,
- SourceQuota: remoteUser.Quota,
- ConflictFlags: string(flagsJSON),
- ResolutionStrategy: result.Strategy,
- Status: result.Status,
- Snapshot: string(snapshotJSON),
- MigrationToken: fmt.Sprintf("batch-%d-user-%d", batchId, remoteUser.Id),
- TargetUserId: result.TargetCNUserID,
- }
- if err := model.UpsertUserMigrationItem(item); err != nil {
- return err
- }
- }
-
- if len(respData) < scanPageSize || total <= (page+1)*scanPageSize {
- break
- }
- page++
- }
-
- if latestSummary != nil {
- summaryJSON, err := common.Marshal(latestSummary)
- if err != nil {
- return err
- }
- if err := model.DB.Model(&model.UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
- "matched_user_count": latestSummary.Matched,
- "excluded_user_count": len(latestSummary.Excluded),
- "selection_summary": string(summaryJSON),
- }).Error; err != nil {
- return err
- }
- }
-
- return model.RefreshUserMigrationBatchStats(batchId)
- }
-
- type conflictResult struct {
- Flags []string
- Strategy string
- Status string
- TargetCNUserID int
- }
-
- func analyzeConflicts(remote *RemoteUserSnapshot, providerSlugMap map[string]int, client migrationClient) (*conflictResult, error) {
- result := &conflictResult{}
- var mergeCandidate *model.User
- hardConflict := false
-
- assignMergeCandidate := func(flag string, user *model.User) {
- if user == nil {
- return
- }
- result.Flags = append(result.Flags, flag)
- if mergeCandidate == nil {
- mergeCandidate = user
- return
- }
- if mergeCandidate.Id != user.Id {
- hardConflict = true
- result.Flags = append(result.Flags, "multiple_merge_targets")
- }
- }
-
- if user, _ := model.GetUserByUsername(remote.Username); user != nil {
- result.Flags = append(result.Flags, "username")
- hardConflict = true
- }
-
- if remote.Email != "" {
- user, err := model.GetUserByEmail(remote.Email)
- if err != nil {
- return nil, err
- }
- assignMergeCandidate("email", user)
- }
-
- type oauthField struct {
- value string
- flag string
- query func(string) (*model.User, error)
- }
- fields := []oauthField{
- {value: remote.GitHubId, flag: "github_id", query: model.GetUserByGitHubId},
- {value: remote.DiscordId, flag: "discord_id", query: model.GetUserByDiscordId},
- {value: remote.OidcId, flag: "oidc_id", query: model.GetUserByOidcId},
- {value: remote.WeChatId, flag: "wechat_id", query: model.GetUserByWeChatId},
- {value: remote.TelegramId, flag: "telegram_id", query: model.GetUserByTelegramId},
- {value: remote.LinuxDOId, flag: "linux_do_id", query: model.GetUserByLinuxDOId},
- }
- for _, field := range fields {
- if field.value == "" {
- continue
- }
- user, err := field.query(field.value)
- if err != nil {
- return nil, err
- }
- assignMergeCandidate(field.flag, user)
- }
-
- for _, binding := range remote.OAuthBindings {
- localProviderID, ok := providerSlugMap[binding.ProviderSlug]
- if !ok {
- continue
- }
- user, err := model.GetUserByOAuthBinding(localProviderID, binding.ProviderUserId)
- if errors.Is(err, gorm.ErrRecordNotFound) {
- continue
- }
- if err != nil {
- return nil, err
- }
- assignMergeCandidate("oauth:"+binding.ProviderSlug, user)
- }
-
- if remote.AffCode != "" {
- if _, err := model.GetUserIdByAffCode(remote.AffCode); err == nil {
- result.Flags = append(result.Flags, "aff_code")
- }
- }
-
- if hardConflict {
- result.Status = model.UserMigrationItemStatusConflict
- return result, nil
- }
-
- if mergeCandidate != nil {
- hasCopy, err := client.HasSyncedCopy(mergeCandidate.Id)
- if err != nil {
- return nil, fmt.Errorf("check synced copy for CN user %d: %w", mergeCandidate.Id, err)
- }
- if hasCopy {
- result.Flags = append(result.Flags, "cn_already_synced_to_ov")
- result.Status = model.UserMigrationItemStatusConflict
- return result, nil
- }
- result.Status = model.UserMigrationItemStatusReady
- result.Strategy = model.UserMigrationStrategyMergeExisting
- result.TargetCNUserID = mergeCandidate.Id
- return result, nil
- }
-
- result.Status = model.UserMigrationItemStatusReady
- result.Strategy = model.UserMigrationStrategyCreateNew
- return result, nil
- }
|