diff --git a/service/user_migration/executor.go b/service/user_migration/executor.go new file mode 100644 index 0000000..51dcb80 --- /dev/null +++ b/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() +} diff --git a/service/user_migration/service_test.go b/service/user_migration/service_test.go index 1b82fb7..8f6146e 100644 --- a/service/user_migration/service_test.go +++ b/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] = ©Req + 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) +} diff --git a/service/user_migration/verifier.go b/service/user_migration/verifier.go new file mode 100644 index 0000000..229d70b --- /dev/null +++ b/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 +}