|
- 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
- }
-
- page := 0
- for {
- resp, err := s.client.ListUsers(page, scanPageSize)
- if err != nil {
- return fmt.Errorf("list users page %d: %w", page, err)
- }
- for _, remoteUser := range resp.Data {
- 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(resp.Data) < scanPageSize {
- break
- }
- page++
- }
-
- 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
- }
|