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