diff --git a/model/user.go b/model/user.go index 142e1de..8d9a1ce 100644 --- a/model/user.go +++ b/model/user.go @@ -348,13 +348,66 @@ func GetUserById(id int, selectAll bool) (*User, error) { return &user, err } +func getUserByStringColumn(column string, value string) (*User, error) { + if value == "" { + return nil, nil + } + var user User + result := DB.Where(column+" = ?", value).Limit(1).Find(&user) + if result.Error != nil { + return nil, result.Error + } + if result.RowsAffected == 0 { + return nil, nil + } + return &user, nil +} + +func GetUserByUsername(username string) (*User, error) { + return getUserByStringColumn("username", username) +} + +func GetUserByEmail(email string) (*User, error) { + return getUserByStringColumn("email", email) +} + +func GetUserByGitHubId(githubId string) (*User, error) { + return getUserByStringColumn("github_id", githubId) +} + +func GetUserByDiscordId(discordId string) (*User, error) { + return getUserByStringColumn("discord_id", discordId) +} + +func GetUserByOidcId(oidcId string) (*User, error) { + return getUserByStringColumn("oidc_id", oidcId) +} + +func GetUserByWeChatId(wechatId string) (*User, error) { + return getUserByStringColumn("wechat_id", wechatId) +} + +func GetUserByTelegramId(telegramId string) (*User, error) { + return getUserByStringColumn("telegram_id", telegramId) +} + +func GetUserByLinuxDOId(linuxDOId string) (*User, error) { + return getUserByStringColumn("linux_do_id", linuxDOId) +} + func GetUserIdByAffCode(affCode string) (int, error) { if affCode == "" { return 0, errors.New("affCode 为空!") } var user User - err := DB.Select("id").First(&user, "aff_code = ?", affCode).Error - return user.Id, err + result := DB.Select("id").Where("aff_code = ?", affCode).Limit(1).Find(&user) + if result.Error != nil { + return 0, result.Error + } + if result.RowsAffected == 0 { + return 0, gorm.ErrRecordNotFound + } + return user.Id, nil } func DeleteUserById(id int) (err error) { diff --git a/model/user_oauth_binding.go b/model/user_oauth_binding.go index c6ce868..c858d27 100644 --- a/model/user_oauth_binding.go +++ b/model/user_oauth_binding.go @@ -40,15 +40,21 @@ func GetUserOAuthBinding(userId, providerId int) (*UserOAuthBinding, error) { // GetUserByOAuthBinding finds a user by provider ID and provider user ID func GetUserByOAuthBinding(providerId int, providerUserId string) (*User, error) { var binding UserOAuthBinding - err := DB.Where("provider_id = ? AND provider_user_id = ?", providerId, providerUserId).First(&binding).Error - if err != nil { - return nil, err + result := DB.Where("provider_id = ? AND provider_user_id = ?", providerId, providerUserId).Limit(1).Find(&binding) + if result.Error != nil { + return nil, result.Error + } + if result.RowsAffected == 0 { + return nil, gorm.ErrRecordNotFound } var user User - err = DB.First(&user, binding.UserId).Error - if err != nil { - return nil, err + userResult := DB.Where("id = ?", binding.UserId).Limit(1).Find(&user) + if userResult.Error != nil { + return nil, userResult.Error + } + if userResult.RowsAffected == 0 { + return nil, gorm.ErrRecordNotFound } return &user, nil } diff --git a/service/user_migration/provider_mapping.go b/service/user_migration/provider_mapping.go new file mode 100644 index 0000000..e7b3338 --- /dev/null +++ b/service/user_migration/provider_mapping.go @@ -0,0 +1,16 @@ +package user_migration + +import "github.com/QuantumNous/new-api/model" + +func buildProviderSlugMap() (map[string]int, error) { + providers, err := model.GetEnabledCustomOAuthProviders() + if err != nil { + return nil, err + } + + result := make(map[string]int, len(providers)) + for _, provider := range providers { + result[provider.Slug] = provider.Id + } + return result, nil +} diff --git a/service/user_migration/scanner.go b/service/user_migration/scanner.go new file mode 100644 index 0000000..87ea455 --- /dev/null +++ b/service/user_migration/scanner.go @@ -0,0 +1,185 @@ +package user_migration + +import ( + "errors" + "fmt" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +const scanPageSize = 100 + +type Service struct { + client migrationClient +} + +func NewService(client migrationClient) *Service { + return &Service{client: client} +} + +func (s *Service) ScanBatch(batchId int) error { + providerSlugMap, err := buildProviderSlugMap() + if err != nil { + return err + } + + page := 0 + for { + resp, err := s.client.ListUsers(page, scanPageSize) + if err != nil { + return fmt.Errorf("list users page %d: %w", page, err) + } + for _, remoteUser := range resp.Data { + result, err := analyzeConflicts(&remoteUser, providerSlugMap, s.client) + if err != nil { + return fmt.Errorf("analyze conflicts for user %d: %w", remoteUser.Id, err) + } + + flagsJSON, err := common.Marshal(result.Flags) + if err != nil { + return err + } + snapshotJSON, err := common.Marshal(remoteUser) + if err != nil { + return err + } + + item := &model.UserMigrationItem{ + BatchId: batchId, + SourceUserId: remoteUser.Id, + SourceUsername: remoteUser.Username, + SourceEmail: remoteUser.Email, + SourceQuota: remoteUser.Quota, + ConflictFlags: string(flagsJSON), + ResolutionStrategy: result.Strategy, + Status: result.Status, + Snapshot: string(snapshotJSON), + MigrationToken: fmt.Sprintf("batch-%d-user-%d", batchId, remoteUser.Id), + TargetUserId: result.TargetCNUserID, + } + if err := model.UpsertUserMigrationItem(item); err != nil { + return err + } + } + + if len(resp.Data) < scanPageSize { + break + } + page++ + } + + return model.RefreshUserMigrationBatchStats(batchId) +} + +type conflictResult struct { + Flags []string + Strategy string + Status string + TargetCNUserID int +} + +func analyzeConflicts(remote *RemoteUserSnapshot, providerSlugMap map[string]int, client migrationClient) (*conflictResult, error) { + result := &conflictResult{} + var mergeCandidate *model.User + hardConflict := false + + assignMergeCandidate := func(flag string, user *model.User) { + if user == nil { + return + } + result.Flags = append(result.Flags, flag) + if mergeCandidate == nil { + mergeCandidate = user + return + } + if mergeCandidate.Id != user.Id { + hardConflict = true + result.Flags = append(result.Flags, "multiple_merge_targets") + } + } + + if user, _ := model.GetUserByUsername(remote.Username); user != nil { + result.Flags = append(result.Flags, "username") + hardConflict = true + } + + if remote.Email != "" { + user, err := model.GetUserByEmail(remote.Email) + if err != nil { + return nil, err + } + assignMergeCandidate("email", user) + } + + type oauthField struct { + value string + flag string + query func(string) (*model.User, error) + } + fields := []oauthField{ + {value: remote.GitHubId, flag: "github_id", query: model.GetUserByGitHubId}, + {value: remote.DiscordId, flag: "discord_id", query: model.GetUserByDiscordId}, + {value: remote.OidcId, flag: "oidc_id", query: model.GetUserByOidcId}, + {value: remote.WeChatId, flag: "wechat_id", query: model.GetUserByWeChatId}, + {value: remote.TelegramId, flag: "telegram_id", query: model.GetUserByTelegramId}, + {value: remote.LinuxDOId, flag: "linux_do_id", query: model.GetUserByLinuxDOId}, + } + for _, field := range fields { + if field.value == "" { + continue + } + user, err := field.query(field.value) + if err != nil { + return nil, err + } + assignMergeCandidate(field.flag, user) + } + + for _, binding := range remote.OAuthBindings { + localProviderID, ok := providerSlugMap[binding.ProviderSlug] + if !ok { + continue + } + user, err := model.GetUserByOAuthBinding(localProviderID, binding.ProviderUserId) + if errors.Is(err, gorm.ErrRecordNotFound) { + continue + } + if err != nil { + return nil, err + } + assignMergeCandidate("oauth:"+binding.ProviderSlug, user) + } + + if remote.AffCode != "" { + if _, err := model.GetUserIdByAffCode(remote.AffCode); err == nil { + result.Flags = append(result.Flags, "aff_code") + } + } + + if hardConflict { + result.Status = model.UserMigrationItemStatusConflict + return result, nil + } + + if mergeCandidate != nil { + hasCopy, err := client.HasSyncedCopy(mergeCandidate.Id) + if err != nil { + return nil, fmt.Errorf("check synced copy for CN user %d: %w", mergeCandidate.Id, err) + } + if hasCopy { + result.Flags = append(result.Flags, "cn_already_synced_to_ov") + result.Status = model.UserMigrationItemStatusConflict + return result, nil + } + result.Status = model.UserMigrationItemStatusReady + result.Strategy = model.UserMigrationStrategyMergeExisting + result.TargetCNUserID = mergeCandidate.Id + return result, nil + } + + result.Status = model.UserMigrationItemStatusReady + result.Strategy = model.UserMigrationStrategyCreateNew + return result, nil +} diff --git a/service/user_migration/service_test.go b/service/user_migration/service_test.go new file mode 100644 index 0000000..1b82fb7 --- /dev/null +++ b/service/user_migration/service_test.go @@ -0,0 +1,239 @@ +package user_migration + +import ( + "fmt" + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +type fakeClient struct { + listPages map[int][]RemoteUserSnapshot + getUserResp *RemoteUserSnapshot + getUserByID map[int]*RemoteUserSnapshot + hasCopyResult bool + hasCopyByID map[int]bool + convertCalls []int + convertReqs map[int]*ConvertRemoteUserRequest +} + +func (f *fakeClient) ListUsers(page, pageSize int) (*ListRemoteUsersResponse, error) { + users := f.listPages[page] + return &ListRemoteUsersResponse{ + Success: true, + Data: users, + Total: len(users), + }, nil +} + +func (f *fakeClient) GetUser(userId int) (*RemoteUserSnapshot, error) { + if snapshot, ok := f.getUserByID[userId]; ok { + return snapshot, nil + } + if f.getUserResp != nil { + return f.getUserResp, nil + } + return nil, fmt.Errorf("user %d not found", userId) +} + +func (f *fakeClient) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error { + f.convertCalls = append(f.convertCalls, userId) + if f.convertReqs == nil { + f.convertReqs = make(map[int]*ConvertRemoteUserRequest) + } + if req != nil { + copyReq := *req + f.convertReqs[userId] = ©Req + } + return nil +} + +func (f *fakeClient) HasSyncedCopy(cnUserId int) (bool, error) { + if value, ok := f.hasCopyByID[cnUserId]; ok { + return value, nil + } + return f.hasCopyResult, nil +} + +func setupMigrationServiceDB(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.CustomOAuthProvider{}, + &model.UserOAuthBinding{}, + &model.UserMigrationBatch{}, + &model.UserMigrationItem{}, + &model.MigrationQuotaGrant{}, + )) + + return db +} + +func TestScanBatch_NoConflict_MarksReady(t *testing.T) { + db := setupMigrationServiceDB(t) + orig := model.DB + model.DB = db + defer func() { + model.DB = orig + }() + + batch := &model.UserMigrationBatch{Name: "wave-1", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1} + require.NoError(t, model.CreateUserMigrationBatch(batch)) + + client := &fakeClient{ + listPages: map[int][]RemoteUserSnapshot{ + 0: {{ + Id: 10000001, + Username: "ov-bob", + Email: "bob@example.com", + Quota: 456, + AffCode: "A002", + }}, + }, + } + svc := &Service{client: client} + require.NoError(t, svc.ScanBatch(batch.Id)) + + items, err := model.ListUserMigrationItemsByBatch(batch.Id) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, model.UserMigrationStrategyCreateNew, items[0].ResolutionStrategy) + require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status) + + var savedBatch model.UserMigrationBatch + require.NoError(t, db.First(&savedBatch, batch.Id).Error) + require.Equal(t, model.UserMigrationBatchStatusReady, savedBatch.Status) +} + +func TestScanBatch_EmailConflict_NoCNSyncedCopy_RecommendsMerge(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-alice", + Password: "hash", + Email: "alice@example.com", + AffCode: "C001", + }).Error) + + batch := &model.UserMigrationBatch{Name: "wave-2", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1} + require.NoError(t, model.CreateUserMigrationBatch(batch)) + + client := &fakeClient{ + listPages: map[int][]RemoteUserSnapshot{ + 0: {{ + Id: 10000001, + Username: "ov-alice", + Email: "alice@example.com", + Quota: 123, + AffCode: "A001", + }}, + }, + hasCopyResult: false, + } + svc := &Service{client: client} + require.NoError(t, svc.ScanBatch(batch.Id)) + + items, err := model.ListUserMigrationItemsByBatch(batch.Id) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, model.UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy) + require.Equal(t, model.UserMigrationItemStatusReady, items[0].Status) + require.Equal(t, 1, items[0].TargetUserId) +} + +func TestScanBatch_EmailConflict_CNHasSyncedCopy_MarksHardConflict(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: 456, + Username: "cn-user456", + Password: "hash", + Email: "shared@example.com", + AffCode: "C456", + }).Error) + + batch := &model.UserMigrationBatch{Name: "wave-3", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1} + require.NoError(t, model.CreateUserMigrationBatch(batch)) + + client := &fakeClient{ + listPages: map[int][]RemoteUserSnapshot{ + 0: {{ + Id: 10000003, + Username: "ov-x", + Email: "shared@example.com", + Quota: 50, + AffCode: "OVX", + }}, + }, + hasCopyResult: true, + } + svc := &Service{client: client} + require.NoError(t, svc.ScanBatch(batch.Id)) + + items, err := model.ListUserMigrationItemsByBatch(batch.Id) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, model.UserMigrationItemStatusConflict, items[0].Status) + require.Contains(t, items[0].ConflictFlags, "cn_already_synced_to_ov") +} + +func TestScanBatch_Paginated(t *testing.T) { + db := setupMigrationServiceDB(t) + orig := model.DB + model.DB = db + defer func() { + model.DB = orig + }() + + batch := &model.UserMigrationBatch{Name: "wave-big", SourceRegion: "ov", TargetRegion: "cn", OperatorId: 1} + require.NoError(t, model.CreateUserMigrationBatch(batch)) + + page0 := make([]RemoteUserSnapshot, scanPageSize) + for i := range page0 { + page0[i] = RemoteUserSnapshot{ + Id: 10000001 + i, + Username: fmt.Sprintf("u%d", i), + AffCode: fmt.Sprintf("A%d", i), + } + } + page1 := []RemoteUserSnapshot{{ + Id: 10000001 + scanPageSize, + Username: "last", + AffCode: "LAST", + }} + + client := &fakeClient{ + listPages: map[int][]RemoteUserSnapshot{ + 0: page0, + 1: page1, + }, + } + svc := &Service{client: client} + require.NoError(t, svc.ScanBatch(batch.Id)) + + items, err := model.ListUserMigrationItemsByBatch(batch.Id) + require.NoError(t, err) + require.Len(t, items, scanPageSize+1) +}