Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

377 Zeilen
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", "conflict_flags").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. if req.ResolutionStrategy == model.UserMigrationStrategyMergeExisting &&
  210. hasMigrationConflictFlag(item.ConflictFlags, "cn_already_synced_to_ov") {
  211. c.JSON(http.StatusConflict, gin.H{
  212. "success": false,
  213. "error": "merge_into_existing is not allowed when target CN user already has a synced OV copy",
  214. })
  215. return
  216. }
  217. updates := map[string]any{
  218. "resolution_strategy": req.ResolutionStrategy,
  219. "status": model.UserMigrationItemStatusReady,
  220. }
  221. if req.TargetUserId != 0 {
  222. updates["target_user_id"] = req.TargetUserId
  223. }
  224. if err := model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(updates).Error; err != nil {
  225. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  226. return
  227. }
  228. if err := model.RefreshUserMigrationBatchStats(item.BatchId); err != nil {
  229. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  230. return
  231. }
  232. c.JSON(http.StatusOK, gin.H{"success": true})
  233. }
  234. func hasMigrationConflictFlag(flagsJSON string, want string) bool {
  235. if flagsJSON == "" || want == "" {
  236. return false
  237. }
  238. var flags []string
  239. if err := common.UnmarshalJsonStr(flagsJSON, &flags); err != nil {
  240. return false
  241. }
  242. for _, flag := range flags {
  243. if flag == want {
  244. return true
  245. }
  246. }
  247. return false
  248. }
  249. func ExecuteUserMigrationBatch(c *gin.Context) {
  250. batchID, _ := strconv.Atoi(c.Param("id"))
  251. locked, err := model.TryLockBatchForExecution(batchID)
  252. if err != nil {
  253. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  254. return
  255. }
  256. if !locked {
  257. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already running or not in an executable state"})
  258. return
  259. }
  260. svc := newMigrationService()
  261. go func() {
  262. _ = svc.ExecuteBatch(batchID)
  263. }()
  264. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "execution started"})
  265. }
  266. func RetryUserMigrationBatch(c *gin.Context) {
  267. batchID, _ := strconv.Atoi(c.Param("id"))
  268. model.DB.Model(&model.UserMigrationItem{}).
  269. Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusFailed).
  270. Updates(map[string]any{"status": model.UserMigrationItemStatusReady, "error_message": ""})
  271. locked, err := model.TryLockBatchForExecution(batchID)
  272. if err != nil || !locked {
  273. c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch cannot be retried now"})
  274. return
  275. }
  276. svc := newMigrationService()
  277. go func() {
  278. _ = svc.ExecuteBatch(batchID)
  279. }()
  280. c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "retry started"})
  281. }
  282. func CancelUserMigrationBatch(c *gin.Context) {
  283. batchID, _ := strconv.Atoi(c.Param("id"))
  284. cancelled, err := model.CancelUserMigrationBatch(batchID)
  285. if err != nil {
  286. if errors.Is(err, gorm.ErrRecordNotFound) {
  287. c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "batch not found"})
  288. return
  289. }
  290. c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()})
  291. return
  292. }
  293. if !cancelled {
  294. c.JSON(http.StatusConflict, gin.H{
  295. "success": false,
  296. "error": "batch cannot be cancelled in current status",
  297. })
  298. return
  299. }
  300. c.JSON(http.StatusOK, gin.H{"success": true, "message": "batch cancelled"})
  301. }
  302. func VerifyUserMigrationBatch(c *gin.Context) {
  303. batchID, _ := strconv.Atoi(c.Param("id"))
  304. var items []model.UserMigrationItem
  305. model.DB.Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusMigrated).Find(&items)
  306. svc := newMigrationVerifier()
  307. results := make([]map[string]any, 0, len(items))
  308. for _, item := range items {
  309. result, err := svc.VerifyItem(&item)
  310. entry := map[string]any{
  311. "source_user_id": item.SourceUserId,
  312. "target_user_id": item.TargetUserId,
  313. }
  314. if err != nil {
  315. entry["error"] = err.Error()
  316. } else {
  317. entry["result"] = result
  318. }
  319. results = append(results, entry)
  320. }
  321. c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
  322. }
  323. func newMigrationService() *user_migration.Service {
  324. settings := system_setting.GetRegionSyncSettings()
  325. endpoint := settings.MasterEndpoint
  326. if settings.IsMaster && len(settings.SlaveEndpoints) > 0 {
  327. endpoint = settings.SlaveEndpoints[0]
  328. }
  329. client := user_migration.NewClient(endpoint, settings.SyncApiKey)
  330. return user_migration.NewService(client)
  331. }