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

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