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

353 строки
11 KiB

  1. package controller
  2. import (
  3. "errors"
  4. "net/http"
  5. "strconv"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/model"
  8. user_migration "github.com/QuantumNous/new-api/service/user_migration"
  9. "github.com/QuantumNous/new-api/setting/system_setting"
  10. "github.com/gin-gonic/gin"
  11. "gorm.io/gorm"
  12. )
  13. type migrationVerifier interface {
  14. VerifyItem(item *model.UserMigrationItem) (*user_migration.VerifyResult, error)
  15. }
  16. type migrationCandidateClient interface {
  17. ListUsers(page, pageSize int, keyword string) (*user_migration.ListRemoteUsersResponse, error)
  18. }
  19. var newMigrationVerifier = func() migrationVerifier {
  20. return newMigrationService()
  21. }
  22. var newMigrationCandidateClient = func() migrationCandidateClient {
  23. settings := system_setting.GetRegionSyncSettings()
  24. endpoint := settings.MasterEndpoint
  25. if settings.IsMaster && len(settings.SlaveEndpoints) > 0 {
  26. endpoint = settings.SlaveEndpoints[0]
  27. }
  28. return user_migration.NewClient(endpoint, settings.SyncApiKey)
  29. }
  30. func CreateUserMigrationBatch(c *gin.Context) {
  31. var req struct {
  32. Name string `json:"name" binding:"required"`
  33. SourceRegion string `json:"source_region" binding:"required"`
  34. TargetRegion string `json:"target_region" binding:"required"`
  35. SelectionMode string `json:"selection_mode"`
  36. SourceUserIDs []int `json:"source_user_ids"`
  37. }
  38. if err := c.ShouldBindJSON(&req); err != nil {
  39. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()})
  40. return
  41. }
  42. operatorID, _ := c.Get("id")
  43. if req.SelectionMode == "" {
  44. req.SelectionMode = user_migration.SelectionModeAll
  45. }
  46. if req.SelectionMode != user_migration.SelectionModeAll && req.SelectionMode != user_migration.SelectionModeExplicitIDs {
  47. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": "invalid selection_mode"})
  48. return
  49. }
  50. payload := user_migration.SelectionPayload{}
  51. if req.SelectionMode == user_migration.SelectionModeExplicitIDs {
  52. if len(req.SourceUserIDs) == 0 {
  53. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": "source_user_ids is required"})
  54. return
  55. }
  56. seen := make(map[int]struct{}, len(req.SourceUserIDs))
  57. payload.SourceUserIDs = make([]int, 0, len(req.SourceUserIDs))
  58. for _, id := range req.SourceUserIDs {
  59. if id <= 0 {
  60. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": "invalid source_user_ids"})
  61. return
  62. }
  63. if _, ok := seen[id]; ok {
  64. continue
  65. }
  66. seen[id] = struct{}{}
  67. payload.SourceUserIDs = append(payload.SourceUserIDs, id)
  68. }
  69. if len(payload.SourceUserIDs) > 1000 {
  70. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": "too many source_user_ids"})
  71. return
  72. }
  73. conflicts, err := model.FindActiveMigrationUserConflicts(payload.SourceUserIDs)
  74. if err != nil {
  75. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  76. return
  77. }
  78. if len(conflicts) > 0 {
  79. c.JSON(http.StatusConflict, gin.H{
  80. "success": false,
  81. "error": "some selected users already exist in active migration batches",
  82. "conflicts": conflicts,
  83. })
  84. return
  85. }
  86. }
  87. payloadJSON, _ := common.Marshal(payload)
  88. batch := &model.UserMigrationBatch{
  89. Name: req.Name,
  90. SourceRegion: req.SourceRegion,
  91. TargetRegion: req.TargetRegion,
  92. SelectionMode: req.SelectionMode,
  93. SelectionPayload: string(payloadJSON),
  94. RequestedUserCount: len(payload.SourceUserIDs),
  95. OperatorId: operatorID.(int),
  96. }
  97. if err := model.CreateUserMigrationBatch(batch); err != nil {
  98. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  99. return
  100. }
  101. c.JSON(http.StatusOK, gin.H{"success": true, "data": batch})
  102. }
  103. func ListMigrationCandidateUsers(c *gin.Context) {
  104. page, _ := strconv.Atoi(c.DefaultQuery("page", "0"))
  105. pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50"))
  106. keyword := c.Query("keyword")
  107. if pageSize <= 0 || pageSize > 200 {
  108. pageSize = 50
  109. }
  110. resp, err := newMigrationCandidateClient().ListUsers(page, pageSize, keyword)
  111. if err != nil {
  112. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  113. return
  114. }
  115. c.JSON(http.StatusOK, gin.H{
  116. "success": true,
  117. "data": resp.Data,
  118. "total": resp.Total,
  119. })
  120. }
  121. func ListUserMigrationBatches(c *gin.Context) {
  122. var batches []model.UserMigrationBatch
  123. if err := model.DB.Order("id desc").Find(&batches).Error; err != nil {
  124. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  125. return
  126. }
  127. c.JSON(http.StatusOK, gin.H{"success": true, "data": batches})
  128. }
  129. func GetUserMigrationBatch(c *gin.Context) {
  130. batchID, _ := strconv.Atoi(c.Param("id"))
  131. var batch model.UserMigrationBatch
  132. if err := model.DB.First(&batch, batchID).Error; err != nil {
  133. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "batch not found"})
  134. return
  135. }
  136. page, _ := strconv.Atoi(c.DefaultQuery("page", "0"))
  137. pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50"))
  138. if pageSize <= 0 {
  139. pageSize = 50
  140. }
  141. if pageSize > 200 {
  142. pageSize = 200
  143. }
  144. var items []model.UserMigrationItem
  145. var total int64
  146. model.DB.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batchID).Count(&total)
  147. model.DB.Where("batch_id = ?", batchID).Order("id asc").Offset(page * pageSize).Limit(pageSize).Find(&items)
  148. c.JSON(http.StatusOK, gin.H{
  149. "success": true,
  150. "data": batch,
  151. "items": items,
  152. "total": total,
  153. })
  154. }
  155. func ScanUserMigrationBatch(c *gin.Context) {
  156. batchID, _ := strconv.Atoi(c.Param("id"))
  157. locked, err := model.TryLockBatchForScan(batchID)
  158. if err != nil {
  159. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  160. return
  161. }
  162. if !locked {
  163. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already being scanned or not in a scannable state"})
  164. return
  165. }
  166. svc := newMigrationService()
  167. go func() {
  168. if err := svc.ScanBatch(batchID); err != nil {
  169. model.DB.Model(&model.UserMigrationBatch{}).Where("id = ?", batchID).
  170. Update("status", model.UserMigrationBatchStatusFailed)
  171. }
  172. }()
  173. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "scan started"})
  174. }
  175. func ResolveUserMigrationItem(c *gin.Context) {
  176. itemID, _ := strconv.Atoi(c.Param("id"))
  177. var item model.UserMigrationItem
  178. if err := model.DB.Select("id", "batch_id").First(&item, itemID).Error; err != nil {
  179. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "item not found"})
  180. return
  181. }
  182. var req struct {
  183. ResolutionStrategy string `json:"resolution_strategy" binding:"required"`
  184. TargetUserId int `json:"target_user_id"`
  185. }
  186. if err := c.ShouldBindJSON(&req); err != nil {
  187. c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()})
  188. return
  189. }
  190. validStrategies := map[string]bool{
  191. model.UserMigrationStrategyCreateNew: true,
  192. model.UserMigrationStrategyMergeExisting: true,
  193. model.UserMigrationStrategySkip: true,
  194. }
  195. if !validStrategies[req.ResolutionStrategy] {
  196. c.JSON(http.StatusBadRequest, gin.H{
  197. "success": false,
  198. "error": "invalid resolution_strategy: must be create_new, merge_into_existing, or skip",
  199. })
  200. return
  201. }
  202. if req.ResolutionStrategy == model.UserMigrationStrategyMergeExisting && req.TargetUserId == 0 {
  203. c.JSON(http.StatusBadRequest, gin.H{
  204. "success": false,
  205. "error": "merge_into_existing requires target_user_id",
  206. })
  207. return
  208. }
  209. updates := map[string]any{
  210. "resolution_strategy": req.ResolutionStrategy,
  211. "status": model.UserMigrationItemStatusReady,
  212. }
  213. if req.TargetUserId != 0 {
  214. updates["target_user_id"] = req.TargetUserId
  215. }
  216. if err := model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(updates).Error; err != nil {
  217. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  218. return
  219. }
  220. if err := model.RefreshUserMigrationBatchStats(item.BatchId); err != nil {
  221. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  222. return
  223. }
  224. c.JSON(http.StatusOK, gin.H{"success": true})
  225. }
  226. func ExecuteUserMigrationBatch(c *gin.Context) {
  227. batchID, _ := strconv.Atoi(c.Param("id"))
  228. locked, err := model.TryLockBatchForExecution(batchID)
  229. if err != nil {
  230. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  231. return
  232. }
  233. if !locked {
  234. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already running or not in an executable state"})
  235. return
  236. }
  237. svc := newMigrationService()
  238. go func() {
  239. _ = svc.ExecuteBatch(batchID)
  240. }()
  241. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "execution started"})
  242. }
  243. func RetryUserMigrationBatch(c *gin.Context) {
  244. batchID, _ := strconv.Atoi(c.Param("id"))
  245. model.DB.Model(&model.UserMigrationItem{}).
  246. Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusFailed).
  247. Updates(map[string]any{"status": model.UserMigrationItemStatusReady, "error_message": ""})
  248. locked, err := model.TryLockBatchForExecution(batchID)
  249. if err != nil || !locked {
  250. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch cannot be retried now"})
  251. return
  252. }
  253. svc := newMigrationService()
  254. go func() {
  255. _ = svc.ExecuteBatch(batchID)
  256. }()
  257. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "retry started"})
  258. }
  259. func CancelUserMigrationBatch(c *gin.Context) {
  260. batchID, _ := strconv.Atoi(c.Param("id"))
  261. cancelled, err := model.CancelUserMigrationBatch(batchID)
  262. if err != nil {
  263. if errors.Is(err, gorm.ErrRecordNotFound) {
  264. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "batch not found"})
  265. return
  266. }
  267. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  268. return
  269. }
  270. if !cancelled {
  271. c.JSON(http.StatusConflict, gin.H{
  272. "success": false,
  273. "error": "batch cannot be cancelled in current status",
  274. })
  275. return
  276. }
  277. c.JSON(http.StatusOK, gin.H{"success": true, "message": "batch cancelled"})
  278. }
  279. func VerifyUserMigrationBatch(c *gin.Context) {
  280. batchID, _ := strconv.Atoi(c.Param("id"))
  281. var items []model.UserMigrationItem
  282. model.DB.Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusMigrated).Find(&items)
  283. svc := newMigrationVerifier()
  284. results := make([]map[string]any, 0, len(items))
  285. for _, item := range items {
  286. result, err := svc.VerifyItem(&item)
  287. entry := map[string]any{
  288. "source_user_id": item.SourceUserId,
  289. "target_user_id": item.TargetUserId,
  290. }
  291. if err != nil {
  292. entry["error"] = err.Error()
  293. } else {
  294. entry["result"] = result
  295. }
  296. results = append(results, entry)
  297. }
  298. c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
  299. }
  300. func newMigrationService() *user_migration.Service {
  301. settings := system_setting.GetRegionSyncSettings()
  302. endpoint := settings.MasterEndpoint
  303. if settings.IsMaster && len(settings.SlaveEndpoints) > 0 {
  304. endpoint = settings.SlaveEndpoints[0]
  305. }
  306. client := user_migration.NewClient(endpoint, settings.SyncApiKey)
  307. return user_migration.NewService(client)
  308. }