Просмотр исходного кода

feat(user-migration): add executor and verifier

master
fengsilin 1 неделю назад
Родитель
Сommit
2ecaa4f818
3 измененных файлов: 475 добавлений и 0 удалений
  1. +208
    -0
      service/user_migration/executor.go
  2. +236
    -0
      service/user_migration/service_test.go
  3. +31
    -0
      service/user_migration/verifier.go

+ 208
- 0
service/user_migration/executor.go Просмотреть файл

@@ -0,0 +1,208 @@
package user_migration

import (
"fmt"
"strings"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"gorm.io/gorm"
)

func (s *Service) ExecuteBatch(batchId int) error {
items, err := model.ListExecutableUserMigrationItems(batchId)
if err != nil {
return err
}

for _, item := range items {
if item.Status == model.UserMigrationItemStatusMigrated {
continue
}
if err := s.executeItem(item); err != nil {
_ = model.MarkUserMigrationItemFailed(item.Id, err.Error())
}
}

return model.RefreshUserMigrationBatchExecutionStats(batchId)
}

func (s *Service) executeItem(item *model.UserMigrationItem) error {
switch item.ResolutionStrategy {
case model.UserMigrationStrategyCreateNew:
return s.executeCreateNew(item)
case model.UserMigrationStrategyMergeExisting:
return s.executeMergeExisting(item)
case model.UserMigrationStrategySkip:
return model.MarkUserMigrationItemSkipped(item.Id)
default:
return fmt.Errorf("unsupported strategy: %q", item.ResolutionStrategy)
}
}

func (s *Service) executeCreateNew(item *model.UserMigrationItem) error {
if item.TargetUserId != 0 {
cnUser, err := model.GetUserById(item.TargetUserId, true)
if err != nil {
return err
}
return s.ensureRemoteConverted(item, cnUser.Quota)
}

snapshot, err := s.client.GetUser(item.SourceUserId)
if err != nil {
return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
}

affCode := snapshot.AffCode
if affCode != "" {
if _, err := model.GetUserIdByAffCode(affCode); err == nil {
affCode = common.GetRandomString(4)
}
}

cnUser := &model.User{
Username: snapshot.Username,
Password: snapshot.Password,
Email: snapshot.Email,
DisplayName: snapshot.DisplayName,
Status: snapshot.Status,
Role: snapshot.Role,
Group: snapshot.Group,
Quota: snapshot.Quota,
AffCode: affCode,
CreatedAt: snapshot.CreatedAt,
Setting: snapshot.Setting,
GitHubId: snapshot.GitHubId,
DiscordId: snapshot.DiscordId,
OidcId: snapshot.OidcId,
WeChatId: snapshot.WeChatId,
TelegramId: snapshot.TelegramId,
LinuxDOId: snapshot.LinuxDOId,
}

err = model.DB.Transaction(func(tx *gorm.DB) error {
if err := model.InsertImportedUserWithHashedPasswordTx(tx, cnUser); err != nil {
if isUniqueConstraintErr(err) {
return fmt.Errorf("execution-time identity drift detected, rescan required: %w", err)
}
return fmt.Errorf("create CN user: %w", err)
}

if err := copyMappedBindingsTx(tx, cnUser.Id, snapshot.OAuthBindings); err != nil {
return err
}

if err := tx.Model(&model.UserMigrationItem{}).Where("id = ?", item.Id).
Update("target_user_id", cnUser.Id).Error; err != nil {
return err
}
item.TargetUserId = cnUser.Id
return nil
})
if err != nil {
return err
}

return s.ensureRemoteConverted(item, cnUser.Quota)
}

