Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

209 wiersze
5.5 KiB

  1. package user_migration
  2. import (
  3. "fmt"
  4. "strings"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "gorm.io/gorm"
  8. )
  9. func (s *Service) ExecuteBatch(batchId int) error {
  10. items, err := model.ListExecutableUserMigrationItems(batchId)
  11. if err != nil {
  12. return err
  13. }
  14. for _, item := range items {
  15. if item.Status == model.UserMigrationItemStatusMigrated {
  16. continue
  17. }
  18. if err := s.executeItem(item); err != nil {
  19. _ = model.MarkUserMigrationItemFailed(item.Id, err.Error())
  20. }
  21. }
  22. return model.RefreshUserMigrationBatchExecutionStats(batchId)
  23. }
  24. func (s *Service) executeItem(item *model.UserMigrationItem) error {
  25. switch item.ResolutionStrategy {
  26. case model.UserMigrationStrategyCreateNew:
  27. return s.executeCreateNew(item)
  28. case model.UserMigrationStrategyMergeExisting:
  29. return s.executeMergeExisting(item)
  30. case model.UserMigrationStrategySkip:
  31. return model.MarkUserMigrationItemSkipped(item.Id)
  32. default:
  33. return fmt.Errorf("unsupported strategy: %q", item.ResolutionStrategy)
  34. }
  35. }
  36. func (s *Service) executeCreateNew(item *model.UserMigrationItem) error {
  37. if item.TargetUserId != 0 {
  38. cnUser, err := model.GetUserById(item.TargetUserId, true)
  39. if err != nil {
  40. return err
  41. }
  42. return s.ensureRemoteConverted(item, cnUser.Quota)
  43. }
  44. snapshot, err := s.client.GetUser(item.SourceUserId)
  45. if err != nil {
  46. return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
  47. }
  48. affCode := snapshot.AffCode
  49. if affCode != "" {
  50. if _, err := model.GetUserIdByAffCode(affCode); err == nil {
  51. affCode = common.GetRandomString(4)
  52. }
  53. }
  54. cnUser := &model.User{
  55. Username: snapshot.Username,
  56. Password: snapshot.Password,
  57. Email: snapshot.Email,
  58. DisplayName: snapshot.DisplayName,
  59. Status: snapshot.Status,
  60. Role: snapshot.Role,
  61. Group: snapshot.Group,
  62. Quota: snapshot.Quota,
  63. AffCode: affCode,
  64. CreatedAt: snapshot.CreatedAt,
  65. Setting: snapshot.Setting,
  66. GitHubId: snapshot.GitHubId,
  67. DiscordId: snapshot.DiscordId,
  68. OidcId: snapshot.OidcId,
  69. WeChatId: snapshot.WeChatId,
  70. TelegramId: snapshot.TelegramId,
  71. LinuxDOId: snapshot.LinuxDOId,
  72. }
  73. err = model.DB.Transaction(func(tx *gorm.DB) error {
  74. if err := model.InsertImportedUserWithHashedPasswordTx(tx, cnUser); err != nil {
  75. if isUniqueConstraintErr(err) {
  76. return fmt.Errorf("execution-time identity drift detected, rescan required: %w", err)
  77. }
  78. return fmt.Errorf("create CN user: %w", err)
  79. }
  80. if err := copyMappedBindingsTx(tx, cnUser.Id, snapshot.OAuthBindings); err != nil {
  81. return err
  82. }
  83. if err := tx.Model(&model.UserMigrationItem{}).Where("id = ?", item.Id).
  84. Update("target_user_id", cnUser.Id).Error; err != nil {
  85. return err
  86. }
  87. item.TargetUserId = cnUser.Id
  88. return nil
  89. })
  90. if err != nil {
  91. return err
  92. }
  93. return s.ensureRemoteConverted(item, cnUser.Quota)
  94. }
  95. func (s *Service) executeMergeExisting(item *model.UserMigrationItem) error {
  96. snapshot, err := s.client.GetUser(item.SourceUserId)
  97. if err != nil {
  98. return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
  99. }
  100. finalQuota := 0
  101. err = model.DB.Transaction(func(tx *gorm.DB) error {
  102. _, err := model.AtomicGrantQuotaIfNotExists(tx, item.TargetUserId, snapshot.Quota, item.MigrationToken)
  103. if err != nil {
  104. return err
  105. }
  106. var cnUser model.User
  107. if err := tx.Select("quota").First(&cnUser, item.TargetUserId).Error; err != nil {
  108. return err
  109. }
  110. finalQuota = cnUser.Quota
  111. return copyMappedBindingsTx(tx, item.TargetUserId, snapshot.OAuthBindings)
  112. })
  113. if err != nil {
  114. return err
  115. }
  116. return s.ensureRemoteConverted(item, finalQuota)
  117. }
  118. func copyMappedBindingsTx(tx *gorm.DB, userId int, bindings []RemoteOAuthBinding) error {
  119. if len(bindings) == 0 {
  120. return nil
  121. }
  122. slugs := make([]string, 0, len(bindings))
  123. for _, binding := range bindings {
  124. slugs = append(slugs, binding.ProviderSlug)
  125. }
  126. providers, err := model.GetCustomOAuthProvidersBySlugs(slugs)
  127. if err != nil {
  128. return err
  129. }
  130. slugToID := make(map[string]int, len(providers))
  131. for _, provider := range providers {
  132. slugToID[provider.Slug] = provider.Id
  133. }
  134. localBindings := make([]*model.UserOAuthBinding, 0, len(bindings))
  135. for _, binding := range bindings {
  136. localID, ok := slugToID[binding.ProviderSlug]
  137. if !ok {
  138. continue
  139. }
  140. localBindings = append(localBindings, &model.UserOAuthBinding{
  141. ProviderId: localID,
  142. ProviderUserId: binding.ProviderUserId,
  143. })
  144. }
  145. return model.CopyOAuthBindingsIfMissingTx(tx, userId, localBindings)
  146. }
  147. func (s *Service) ensureRemoteConverted(item *model.UserMigrationItem, syncedQuota int) error {
  148. remoteUser, err := s.client.GetUser(item.SourceUserId)
  149. if err != nil {
  150. return err
  151. }
  152. if remoteUser.Source != common.UserSourceSynced || remoteUser.RemoteUserId != item.TargetUserId {
  153. if err := s.client.ConvertToSynced(item.SourceUserId, &ConvertRemoteUserRequest{
  154. RemoteUserId: item.TargetUserId,
  155. SyncedQuota: syncedQuota,
  156. }); err != nil {
  157. return err
  158. }
  159. }
  160. return markItemMigrated(item.Id)
  161. }
  162. func markItemMigrated(itemID int) error {
  163. return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{
  164. "status": model.UserMigrationItemStatusMigrated,
  165. "migrated_at": timeNow(),
  166. "error_message": "",
  167. }).Error
  168. }
  169. func isUniqueConstraintErr(err error) bool {
  170. if err == nil {
  171. return false
  172. }
  173. message := err.Error()
  174. return strings.Contains(message, "UNIQUE constraint failed") ||
  175. strings.Contains(message, "Duplicate entry") ||
  176. strings.Contains(strings.ToLower(message), "duplicate key value")
  177. }
  178. func timeNow() int64 {
  179. return common.GetTimestamp()
  180. }