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

234 строки
7.0 KiB

  1. package controller
  2. import (
  3. "net/http"
  4. "strconv"
  5. "github.com/QuantumNous/new-api/model"
  6. user_migration "github.com/QuantumNous/new-api/service/user_migration"
  7. "github.com/QuantumNous/new-api/setting/system_setting"
  8. "github.com/gin-gonic/gin"
  9. )
  10. func CreateUserMigrationBatch(c *gin.Context) {
  11. var req struct {
  12. Name string `json:"name" binding:"required"`
  13. SourceRegion string `json:"source_region" binding:"required"`
  14. TargetRegion string `json:"target_region" binding:"required"`
  15. }
  16. if err := c.ShouldBindJSON(&req); err != nil {
  17. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()})
  18. return
  19. }
  20. operatorID, _ := c.Get("id")
  21. batch := &model.UserMigrationBatch{
  22. Name: req.Name,
  23. SourceRegion: req.SourceRegion,
  24. TargetRegion: req.TargetRegion,
  25. OperatorId: operatorID.(int),
  26. }
  27. if err := model.CreateUserMigrationBatch(batch); err != nil {
  28. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  29. return
  30. }
  31. c.JSON(http.StatusOK, gin.H{"success": true, "data": batch})
  32. }
  33. func ListUserMigrationBatches(c *gin.Context) {
  34. var batches []model.UserMigrationBatch
  35. if err := model.DB.Order("id desc").Find(&batches).Error; err != nil {
  36. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  37. return
  38. }
  39. c.JSON(http.StatusOK, gin.H{"success": true, "data": batches})
  40. }
  41. func GetUserMigrationBatch(c *gin.Context) {
  42. batchID, _ := strconv.Atoi(c.Param("id"))
  43. var batch model.UserMigrationBatch
  44. if err := model.DB.First(&batch, batchID).Error; err != nil {
  45. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "batch not found"})
  46. return
  47. }
  48. page, _ := strconv.Atoi(c.DefaultQuery("page", "0"))
  49. pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50"))
  50. if pageSize <= 0 {
  51. pageSize = 50
  52. }
  53. if pageSize > 200 {
  54. pageSize = 200
  55. }
  56. var items []model.UserMigrationItem
  57. var total int64
  58. model.DB.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batchID).Count(&total)
  59. model.DB.Where("batch_id = ?", batchID).Order("id asc").Offset(page * pageSize).Limit(pageSize).Find(&items)
  60. c.JSON(http.StatusOK, gin.H{
  61. "success": true,
  62. "data": batch,
  63. "items": items,
  64. "total": total,
  65. })
  66. }
  67. func ScanUserMigrationBatch(c *gin.Context) {
  68. batchID, _ := strconv.Atoi(c.Param("id"))
  69. locked, err := model.TryLockBatchForScan(batchID)
  70. if err != nil {
  71. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  72. return
  73. }
  74. if !locked {
  75. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already being scanned or not in a scannable state"})
  76. return
  77. }
  78. svc := newMigrationService()
  79. go func() {
  80. if err := svc.ScanBatch(batchID); err != nil {
  81. model.DB.Model(&model.UserMigrationBatch{}).Where("id = ?", batchID).
  82. Update("status", model.UserMigrationBatchStatusFailed)
  83. }
  84. }()
  85. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "scan started"})
  86. }
  87. func ResolveUserMigrationItem(c *gin.Context) {
  88. itemID, _ := strconv.Atoi(c.Param("id"))
  89. var item model.UserMigrationItem
  90. if err := model.DB.Select("id", "batch_id").First(&item, itemID).Error; err != nil {
  91. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "item not found"})
  92. return
  93. }
  94. var req struct {
  95. ResolutionStrategy string `json:"resolution_strategy" binding:"required"`
  96. TargetUserId int `json:"target_user_id"`
  97. }
  98. if err := c.ShouldBindJSON(&req); err != nil {
  99. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()})
  100. return
  101. }
  102. validStrategies := map[string]bool{
  103. model.UserMigrationStrategyCreateNew: true,
  104. model.UserMigrationStrategyMergeExisting: true,
  105. model.UserMigrationStrategySkip: true,
  106. }
  107. if !validStrategies[req.ResolutionStrategy] {
  108. c.JSON(http.StatusBadRequest, gin.H{
  109. "success": false,
  110. "error": "invalid resolution_strategy: must be create_new, merge_into_existing, or skip",
  111. })
  112. return
  113. }
  114. if req.ResolutionStrategy == model.UserMigrationStrategyMergeExisting && req.TargetUserId == 0 {
  115. c.JSON(http.StatusBadRequest, gin.H{
  116. "success": false,
  117. "error": "merge_into_existing requires target_user_id",
  118. })
  119. return
  120. }
  121. updates := map[string]any{
  122. "resolution_strategy": req.ResolutionStrategy,
  123. "status": model.UserMigrationItemStatusReady,
  124. }
  125. if req.TargetUserId != 0 {
  126. updates["target_user_id"] = req.TargetUserId
  127. }
  128. if err := model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(updates).Error; err != nil {
  129. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  130. return
  131. }
  132. if err := model.RefreshUserMigrationBatchStats(item.BatchId); err != nil {
  133. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  134. return
  135. }
  136. c.JSON(http.StatusOK, gin.H{"success": true})
  137. }
  138. func ExecuteUserMigrationBatch(c *gin.Context) {
  139. batchID, _ := strconv.Atoi(c.Param("id"))
  140. locked, err := model.TryLockBatchForExecution(batchID)
  141. if err != nil {
  142. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  143. return
  144. }
  145. if !locked {
  146. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already running or not in an executable state"})
  147. return
  148. }
  149. svc := newMigrationService()
  150. go func() {
  151. _ = svc.ExecuteBatch(batchID)
  152. }()
  153. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "execution started"})
  154. }
  155. func RetryUserMigrationBatch(c *gin.Context) {
  156. batchID, _ := strconv.Atoi(c.Param("id"))
  157. model.DB.Model(&model.UserMigrationItem{}).
  158. Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusFailed).
  159. Updates(map[string]any{"status": model.UserMigrationItemStatusReady, "error_message": ""})
  160. locked, err := model.TryLockBatchForExecution(batchID)
  161. if err != nil || !locked {
  162. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch cannot be retried now"})
  163. return
  164. }
  165. svc := newMigrationService()
  166. go func() {
  167. _ = svc.ExecuteBatch(batchID)
  168. }()
  169. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "retry started"})
  170. }
  171. func VerifyUserMigrationBatch(c *gin.Context) {
  172. batchID, _ := strconv.Atoi(c.Param("id"))
  173. var items []model.UserMigrationItem
  174. model.DB.Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusMigrated).Find(&items)
  175. svc := newMigrationService()
  176. results := make([]map[string]any, 0, len(items))
  177. for _, item := range items {
  178. result, err := svc.VerifyItem(&item)
  179. entry := map[string]any{
  180. "source_user_id": item.SourceUserId,
  181. "target_user_id": item.TargetUserId,
  182. }
  183. if err != nil {
  184. entry["error"] = err.Error()
  185. } else {
  186. entry["result"] = result
  187. }
  188. results = append(results, entry)
  189. }
  190. c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
  191. }
  192. func newMigrationService() *user_migration.Service {
  193. settings := system_setting.GetRegionSyncSettings()
  194. endpoint := settings.MasterEndpoint
  195. if settings.IsMaster && len(settings.SlaveEndpoints) > 0 {
  196. endpoint = settings.SlaveEndpoints[0]
  197. }
  198. client := user_migration.NewClient(endpoint, settings.SyncApiKey)
  199. return user_migration.NewService(client)
  200. }