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.
 
 
 

227 lines
6.0 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. var batch model.UserMigrationBatch
  22. if err := model.DB.First(&batch, batchId).Error; err != nil {
  23. return err
  24. }
  25. var selectionPayload SelectionPayload
  26. if batch.SelectionPayload != "" {
  27. if err := common.UnmarshalJsonStr(batch.SelectionPayload, &selectionPayload); err != nil {
  28. return fmt.Errorf("parse selection payload: %w", err)
  29. }
  30. }
  31. page := 0
  32. var latestSummary *SelectionSummary
  33. for {
  34. var respData []RemoteUserSnapshot
  35. var total int
  36. if batch.SelectionMode == SelectionModeExplicitIDs {
  37. resp, err := s.client.QueryUsers(page, scanPageSize, selectionPayload.SourceUserIDs)
  38. if err != nil {
  39. return fmt.Errorf("query users page %d: %w", page, err)
  40. }
  41. respData = resp.Data
  42. total = resp.Total
  43. latestSummary = resp.SelectionSummary
  44. } else {
  45. resp, err := s.client.ListUsers(page, scanPageSize, "")
  46. if err != nil {
  47. return fmt.Errorf("list users page %d: %w", page, err)
  48. }
  49. respData = resp.Data
  50. total = resp.Total
  51. }
  52. for _, remoteUser := range respData {
  53. result, err := analyzeConflicts(&remoteUser, providerSlugMap, s.client)
  54. if err != nil {
  55. return fmt.Errorf("analyze conflicts for user %d: %w", remoteUser.Id, err)
  56. }
  57. flagsJSON, err := common.Marshal(result.Flags)
  58. if err != nil {
  59. return err
  60. }
  61. snapshotJSON, err := common.Marshal(remoteUser)
  62. if err != nil {
  63. return err
  64. }
  65. item := &model.UserMigrationItem{
  66. BatchId: batchId,
  67. SourceUserId: remoteUser.Id,
  68. SourceUsername: remoteUser.Username,
  69. SourceEmail: remoteUser.Email,
  70. SourceQuota: remoteUser.Quota,
  71. ConflictFlags: string(flagsJSON),
  72. ResolutionStrategy: result.Strategy,
  73. Status: result.Status,
  74. Snapshot: string(snapshotJSON),
  75. MigrationToken: fmt.Sprintf("batch-%d-user-%d", batchId, remoteUser.Id),
  76. TargetUserId: result.TargetCNUserID,
  77. }
  78. if err := model.UpsertUserMigrationItem(item); err != nil {
  79. return err
  80. }
  81. }
  82. if len(respData) < scanPageSize || total <= (page+1)*scanPageSize {
  83. break
  84. }
  85. page++
  86. }
  87. if latestSummary != nil {
  88. summaryJSON, err := common.Marshal(latestSummary)
  89. if err != nil {
  90. return err
  91. }
  92. if err := model.DB.Model(&model.UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
  93. "matched_user_count": latestSummary.Matched,
  94. "excluded_user_count": len(latestSummary.Excluded),
  95. "selection_summary": string(summaryJSON),
  96. }).Error; err != nil {
  97. return err
  98. }
  99. }
  100. return model.RefreshUserMigrationBatchStats(batchId)
  101. }
  102. type conflictResult struct {
  103. Flags []string
  104. Strategy string
  105. Status string
  106. TargetCNUserID int
  107. }
  108. func analyzeConflicts(remote *RemoteUserSnapshot, providerSlugMap map[string]int, client migrationClient) (*conflictResult, error) {
  109. result := &conflictResult{}
  110. var mergeCandidate *model.User
  111. hardConflict := false
  112. assignMergeCandidate := func(flag string, user *model.User) {
  113. if user == nil {
  114. return
  115. }
  116. result.Flags = append(result.Flags, flag)
  117. if mergeCandidate == nil {
  118. mergeCandidate = user
  119. return
  120. }
  121. if mergeCandidate.Id != user.Id {
  122. hardConflict = true
  123. result.Flags = append(result.Flags, "multiple_merge_targets")
  124. }
  125. }
  126. if user, _ := model.GetUserByUsername(remote.Username); user != nil {
  127. result.Flags = append(result.Flags, "username")
  128. hardConflict = true
  129. }
  130. if remote.Email != "" {
  131. user, err := model.GetUserByEmail(remote.Email)
  132. if err != nil {
  133. return nil, err
  134. }
  135. assignMergeCandidate("email", user)
  136. }
  137. type oauthField struct {
  138. value string
  139. flag string
  140. query func(string) (*model.User, error)
  141. }
  142. fields := []oauthField{
  143. {value: remote.GitHubId, flag: "github_id", query: model.GetUserByGitHubId},
  144. {value: remote.DiscordId, flag: "discord_id", query: model.GetUserByDiscordId},
  145. {value: remote.OidcId, flag: "oidc_id", query: model.GetUserByOidcId},
  146. {value: remote.WeChatId, flag: "wechat_id", query: model.GetUserByWeChatId},
  147. {value: remote.TelegramId, flag: "telegram_id", query: model.GetUserByTelegramId},
  148. {value: remote.LinuxDOId, flag: "linux_do_id", query: model.GetUserByLinuxDOId},
  149. }
  150. for _, field := range fields {
  151. if field.value == "" {
  152. continue
  153. }
  154. user, err := field.query(field.value)
  155. if err != nil {
  156. return nil, err
  157. }
  158. assignMergeCandidate(field.flag, user)
  159. }
  160. for _, binding := range remote.OAuthBindings {
  161. localProviderID, ok := providerSlugMap[binding.ProviderSlug]
  162. if !ok {
  163. continue
  164. }
  165. user, err := model.GetUserByOAuthBinding(localProviderID, binding.ProviderUserId)
  166. if errors.Is(err, gorm.ErrRecordNotFound) {
  167. continue
  168. }
  169. if err != nil {
  170. return nil, err
  171. }
  172. assignMergeCandidate("oauth:"+binding.ProviderSlug, user)
  173. }
  174. if remote.AffCode != "" {
  175. if _, err := model.GetUserIdByAffCode(remote.AffCode); err == nil {
  176. result.Flags = append(result.Flags, "aff_code")
  177. }
  178. }
  179. if hardConflict {
  180. result.Status = model.UserMigrationItemStatusConflict
  181. return result, nil
  182. }
  183. if mergeCandidate != nil {
  184. hasCopy, err := client.HasSyncedCopy(mergeCandidate.Id)
  185. if err != nil {
  186. return nil, fmt.Errorf("check synced copy for CN user %d: %w", mergeCandidate.Id, err)
  187. }
  188. if hasCopy {
  189. result.Flags = append(result.Flags, "cn_already_synced_to_ov")
  190. result.Status = model.UserMigrationItemStatusConflict
  191. return result, nil
  192. }
  193. result.Status = model.UserMigrationItemStatusReady
  194. result.Strategy = model.UserMigrationStrategyMergeExisting
  195. result.TargetCNUserID = mergeCandidate.Id
  196. return result, nil
  197. }
  198. result.Status = model.UserMigrationItemStatusReady
  199. result.Strategy = model.UserMigrationStrategyCreateNew
  200. return result, nil
  201. }