func (s *Service) executeMergeExisting(item *model.UserMigrationItem) error {
snapshot, err := s.client.GetUser(item.SourceUserId)
if err != nil {
return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
}

finalQuota := 0
err = model.DB.Transaction(func(tx *gorm.DB) error {
_, err := model.AtomicGrantQuotaIfNotExists(tx, item.TargetUserId, snapshot.Quota, item.MigrationToken)
if err != nil {
return err
}

var cnUser model.User
if err := tx.Select("quota").First(&cnUser, item.TargetUserId).Error; err != nil {
return err
}
finalQuota = cnUser.Quota

return copyMappedBindingsTx(tx, item.TargetUserId, snapshot.OAuthBindings)
})
if err != nil {
return err
}

return s.ensureRemoteConverted(item, finalQuota)
}

func copyMappedBindingsTx(tx *gorm.DB, userId int, bindings []RemoteOAuthBinding) error {
if len(bindings) == 0 {
return nil
}

slugs := make([]string, 0, len(bindings))
for _, binding := range bindings {
slugs = append(slugs, binding.ProviderSlug)
}
providers, err := model.GetCustomOAuthProvidersBySlugs(slugs)
if err != nil {
return err
}
slugToID := make(map[string]int, len(providers))
for _, provider := range providers {
slugToID[provider.Slug] = provider.Id
}

localBindings := make([]*model.UserOAuthBinding, 0, len(bindings))
for _, binding := range bindings {
localID, ok := slugToID[binding.ProviderSlug]
if !ok {
continue
}
localBindings = append(localBindings, &model.UserOAuthBinding{
ProviderId: localID,
ProviderUserId: binding.ProviderUserId,
})
}
return model.CopyOAuthBindingsIfMissingTx(tx, userId, localBindings)
}

func (s *Service) ensureRemoteConverted(item *model.UserMigrationItem, syncedQuota int) error {
remoteUser, err := s.client.GetUser(item.SourceUserId)
if err != nil {
return err
}

if remoteUser.Source != common.UserSourceSynced || remoteUser.RemoteUserId != item.TargetUserId {
if err := s.client.ConvertToSynced(item.SourceUserId, &ConvertRemoteUserRequest{
RemoteUserId: item.TargetUserId,
SyncedQuota: syncedQuota,
}); err != nil {
return err
}
}

return markItemMigrated(item.Id)
}

func markItemMigrated(itemID int) error {
return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{
"status": model.UserMigrationItemStatusMigrated,
"migrated_at": timeNow(),
"error_message": "",
}).Error
}

func isUniqueConstraintErr(err error) bool {
if err == nil {
return false
}
message := err.Error()
return strings.Contains(message, "UNIQUE constraint failed") ||
strings.Contains(message, "Duplicate entry") ||
strings.Contains(strings.ToLower(message), "duplicate key value")
}

func timeNow() int64 {
return common.GetTimestamp()
}

+ 236
- 0
service/user_migration/service_test.go Просмотреть файл

@@ -4,6 +4,7 @@ import (
"fmt"
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
@@ -47,6 +48,14 @@ func (f *fakeClient) ConvertToSynced(userId int, req *ConvertRemoteUserRequest)
if req != nil {
copyReq := *req
f.convertReqs[userId] = &copyReq
if snapshot, ok := f.getUserByID[userId]; ok && snapshot != nil {
snapshot.Source = common.UserSourceSynced
snapshot.RemoteUserId = req.RemoteUserId
}
if f.getUserResp != nil && f.getUserResp.Id == userId {
f.getUserResp.Source = common.UserSourceSynced
f.getUserResp.RemoteUserId = req.RemoteUserId
}
}
return nil
}
@@ -237,3 +246,230 @@ func TestScanBatch_Paginated(t *testing.T) {
require.NoError(t, err)
require.Len(t, items, scanPageSize+1)
}

