Browse Source

feat(user-migration): add root admin migration API

master
fengsilin 1 week ago
parent
commit
7282142316
3 changed files with 383 additions and 0 deletions
  1. +233
    -0
      controller/user_migration.go
  2. +137
    -0
      controller/user_migration_test.go
  3. +13
    -0
      router/api-router.go

+ 233
- 0
controller/user_migration.go View File

@@ -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)
}

+ 137
- 0
controller/user_migration_test.go View File

@@ -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)
}

+ 13
- 0
router/api-router.go View File

@@ -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)
}
}
}

Loading…
Cancel
Save