You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

186 lines
4.7 KiB

  1. package user_migration
  2. import (
  3. "errors"
  4. "fmt"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "gorm.io/gorm"
  8. )
  9. const scanPageSize = 100
  10. type Service struct {
  11. client migrationClient
  12. }
  13. func NewService(client migrationClient) *Service {
  14. return &Service{client: client}
  15. }
  16. func (s *Service) ScanBatch(batchId int) error {
  17. providerSlugMap, err := buildProviderSlugMap()
  18. if err != nil {
  19. return err
  20. }
  21. page := 0
  22. for {
  23. resp, err := s.client.ListUsers(page, scanPageSize)
  24. if err != nil {
  25. return fmt.Errorf("list users page %d: %w", page, err)
  26. }
  27. for _, remoteUser := range resp.Data {
  28. result, err := analyzeConflicts(&remoteUser, providerSlugMap, s.client)
  29. if err != nil {
  30. return fmt.Errorf("analyze conflicts for user %d: %w", remoteUser.Id, err)
  31. }
  32. flagsJSON, err := common.Marshal(result.Flags)
  33. if err != nil {
  34. return err
  35. }
  36. snapshotJSON, err := common.Marshal(remoteUser)
  37. if err != nil {
  38. return err
  39. }
  40. item := &model.UserMigrationItem{
  41. BatchId: batchId,
  42. SourceUserId: remoteUser.Id,
  43. SourceUsername: remoteUser.Username,
  44. SourceEmail: remoteUser.Email,
  45. SourceQuota: remoteUser.Quota,
  46. ConflictFlags: string(flagsJSON),
  47. ResolutionStrategy: result.Strategy,
  48. Status: result.Status,
  49. Snapshot: string(snapshotJSON),
  50. MigrationToken: fmt.Sprintf("batch-%d-user-%d", batchId, remoteUser.Id),
  51. TargetUserId: result.TargetCNUserID,
  52. }
  53. if err := model.UpsertUserMigrationItem(item); err != nil {
  54. return err
  55. }
  56. }
  57. if len(resp.Data) < scanPageSize {
  58. break
  59. }
  60. page++
  61. }
  62. return model.RefreshUserMigrationBatchStats(batchId)
  63. }
  64. type conflictResult struct {
  65. Flags []string
  66. Strategy string
  67. Status string
  68. TargetCNUserID int
  69. }
  70. func analyzeConflicts(remote *RemoteUserSnapshot, providerSlugMap map[string]int, client migrationClient) (*conflictResult, error) {
  71. result := &conflictResult{}
  72. var mergeCandidate *model.User
  73. hardConflict := false
  74. assignMergeCandidate := func(flag string, user *model.User) {
  75. if user == nil {
  76. return
  77. }
  78. result.Flags = append(result.Flags, flag)
  79. if mergeCandidate == nil {
  80. mergeCandidate = user
  81. return
  82. }
  83. if mergeCandidate.Id != user.Id {
  84. hardConflict = true
  85. result.Flags = append(result.Flags, "multiple_merge_targets")
  86. }
  87. }
  88. if user, _ := model.GetUserByUsername(remote.Username); user != nil {
  89. result.Flags = append(result.Flags, "username")
  90. hardConflict = true
  91. }
  92. if remote.Email != "" {
  93. user, err := model.GetUserByEmail(remote.Email)
  94. if err != nil {
  95. return nil, err
  96. }
  97. assignMergeCandidate("email", user)
  98. }
  99. type oauthField struct {
  100. value string
  101. flag string
  102. query func(string) (*model.User, error)
  103. }
  104. fields := []oauthField{
  105. {value: remote.GitHubId, flag: "github_id", query: model.GetUserByGitHubId},
  106. {value: remote.DiscordId, flag: "discord_id", query: model.GetUserByDiscordId},
  107. {value: remote.OidcId, flag: "oidc_id", query: model.GetUserByOidcId},
  108. {value: remote.WeChatId, flag: "wechat_id", query: model.GetUserByWeChatId},
  109. {value: remote.TelegramId, flag: "telegram_id", query: model.GetUserByTelegramId},
  110. {value: remote.LinuxDOId, flag: "linux_do_id", query: model.GetUserByLinuxDOId},
  111. }
  112. for _, field := range fields {
  113. if field.value == "" {
  114. continue
  115. }
  116. user, err := field.query(field.value)
  117. if err != nil {
  118. return nil, err
  119. }
  120. assignMergeCandidate(field.flag, user)
  121. }
  122. for _, binding := range remote.OAuthBindings {
  123. localProviderID, ok := providerSlugMap[binding.ProviderSlug]
  124. if !ok {
  125. continue
  126. }
  127. user, err := model.GetUserByOAuthBinding(localProviderID, binding.ProviderUserId)
  128. if errors.Is(err, gorm.ErrRecordNotFound) {
  129. continue
  130. }
  131. if err != nil {
  132. return nil, err
  133. }
  134. assignMergeCandidate("oauth:"+binding.ProviderSlug, user)
  135. }
  136. if remote.AffCode != "" {
  137. if _, err := model.GetUserIdByAffCode(remote.AffCode); err == nil {
  138. result.Flags = append(result.Flags, "aff_code")
  139. }
  140. }
  141. if hardConflict {
  142. result.Status = model.UserMigrationItemStatusConflict
  143. return result, nil
  144. }
  145. if mergeCandidate != nil {
  146. hasCopy, err := client.HasSyncedCopy(mergeCandidate.Id)
  147. if err != nil {
  148. return nil, fmt.Errorf("check synced copy for CN user %d: %w", mergeCandidate.Id, err)
  149. }
  150. if hasCopy {
  151. result.Flags = append(result.Flags, "cn_already_synced_to_ov")
  152. result.Status = model.UserMigrationItemStatusConflict
  153. return result, nil
  154. }
  155. result.Status = model.UserMigrationItemStatusReady
  156. result.Strategy = model.UserMigrationStrategyMergeExisting
  157. result.TargetCNUserID = mergeCandidate.Id
  158. return result, nil
  159. }
  160. result.Status = model.UserMigrationItemStatusReady
  161. result.Strategy = model.UserMigrationStrategyCreateNew
  162. return result, nil
  163. }