func TestExecuteBatch_CreateNew_CreatesUserAndConvertsRemote(t *testing.T) {
db := setupMigrationServiceDB(t)
orig := model.DB
model.DB = db
defer func() {
model.DB = orig
}()

batch := &model.UserMigrationBatch{
Name: "wave-exec",
Status: model.UserMigrationBatchStatusReady,
}
batch.SourceRegion = "ov"
batch.TargetRegion = "cn"
batch.OperatorId = 1
require.NoError(t, model.CreateUserMigrationBatch(batch))

item := &model.UserMigrationItem{
BatchId: batch.Id,
SourceUserId: 10000001,
SourceUsername: "ov-alice",
SourceEmail: "alice@example.com",
SourceQuota: 300,
ResolutionStrategy: model.UserMigrationStrategyCreateNew,
Status: model.UserMigrationItemStatusReady,
MigrationToken: "batch-1-user-10000001",
}
require.NoError(t, model.UpsertUserMigrationItem(item))

client := &fakeClient{
getUserResp: &RemoteUserSnapshot{
Id: 10000001,
Username: "ov-alice",
Password: "hashed-password",
Email: "alice@example.com",
DisplayName: "Alice",
Status: 1,
Role: 1,
Group: "default",
Quota: 300,
AffCode: "OVA1",
},
}
svc := &Service{client: client}

locked, err := model.TryLockBatchForExecution(batch.Id)
require.NoError(t, err)
require.True(t, locked)
require.NoError(t, svc.ExecuteBatch(batch.Id))

var cnUser model.User
require.NoError(t, db.Where("username = ?", "ov-alice").First(&cnUser).Error)
require.Equal(t, 300, cnUser.Quota)
require.Less(t, cnUser.Id, 10000000)

require.Len(t, client.convertCalls, 1)
require.Equal(t, 10000001, client.convertCalls[0])
}

func TestExecuteBatch_CreateNew_IdempotentOnRetry(t *testing.T) {
db := setupMigrationServiceDB(t)
orig := model.DB
model.DB = db
defer func() {
model.DB = orig
}()

batch := &model.UserMigrationBatch{
Name: "wave-retry",
Status: model.UserMigrationBatchStatusReady,
}
batch.SourceRegion = "ov"
batch.TargetRegion = "cn"
batch.OperatorId = 1
require.NoError(t, model.CreateUserMigrationBatch(batch))

item := &model.UserMigrationItem{
BatchId: batch.Id,
SourceUserId: 10000002,
SourceUsername: "ov-carol",
SourceEmail: "carol@example.com",
SourceQuota: 150,
ResolutionStrategy: model.UserMigrationStrategyCreateNew,
Status: model.UserMigrationItemStatusReady,
MigrationToken: "batch-1-user-10000002",
}
require.NoError(t, model.UpsertUserMigrationItem(item))

snapshot := &RemoteUserSnapshot{
Id: 10000002,
Username: "ov-carol",
Password: "hash",
Email: "carol@example.com",
Quota: 150,
AffCode: "OVC1",
}
client := &fakeClient{getUserResp: snapshot}
svc := &Service{client: client}

locked, err := model.TryLockBatchForExecution(batch.Id)
require.NoError(t, err)
require.True(t, locked)
require.NoError(t, svc.ExecuteBatch(batch.Id))

require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
Update("status", model.UserMigrationBatchStatusFailed).Error)
require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
Update("status", model.UserMigrationItemStatusReady).Error)

locked, err = model.TryLockBatchForExecution(batch.Id)
require.NoError(t, err)
require.True(t, locked)
require.NoError(t, svc.ExecuteBatch(batch.Id))

var count int64
require.NoError(t, db.Model(&model.User{}).Where("username = ?", "ov-carol").Count(&count).Error)
require.EqualValues(t, 1, count)
require.Len(t, client.convertCalls, 1)
}

