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.
 
 
 

378 rivejä
13 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "slices"
  5. "github.com/QuantumNous/new-api/common"
  6. "gorm.io/gorm"
  7. "gorm.io/gorm/clause"
  8. )
  9. const (
  10. UserMigrationBatchStatusDraft = "draft"
  11. UserMigrationBatchStatusScanning = "scanning"
  12. UserMigrationBatchStatusScanned = "scanned"
  13. UserMigrationBatchStatusReady = "ready"
  14. UserMigrationBatchStatusRunning = "running"
  15. UserMigrationBatchStatusCompleted = "completed"
  16. UserMigrationBatchStatusFailed = "failed"
  17. UserMigrationBatchStatusCancelled = "cancelled"
  18. UserMigrationItemStatusPendingScan = "pending_scan"
  19. UserMigrationItemStatusConflict = "conflict"
  20. UserMigrationItemStatusReady = "ready"
  21. UserMigrationItemStatusMigrating = "migrating"
  22. UserMigrationItemStatusMigrated = "migrated"
  23. UserMigrationItemStatusSkipped = "skipped"
  24. UserMigrationItemStatusFailed = "failed"
  25. UserMigrationStrategyCreateNew = "create_new"
  26. UserMigrationStrategyMergeExisting = "merge_into_existing"
  27. UserMigrationStrategySkip = "skip"
  28. )
  29. type UserMigrationBatch struct {
  30. Id int `json:"id" gorm:"primaryKey"`
  31. Name string `json:"name" gorm:"type:varchar(128);not null"`
  32. SourceRegion string `json:"source_region" gorm:"type:varchar(32);not null"`
  33. TargetRegion string `json:"target_region" gorm:"type:varchar(32);not null"`
  34. SelectionMode string `json:"selection_mode" gorm:"type:varchar(32);not null;default:'all'"`
  35. SelectionPayload string `json:"selection_payload" gorm:"type:text"`
  36. RequestedUserCount int `json:"requested_user_count" gorm:"default:0"`
  37. MatchedUserCount int `json:"matched_user_count" gorm:"default:0"`
  38. ExcludedUserCount int `json:"excluded_user_count" gorm:"default:0"`
  39. SelectionSummary string `json:"selection_summary" gorm:"type:text"`
  40. Status string `json:"status" gorm:"type:varchar(32);not null;default:'draft';index"`
  41. TotalUsers int `json:"total_users" gorm:"default:0"`
  42. ConflictUsers int `json:"conflict_users" gorm:"default:0"`
  43. MigratedUsers int `json:"migrated_users" gorm:"default:0"`
  44. FailedUsers int `json:"failed_users" gorm:"default:0"`
  45. OperatorId int `json:"operator_id" gorm:"index"`
  46. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  47. UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
  48. }
  49. type UserMigrationItem struct {
  50. Id int `json:"id" gorm:"primaryKey"`
  51. BatchId int `json:"batch_id" gorm:"index:idx_batch_source_user,unique"`
  52. SourceUserId int `json:"source_user_id" gorm:"index:idx_batch_source_user,unique"`
  53. TargetUserId int `json:"target_user_id" gorm:"default:0;index"`
  54. SourceUsername string `json:"source_username" gorm:"type:varchar(128);not null"`
  55. SourceEmail string `json:"source_email" gorm:"type:varchar(255);default:''"`
  56. SourceQuota int `json:"source_quota" gorm:"default:0"`
  57. ConflictFlags string `json:"conflict_flags" gorm:"type:text"`
  58. ResolutionStrategy string `json:"resolution_strategy" gorm:"type:varchar(32);default:''"`
  59. Status string `json:"status" gorm:"type:varchar(32);not null;default:'pending_scan';index"`
  60. ErrorMessage string `json:"error_message" gorm:"type:text"`
  61. Snapshot string `json:"snapshot" gorm:"type:text"`
  62. MigrationToken string `json:"migration_token" gorm:"type:varchar(64);default:'';index"`
  63. MigratedAt int64 `json:"migrated_at" gorm:"bigint;default:0"`
  64. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  65. UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
  66. }
  67. type MigrationQuotaGrant struct {
  68. Id int `json:"id" gorm:"primaryKey"`
  69. MigrationToken string `json:"migration_token" gorm:"type:varchar(64);not null;uniqueIndex"`
  70. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  71. }
  72. type ActiveMigrationUserConflict struct {
  73. SourceUserID int `json:"source_user_id"`
  74. BatchID int `json:"batch_id"`
  75. BatchStatus string `json:"batch_status"`
  76. }
  77. func (b *UserMigrationBatch) BeforeCreate(tx *gorm.DB) error {
  78. now := common.GetTimestamp()
  79. if b.Status == "" {
  80. b.Status = UserMigrationBatchStatusDraft
  81. }
  82. if b.SelectionMode == "" {
  83. b.SelectionMode = "all"
  84. }
  85. b.CreatedAt = now
  86. b.UpdatedAt = now
  87. return nil
  88. }
  89. func (b *UserMigrationBatch) BeforeUpdate(tx *gorm.DB) error {
  90. b.UpdatedAt = common.GetTimestamp()
  91. return nil
  92. }
  93. func (i *UserMigrationItem) BeforeCreate(tx *gorm.DB) error {
  94. now := common.GetTimestamp()
  95. if i.Status == "" {
  96. i.Status = UserMigrationItemStatusPendingScan
  97. }
  98. i.CreatedAt = now
  99. i.UpdatedAt = now
  100. return nil
  101. }
  102. func (i *UserMigrationItem) BeforeUpdate(tx *gorm.DB) error {
  103. i.UpdatedAt = common.GetTimestamp()
  104. return nil
  105. }
  106. func (g *MigrationQuotaGrant) BeforeCreate(tx *gorm.DB) error {
  107. g.CreatedAt = common.GetTimestamp()
  108. return nil
  109. }
  110. func CreateUserMigrationBatch(batch *UserMigrationBatch) error {
  111. return DB.Create(batch).Error
  112. }
  113. func CancelUserMigrationBatch(batchId int) (bool, error) {
  114. var batch UserMigrationBatch
  115. if err := DB.First(&batch, batchId).Error; err != nil {
  116. return false, err
  117. }
  118. switch batch.Status {
  119. case UserMigrationBatchStatusCancelled:
  120. return true, nil
  121. case UserMigrationBatchStatusDraft,
  122. UserMigrationBatchStatusScanned,
  123. UserMigrationBatchStatusReady,
  124. UserMigrationBatchStatusFailed:
  125. return true, DB.Model(&UserMigrationBatch{}).
  126. Where("id = ?", batchId).
  127. Update("status", UserMigrationBatchStatusCancelled).Error
  128. default:
  129. return false, nil
  130. }
  131. }
  132. func FindActiveMigrationUserConflicts(sourceUserIDs []int) ([]ActiveMigrationUserConflict, error) {
  133. if len(sourceUserIDs) == 0 {
  134. return nil, nil
  135. }
  136. activeStatuses := []string{
  137. UserMigrationBatchStatusDraft,
  138. UserMigrationBatchStatusScanning,
  139. UserMigrationBatchStatusScanned,
  140. UserMigrationBatchStatusReady,
  141. UserMigrationBatchStatusRunning,
  142. UserMigrationBatchStatusFailed,
  143. }
  144. type activeBatchSelection struct {
  145. ID int
  146. Status string
  147. SelectionMode string
  148. SelectionPayload string
  149. }
  150. conflictMap := make(map[string]ActiveMigrationUserConflict)
  151. addConflict := func(conflict ActiveMigrationUserConflict) {
  152. key := fmt.Sprintf("%d:%d", conflict.SourceUserID, conflict.BatchID)
  153. if _, exists := conflictMap[key]; exists {
  154. return
  155. }
  156. conflictMap[key] = conflict
  157. }
  158. var itemConflicts []ActiveMigrationUserConflict
  159. err := DB.Table("user_migration_items AS umi").
  160. Select("umi.source_user_id, umb.id AS batch_id, umb.status AS batch_status").
  161. Joins("JOIN user_migration_batches AS umb ON umb.id = umi.batch_id").
  162. Where("umi.source_user_id IN ?", sourceUserIDs).
  163. Where("umb.status IN ?", activeStatuses).
  164. Order("umi.source_user_id ASC, umb.id ASC").
  165. Scan(&itemConflicts).Error
  166. if err != nil {
  167. return nil, err
  168. }
  169. for _, conflict := range itemConflicts {
  170. addConflict(conflict)
  171. }
  172. var activeBatches []activeBatchSelection
  173. if err := DB.Model(&UserMigrationBatch{}).
  174. Select("id, status, selection_mode, selection_payload").
  175. Where("status IN ?", activeStatuses).
  176. Where("selection_mode = ?", "explicit_ids").
  177. Find(&activeBatches).Error; err != nil {
  178. return nil, err
  179. }
  180. sourceIDSet := make(map[int]struct{}, len(sourceUserIDs))
  181. for _, id := range sourceUserIDs {
  182. sourceIDSet[id] = struct{}{}
  183. }
  184. for _, batch := range activeBatches {
  185. if batch.SelectionPayload == "" {
  186. continue
  187. }
  188. var payload struct {
  189. SourceUserIDs []int `json:"source_user_ids"`
  190. }
  191. if err := common.UnmarshalJsonStr(batch.SelectionPayload, &payload); err != nil {
  192. return nil, err
  193. }
  194. for _, sourceUserID := range payload.SourceUserIDs {
  195. if _, ok := sourceIDSet[sourceUserID]; !ok {
  196. continue
  197. }
  198. addConflict(ActiveMigrationUserConflict{
  199. SourceUserID: sourceUserID,
  200. BatchID: batch.ID,
  201. BatchStatus: batch.Status,
  202. })
  203. }
  204. }
  205. conflicts := make([]ActiveMigrationUserConflict, 0, len(conflictMap))
  206. for _, conflict := range conflictMap {
  207. conflicts = append(conflicts, conflict)
  208. }
  209. slices.SortFunc(conflicts, func(a, b ActiveMigrationUserConflict) int {
  210. if a.SourceUserID != b.SourceUserID {
  211. return a.SourceUserID - b.SourceUserID
  212. }
  213. return a.BatchID - b.BatchID
  214. })
  215. return conflicts, nil
  216. }
  217. func UpsertUserMigrationItem(item *UserMigrationItem) error {
  218. var existing UserMigrationItem
  219. result := DB.Where("batch_id = ? AND source_user_id = ?", item.BatchId, item.SourceUserId).Limit(1).Find(&existing)
  220. if result.Error != nil {
  221. return result.Error
  222. }
  223. if result.RowsAffected > 0 {
  224. updates := map[string]any{
  225. "source_username": item.SourceUsername,
  226. "source_email": item.SourceEmail,
  227. "source_quota": item.SourceQuota,
  228. "conflict_flags": item.ConflictFlags,
  229. "status": item.Status,
  230. "snapshot": item.Snapshot,
  231. }
  232. if existing.ResolutionStrategy == "" && item.ResolutionStrategy != "" {
  233. updates["resolution_strategy"] = item.ResolutionStrategy
  234. }
  235. if existing.TargetUserId == 0 && item.TargetUserId != 0 {
  236. updates["target_user_id"] = item.TargetUserId
  237. }
  238. item.Id = existing.Id
  239. return DB.Model(&existing).Updates(updates).Error
  240. }
  241. return DB.Create(item).Error
  242. }
  243. func ListUserMigrationItemsByBatch(batchId int) ([]*UserMigrationItem, error) {
  244. var items []*UserMigrationItem
  245. err := DB.Where("batch_id = ?", batchId).Order("id asc").Find(&items).Error
  246. return items, err
  247. }
  248. func ListExecutableUserMigrationItems(batchId int) ([]*UserMigrationItem, error) {
  249. var items []*UserMigrationItem
  250. err := DB.Where("batch_id = ? AND status IN ?", batchId, []string{
  251. UserMigrationItemStatusReady,
  252. UserMigrationItemStatusMigrating,
  253. UserMigrationItemStatusFailed,
  254. }).Order("id asc").Find(&items).Error
  255. return items, err
  256. }
  257. func MarkUserMigrationItemFailed(id int, errMsg string) error {
  258. return DB.Model(&UserMigrationItem{}).Where("id = ?", id).Updates(map[string]any{
  259. "status": UserMigrationItemStatusFailed,
  260. "error_message": errMsg,
  261. }).Error
  262. }
  263. func MarkUserMigrationItemSkipped(id int) error {
  264. return DB.Model(&UserMigrationItem{}).Where("id = ?", id).
  265. Update("status", UserMigrationItemStatusSkipped).Error
  266. }
  267. func RefreshUserMigrationBatchStats(batchId int) error {
  268. var total, conflicts, pendingScan int64
  269. DB.Model(&UserMigrationItem{}).Where("batch_id = ?", batchId).Count(&total)
  270. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusConflict).Count(&conflicts)
  271. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusPendingScan).Count(&pendingScan)
  272. status := UserMigrationBatchStatusScanned
  273. if conflicts == 0 && pendingScan == 0 && total > 0 {
  274. status = UserMigrationBatchStatusReady
  275. }
  276. return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
  277. "total_users": int(total),
  278. "conflict_users": int(conflicts),
  279. "status": status,
  280. }).Error
  281. }
  282. func RefreshUserMigrationBatchExecutionStats(batchId int) error {
  283. var migrated, failed, pending, migrating int64
  284. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrated).Count(&migrated)
  285. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusFailed).Count(&failed)
  286. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status IN ?", batchId, []string{
  287. UserMigrationItemStatusReady,
  288. UserMigrationItemStatusMigrating,
  289. }).Count(&pending)
  290. DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrating).Count(&migrating)
  291. status := UserMigrationBatchStatusCompleted
  292. if failed > 0 {
  293. status = UserMigrationBatchStatusFailed
  294. }
  295. if pending > 0 {
  296. status = UserMigrationBatchStatusRunning
  297. }
  298. if pending == migrating && migrating > 0 {
  299. status = UserMigrationBatchStatusFailed
  300. }
  301. return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
  302. "migrated_users": int(migrated),
  303. "failed_users": int(failed),
  304. "status": status,
  305. }).Error
  306. }
  307. func TryLockBatchForExecution(batchId int) (bool, error) {
  308. result := DB.Model(&UserMigrationBatch{}).
  309. Where("id = ? AND status IN ?", batchId, []string{
  310. UserMigrationBatchStatusReady,
  311. UserMigrationBatchStatusFailed,
  312. }).
  313. Update("status", UserMigrationBatchStatusRunning)
  314. return result.RowsAffected == 1, result.Error
  315. }
  316. func TryLockBatchForScan(batchId int) (bool, error) {
  317. result := DB.Model(&UserMigrationBatch{}).
  318. Where("id = ? AND status IN ?", batchId, []string{
  319. UserMigrationBatchStatusDraft,
  320. UserMigrationBatchStatusScanned,
  321. UserMigrationBatchStatusReady,
  322. }).
  323. Update("status", UserMigrationBatchStatusScanning)
  324. return result.RowsAffected == 1, result.Error
  325. }
  326. func AtomicGrantQuotaIfNotExists(tx *gorm.DB, userId, delta int, token string) (bool, error) {
  327. result := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&MigrationQuotaGrant{MigrationToken: token})
  328. if result.Error != nil {
  329. return false, result.Error
  330. }
  331. if result.RowsAffected == 0 {
  332. return true, nil
  333. }
  334. return false, tx.Model(&User{}).Where("id = ?", userId).
  335. Update("quota", gorm.Expr("quota + ?", delta)).Error
  336. }