Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

240 строки
8.6 KiB

  1. package model
  2. import (
  3. "github.com/QuantumNous/new-api/common"
  4. "gorm.io/gorm"
  5. "gorm.io/gorm/clause"
  6. )
  7. const (
  8. UserMigrationBatchStatusDraft = "draft"
  9. UserMigrationBatchStatusScanned = "scanned"
  10. UserMigrationBatchStatusReady = "ready"
  11. UserMigrationBatchStatusRunning = "running"
  12. UserMigrationBatchStatusCompleted = "completed"
  13. UserMigrationBatchStatusFailed = "failed"
  14. UserMigrationBatchStatusCancelled = "cancelled"
  15. UserMigrationItemStatusPendingScan = "pending_scan"
  16. UserMigrationItemStatusConflict = "conflict"
  17. UserMigrationItemStatusReady = "ready"
  18. UserMigrationItemStatusMigrating = "migrating"
  19. UserMigrationItemStatusMigrated = "migrated"
  20. UserMigrationItemStatusSkipped = "skipped"
  21. UserMigrationItemStatusFailed = "failed"
  22. UserMigrationStrategyCreateNew = "create_new"
  23. UserMigrationStrategyMergeExisting = "merge_into_existing"
  24. UserMigrationStrategySkip = "skip"
  25. )
  26. type UserMigrationBatch struct {
  27. Id int `json:"id" gorm:"primaryKey"`
  28. Name string `json:"name" gorm:"type:varchar(128);not null"`
  29. SourceRegion string `json:"source_region" gorm:"type:varchar(32);not null"`
  30. TargetRegion string `json:"target_region" gorm:"type:varchar(32);not null"`
  31. Status string `json:"status" gorm:"type:varchar(32);not null;default:'draft';index"`
  32. TotalUsers int `json:"total_users" gorm:"default:0"`
  33. ConflictUsers int `json:"conflict_users" gorm:"default:0"`
  34. MigratedUsers int `json:"migrated_users" gorm:"default:0"`
  35. FailedUsers int `json:"failed_users" gorm:"default:0"`
  36. OperatorId int `json:"operator_id" gorm:"index"`
  37. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  38. UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
  39. }
  40. type UserMigrationItem struct {
  41. Id int `json:"id" gorm:"primaryKey"`
  42. BatchId int `json:"batch_id" gorm:"index:idx_batch_source_user,unique"`
  43. SourceUserId int `json:"source_user_id" gorm:"index:idx_batch_source_user,unique"`
  44. TargetUserId int `json:"target_user_id" gorm:"default:0;index"`
  45. SourceUsername string `json:"source_username" gorm:"type:varchar(128);not null"`
  46. SourceEmail string `json:"source_email" gorm:"type:varchar(255);default:''"`
  47. SourceQuota int `json:"source_quota" gorm:"default:0"`
  48. ConflictFlags string `json:"conflict_flags" gorm:"type:text"`
  49. ResolutionStrategy string `json:"resolution_strategy" gorm:"type:varchar(32);default:''"`
  50. Status string `json:"status" gorm:"type:varchar(32);not null;default:'pending_scan';index"`
  51. ErrorMessage string `json:"error_message" gorm:"type:text"`
  52. Snapshot string `json:"snapshot" gorm:"type:text"`
  53. MigrationToken string `json:"migration_token" gorm:"type:varchar(64);default:'';index"`
  54. MigratedAt int64 `json:"migrated_at" gorm:"bigint;default:0"`
  55. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  56. UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
  57. }
  58. type MigrationQuotaGrant struct {
  59. Id int `json:"id" gorm:"primaryKey"`
  60. MigrationToken string `json:"migration_token" gorm:"type:varchar(64);not null;uniqueIndex"`
  61. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  62. }
  63. func (b *UserMigrationBatch) BeforeCreate(tx *gorm.DB) error {
  64. now := common.GetTimestamp()
  65. if b.Status == "" {
  66. b.Status = UserMigrationBatchStatusDraft
  67. }
  68. b.CreatedAt = now
  69. b.UpdatedAt = now
  70. return nil
  71. }
  72. func (b *UserMigrationBatch) BeforeUpdate(tx *gorm.DB) error {
  73. b.UpdatedAt = common.GetTimestamp()
  74. return nil
  75. }
  76. func (i *UserMigrationItem) BeforeCreate(tx *gorm.DB) error {
  77. now := common.GetTimestamp()
  78. if i.Status == "" {
  79. i.Status = UserMigrationItemStatusPendingScan
  80. }
  81. i.CreatedAt = now
  82. i.UpdatedAt = now
  83. return nil
  84. }
  85. func (i *UserMigrationItem) BeforeUpdate(tx *gorm.DB) error {
  86. i.UpdatedAt = common.GetTimestamp()
  87. return nil
  88. }
  89. func (g *MigrationQuotaGrant) BeforeCreate(tx *gorm.DB) error {
  90. g.CreatedAt = common.GetTimestamp()
  91. return nil
  92. }
  93. func CreateUserMigrationBatch(batch *UserMigrationBatch) error {
  94. return DB.Create(batch).Error
  95. }
  96. func UpsertUserMigrationItem(item *UserMigrationItem) error {
  97. var existing UserMigrationItem
  98. result := DB.Where("batch_id = ? AND source_user_id = ?", item.BatchId, item.SourceUserId).Limit(1).Find(&existing)
  99. if result.Error != nil {
  100. return result.Error
  101. }
  102. if result.RowsAffected > 0 {
  103. updates := map[string]any{
  104. "source_username": item.SourceUsername,
  105. "source_email": item.SourceEmail,
  106. "source_quota": item.SourceQuota,
  107. "conflict_flags": item.ConflictFlags,
  108. "status": item.Status,
  109. "snapshot": item.Snapshot,
  110. }
  111. if existing.ResolutionStrategy == "" && item.ResolutionStrategy != "" {
  112. updates["resolution_strategy"] = item.ResolutionStrategy
  113. }
  114. if existing.TargetUserId == 0 && item.TargetUserId != 0 {
  115. updates["target_user_id"] = item.TargetUserId
  116. }
  117. item.Id = existing.Id
  118. return DB.Model(&existing).Updates(updates).Error
  119. }
  120. return DB.Create(item).Error
  121. }
  122. func ListUserMigrationItemsByBatch(batchId int) ([]*UserMigrationItem, error) {
  123. var items []*UserMigrationItem
  124. err := DB.Where("batch_id = ?", batchId).Order("id asc").Find(&items).Error
  125. return items, err
  126. }
  127. func ListExecutableUserMigrationItems(batchId int) ([]*UserMigrationItem, error) {
  128. var items []*UserMigrationItem
  129. err := DB.Where("batch_id = ? AND status IN ?", batchId, []string{
  130. UserMigrationItemStatusReady,
  131. UserMigrationItemStatusMigrating,
  132. UserMigrationItemStatusFailed,
  133. }).Order("id asc").Find(&items).Error
  134. return items, err
  135. }
  136. func MarkUserMigrationItemFailed(id int, errMsg string) error {
  137. return DB.Model(&UserMigrationItem{}).Where("id = ?", id).Updates(map[string]any{
  138. "status": UserMigrationItemStatusFailed,
  139. "error_message": errMsg,
  140. }).Error
  141. }
  142. func MarkUserMigrationItemSkipped(id int) error {
  143. return DB.Model(&UserMigrationItem{}).Where("id = ?", id).
  144. Update("status", UserMigrationItemStatusSkipped).Error
  145. }
  146. func RefreshUserMigrationBatchStats(batchId int) error {
  147. var total, conflicts, pendingScan int64
  148. DB.Model(&UserMigrationItem{}).Where("batch_id = ?", batchId).Count(&total)
  149. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusConflict).Count(&conflicts)
  150. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusPendingScan).Count(&pendingScan)
  151. status := UserMigrationBatchStatusScanned
  152. if conflicts == 0 && pendingScan == 0 && total > 0 {
  153. status = UserMigrationBatchStatusReady
  154. }
  155. return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
  156. "total_users": int(total),
  157. "conflict_users": int(conflicts),
  158. "status": status,
  159. }).Error
  160. }
  161. func RefreshUserMigrationBatchExecutionStats(batchId int) error {
  162. var migrated, failed, pending int64
  163. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrated).Count(&migrated)
  164. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusFailed).Count(&failed)
  165. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status IN ?", batchId, []string{
  166. UserMigrationItemStatusReady,
  167. UserMigrationItemStatusMigrating,
  168. }).Count(&pending)
  169. status := UserMigrationBatchStatusCompleted
  170. if failed > 0 {
  171. status = UserMigrationBatchStatusFailed
  172. }
  173. if pending > 0 {
  174. status = UserMigrationBatchStatusRunning
  175. }
  176. return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
  177. "migrated_users": int(migrated),
  178. "failed_users": int(failed),
  179. "status": status,
  180. }).Error
  181. }
  182. func TryLockBatchForExecution(batchId int) (bool, error) {
  183. result := DB.Model(&UserMigrationBatch{}).
  184. Where("id = ? AND status IN ?", batchId, []string{
  185. UserMigrationBatchStatusReady,
  186. UserMigrationBatchStatusFailed,
  187. }).
  188. Update("status", UserMigrationBatchStatusRunning)
  189. return result.RowsAffected == 1, result.Error
  190. }
  191. func TryLockBatchForScan(batchId int) (bool, error) {
  192. result := DB.Model(&UserMigrationBatch{}).
  193. Where("id = ? AND status IN ?", batchId, []string{
  194. UserMigrationBatchStatusDraft,
  195. UserMigrationBatchStatusScanned,
  196. UserMigrationBatchStatusReady,
  197. }).
  198. Update("status", UserMigrationBatchStatusScanned)
  199. return result.RowsAffected == 1, result.Error
  200. }
  201. func AtomicGrantQuotaIfNotExists(tx *gorm.DB, userId, delta int, token string) (bool, error) {
  202. result := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&MigrationQuotaGrant{MigrationToken: token})
  203. if result.Error != nil {
  204. return false, result.Error
  205. }
  206. if result.RowsAffected == 0 {
  207. return true, nil
  208. }
  209. return false, tx.Model(&User{}).Where("id = ?", userId).
  210. Update("quota", gorm.Expr("quota + ?", delta)).Error
  211. }