func TestExecuteBatch_MergeExisting_IdempotentQuota(t *testing.T) {
db := setupMigrationServiceDB(t)
orig := model.DB
model.DB = db
defer func() {
model.DB = orig
}()

require.NoError(t, db.Create(&model.User{
Id: 1,
Username: "cn-bob",
Password: "hash",
Email: "bob@example.com",
AffCode: "CNB1",
Quota: 500,
}).Error)

batch := &model.UserMigrationBatch{
Name: "wave-merge",
Status: model.UserMigrationBatchStatusReady,
}
batch.SourceRegion = "ov"
batch.TargetRegion = "cn"
batch.OperatorId = 1
require.NoError(t, model.CreateUserMigrationBatch(batch))

item := &model.UserMigrationItem{
BatchId: batch.Id,
SourceUserId: 10000002,
TargetUserId: 1,
SourceUsername: "ov-bob",
SourceEmail: "bob@example.com",
SourceQuota: 200,
ResolutionStrategy: model.UserMigrationStrategyMergeExisting,
Status: model.UserMigrationItemStatusReady,
MigrationToken: "batch-1-user-10000002",
}
require.NoError(t, model.UpsertUserMigrationItem(item))

client := &fakeClient{
getUserResp: &RemoteUserSnapshot{
Id: 10000002,
Username: "ov-bob",
Password: "hashed",
Email: "bob@example.com",
Quota: 200,
},
}
svc := &Service{client: client}

for i := 0; i < 2; i++ {
require.NoError(t, db.Model(&model.UserMigrationBatch{}).Where("id = ?", batch.Id).
Update("status", model.UserMigrationBatchStatusReady).Error)
require.NoError(t, db.Model(&model.UserMigrationItem{}).Where("batch_id = ?", batch.Id).
Update("status", model.UserMigrationItemStatusReady).Error)

locked, err := model.TryLockBatchForExecution(batch.Id)
require.NoError(t, err)
require.True(t, locked)
require.NoError(t, svc.ExecuteBatch(batch.Id))
}

var saved model.User
require.NoError(t, db.First(&saved, 1).Error)
require.Equal(t, 700, saved.Quota)
require.Len(t, client.convertCalls, 1)
}

func TestVerifyItem_CorrectlyChecksRemoteConverted(t *testing.T) {
db := setupMigrationServiceDB(t)
orig := model.DB
model.DB = db
defer func() {
model.DB = orig
}()

require.NoError(t, db.Create(&model.User{
Id: 789,
Username: "cn-migrated",
Password: "hash",
AffCode: "CM1",
Quota: 300,
}).Error)

item := &model.UserMigrationItem{
BatchId: 1,
SourceUserId: 10000001,
TargetUserId: 789,
SourceQuota: 300,
Status: model.UserMigrationItemStatusMigrated,
}

client := &fakeClient{
getUserResp: &RemoteUserSnapshot{
Id: 10000001,
Source: common.UserSourceSynced,
RemoteUserId: 789,
},
}
svc := &Service{client: client}
result, err := svc.VerifyItem(item)
require.NoError(t, err)
require.True(t, result.TargetUserExists)
require.True(t, result.RemoteConverted)
require.True(t, result.QuotaMatched)
}

+ 31
- 0
service/user_migration/verifier.go Просмотреть файл

@@ -0,0 +1,31 @@
package user_migration

import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
)

type VerifyResult struct {
TargetUserExists bool `json:"target_user_exists"`
RemoteConverted bool `json:"remote_converted"`
QuotaMatched bool `json:"quota_matched"`
}

func (s *Service) VerifyItem(item *model.UserMigrationItem) (*VerifyResult, error) {
targetUser, err := model.GetUserById(item.TargetUserId, true)
if err != nil {
return nil, err
}

remoteUser, err := s.client.GetUser(item.SourceUserId)
if err != nil {
return nil, err
}

return &VerifyResult{
TargetUserExists: targetUser != nil,
RemoteConverted: remoteUser.Source == common.UserSourceSynced &&
remoteUser.RemoteUserId == item.TargetUserId,
QuotaMatched: targetUser != nil && targetUser.Quota >= item.SourceQuota,
}, nil
}

Загрузка…
Отмена
Сохранить