diff --git a/controller/user_migration.go b/controller/user_migration.go new file mode 100644 index 0000000..372981c --- /dev/null +++ b/controller/user_migration.go @@ -0,0 +1,233 @@ +package controller + +import ( + "net/http" + "strconv" + + "github.com/QuantumNous/new-api/model" + user_migration "github.com/QuantumNous/new-api/service/user_migration" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" +) + +func CreateUserMigrationBatch(c *gin.Context) { + var req struct { + Name string `json:"name" binding:"required"` + SourceRegion string `json:"source_region" binding:"required"` + TargetRegion string `json:"target_region" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()}) + return + } + + operatorID, _ := c.Get("id") + batch := &model.UserMigrationBatch{ + Name: req.Name, + SourceRegion: req.SourceRegion, + TargetRegion: req.TargetRegion, + OperatorId: operatorID.(int), + } + if err := model.CreateUserMigrationBatch(batch); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": batch}) +} + +func ListUserMigrationBatches(c *gin.Context) { + var batches []model.UserMigrationBatch + if err := model.DB.Order("id desc").Find(&batches).Error; err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": batches}) +} + +func GetUserMigrationBatch(c *gin.Context) { + batchID, _ := strconv.Atoi(c.Param("id")) + + var batch model.UserMigrationBatch + if err := model.DB.First(&batch, batchID).Error; err != nil { + c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "batch not found"}) + return + } + + page, _ := strconv.Atoi(c.DefaultQuery("page", "0")) + pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50")) + if pageSize <= 0 { + pageSize = 50 + } + if pageSize > 200 { + pageSize = 200 + } + + var items []model.UserMigrationItem + var total int64 + model.DB.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batchID).Count(&total) + model.DB.Where("batch_id = ?", batchID).Order("id asc").Offset(page * pageSize).Limit(pageSize).Find(&items) + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": batch, + "items": items, + "total": total, + }) +} + +func ScanUserMigrationBatch(c *gin.Context) { + batchID, _ := strconv.Atoi(c.Param("id")) + locked, err := model.TryLockBatchForScan(batchID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + if !locked { + c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already being scanned or not in a scannable state"}) + return + } + + svc := newMigrationService() + go func() { + if err := svc.ScanBatch(batchID); err != nil { + model.DB.Model(&model.UserMigrationBatch{}).Where("id = ?", batchID). + Update("status", model.UserMigrationBatchStatusFailed) + } + }() + + c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "scan started"}) +} + +func ResolveUserMigrationItem(c *gin.Context) { + itemID, _ := strconv.Atoi(c.Param("id")) + + var item model.UserMigrationItem + if err := model.DB.Select("id", "batch_id").First(&item, itemID).Error; err != nil { + c.JSON(http.StatusNotFound, gin.H{"success": false, "error": "item not found"}) + return + } + + var req struct { + ResolutionStrategy string `json:"resolution_strategy" binding:"required"` + TargetUserId int `json:"target_user_id"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": err.Error()}) + return + } + + validStrategies := map[string]bool{ + model.UserMigrationStrategyCreateNew: true, + model.UserMigrationStrategyMergeExisting: true, + model.UserMigrationStrategySkip: true, + } + if !validStrategies[req.ResolutionStrategy] { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": "invalid resolution_strategy: must be create_new, merge_into_existing, or skip", + }) + return + } + if req.ResolutionStrategy == model.UserMigrationStrategyMergeExisting && req.TargetUserId == 0 { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "error": "merge_into_existing requires target_user_id", + }) + return + } + + updates := map[string]any{ + "resolution_strategy": req.ResolutionStrategy, + "status": model.UserMigrationItemStatusReady, + } + if req.TargetUserId != 0 { + updates["target_user_id"] = req.TargetUserId + } + if err := model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(updates).Error; err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + if err := model.RefreshUserMigrationBatchStats(item.BatchId); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +func ExecuteUserMigrationBatch(c *gin.Context) { + batchID, _ := strconv.Atoi(c.Param("id")) + locked, err := model.TryLockBatchForExecution(batchID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": err.Error()}) + return + } + if !locked { + c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch is already running or not in an executable state"}) + return + } + + svc := newMigrationService() + go func() { + _ = svc.ExecuteBatch(batchID) + }() + + c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "execution started"}) +} + +func RetryUserMigrationBatch(c *gin.Context) { + batchID, _ := strconv.Atoi(c.Param("id")) + + model.DB.Model(&model.UserMigrationItem{}). + Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusFailed). + Updates(map[string]any{"status": model.UserMigrationItemStatusReady, "error_message": ""}) + + locked, err := model.TryLockBatchForExecution(batchID) + if err != nil || !locked { + c.JSON(http.StatusConflict, gin.H{"success": false, "error": "batch cannot be retried now"}) + return + } + + svc := newMigrationService() + go func() { + _ = svc.ExecuteBatch(batchID) + }() + + c.JSON(http.StatusAccepted, gin.H{"success": true, "message": "retry started"}) +} + +func VerifyUserMigrationBatch(c *gin.Context) { + batchID, _ := strconv.Atoi(c.Param("id")) + + var items []model.UserMigrationItem + model.DB.Where("batch_id = ? AND status = ?", batchID, model.UserMigrationItemStatusMigrated).Find(&items) + + svc := newMigrationService() + results := make([]map[string]any, 0, len(items)) + for _, item := range items { + result, err := svc.VerifyItem(&item) + entry := map[string]any{ + "source_user_id": item.SourceUserId, + "target_user_id": item.TargetUserId, + } + if err != nil { + entry["error"] = err.Error() + } else { + entry["result"] = result + } + results = append(results, entry) + } + + c.JSON(http.StatusOK, gin.H{"success": true, "data": results}) +} + +func newMigrationService() *user_migration.Service { + settings := system_setting.GetRegionSyncSettings() + endpoint := settings.MasterEndpoint + if settings.IsMaster && len(settings.SlaveEndpoints) > 0 { + endpoint = settings.SlaveEndpoints[0] + } + + client := user_migration.NewClient(endpoint, settings.SyncApiKey) + return user_migration.NewService(client) +} diff --git a/controller/user_migration_test.go b/controller/user_migration_test.go new file mode 100644 index 0000000..94ea99c --- /dev/null +++ b/controller/user_migration_test.go @@ -0,0 +1,137 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupUserMigrationControllerDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.UserMigrationBatch{}, + &model.UserMigrationItem{}, + &model.MigrationQuotaGrant{}, + )) + + return db +} + +func testRootAuthMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + c.Set("id", 1) + c.Set("role", 100) + c.Next() + } +} + +func TestCreateUserMigrationBatch_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testRootAuthMiddleware()) + router.POST("/api/user-migrations/batches", CreateUserMigrationBatch) + + db := setupUserMigrationControllerDB(t) + orig := model.DB + model.DB = db + defer func() { + model.DB = orig + }() + + body := map[string]any{ + "name": "wave-root-1", + "source_region": "overseas", + "target_region": "cn", + } + bodyBytes, err := json.Marshal(body) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches", bytes.NewBuffer(bodyBytes)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.True(t, resp["success"].(bool)) + require.NotNil(t, resp["data"]) +} + +func TestExecuteUserMigrationBatch_AlreadyRunning(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testRootAuthMiddleware()) + router.POST("/api/user-migrations/batches/:id/execute", ExecuteUserMigrationBatch) + + db := setupUserMigrationControllerDB(t) + orig := model.DB + model.DB = db + defer func() { + model.DB = orig + }() + + batch := &model.UserMigrationBatch{ + Name: "b", + SourceRegion: "ov", + TargetRegion: "cn", + OperatorId: 1, + Status: model.UserMigrationBatchStatusRunning, + } + require.NoError(t, db.Create(batch).Error) + + req := httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusConflict, w.Code) +} + +func TestResolveUserMigrationItem_InvalidStrategy(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(testRootAuthMiddleware()) + router.POST("/api/user-migrations/items/:id/resolve", ResolveUserMigrationItem) + + db := setupUserMigrationControllerDB(t) + orig := model.DB + model.DB = db + defer func() { + model.DB = orig + }() + + require.NoError(t, db.Create(&model.UserMigrationItem{ + BatchId: 1, + SourceUserId: 10000001, + SourceUsername: "u", + Status: model.UserMigrationItemStatusConflict, + }).Error) + + req := httptest.NewRequest( + http.MethodPost, + "/api/user-migrations/items/1/resolve", + bytes.NewBufferString(`{"resolution_strategy":"invalid_strategy"}`), + ) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) +} diff --git a/router/api-router.go b/router/api-router.go index 0abe810..f76842b 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -414,5 +414,18 @@ func SetApiRouter(router *gin.Engine) { migrationRoute.GET("/users/:id", controller.GetMigrationUser) migrationRoute.POST("/users/:id/convert-to-synced", controller.ConvertMigrationUserToSynced) } + + userMigrationRoute := apiRouter.Group("/user-migrations") + userMigrationRoute.Use(middleware.RootAuth()) + { + userMigrationRoute.POST("/batches", controller.CreateUserMigrationBatch) + userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches) + userMigrationRoute.GET("/batches/:id", controller.GetUserMigrationBatch) + userMigrationRoute.POST("/batches/:id/scan", controller.ScanUserMigrationBatch) + userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch) + userMigrationRoute.POST("/batches/:id/retry", controller.RetryUserMigrationBatch) + userMigrationRoute.POST("/batches/:id/verify", controller.VerifyUserMigrationBatch) + userMigrationRoute.POST("/items/:id/resolve", controller.ResolveUserMigrationItem) + } } }