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

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