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

feat(user-migration): add batch item and quota grant models

master
fengsilin 1 неделю назад
Родитель
Сommit
26cb2273a4
3 измененных файлов: 395 добавлений и 0 удалений
  1. +6
    -0
      model/main.go
  2. +239
    -0
      model/user_migration.go
  3. +150
    -0
      model/user_migration_test.go

+ 6
- 0
model/main.go Просмотреть файл

@@ -286,6 +286,9 @@ func migrateDB() error {
&QuotaSyncLog{},
&EmailQuotaRule{},
&UserModelRateLimit{},
&UserMigrationBatch{},
&UserMigrationItem{},
&MigrationQuotaGrant{},
)
if err != nil {
return err
@@ -350,6 +353,9 @@ func migrateDBFast() error {
{&PendingSyncRecord{}, "PendingSyncRecord"},
{&QuotaSyncLog{}, "QuotaSyncLog"},
{&EmailQuotaRule{}, "EmailQuotaRule"},
{&UserMigrationBatch{}, "UserMigrationBatch"},
{&UserMigrationItem{}, "UserMigrationItem"},
{&MigrationQuotaGrant{}, "MigrationQuotaGrant"},
}
// 动态计算migration数量,确保errChan缓冲区足够大
errChan := make(chan error, len(migrations))


+ 239
- 0
model/user_migration.go Просмотреть файл

@@ -0,0 +1,239 @@
package model

import (
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)

const (
UserMigrationBatchStatusDraft = "draft"
UserMigrationBatchStatusScanned = "scanned"
UserMigrationBatchStatusReady = "ready"
UserMigrationBatchStatusRunning = "running"
UserMigrationBatchStatusCompleted = "completed"
UserMigrationBatchStatusFailed = "failed"
UserMigrationBatchStatusCancelled = "cancelled"

UserMigrationItemStatusPendingScan = "pending_scan"
UserMigrationItemStatusConflict = "conflict"
UserMigrationItemStatusReady = "ready"
UserMigrationItemStatusMigrating = "migrating"
UserMigrationItemStatusMigrated = "migrated"
UserMigrationItemStatusSkipped = "skipped"
UserMigrationItemStatusFailed = "failed"

UserMigrationStrategyCreateNew = "create_new"
UserMigrationStrategyMergeExisting = "merge_into_existing"
UserMigrationStrategySkip = "skip"
)

type UserMigrationBatch struct {
Id int `json:"id" gorm:"primaryKey"`
Name string `json:"name" gorm:"type:varchar(128);not null"`
SourceRegion string `json:"source_region" gorm:"type:varchar(32);not null"`
TargetRegion string `json:"target_region" gorm:"type:varchar(32);not null"`
Status string `json:"status" gorm:"type:varchar(32);not null;default:'draft';index"`
TotalUsers int `json:"total_users" gorm:"default:0"`
ConflictUsers int `json:"conflict_users" gorm:"default:0"`
MigratedUsers int `json:"migrated_users" gorm:"default:0"`
FailedUsers int `json:"failed_users" gorm:"default:0"`
OperatorId int `json:"operator_id" gorm:"index"`
CreatedAt int64 `json:"created_at" gorm:"bigint"`
UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
}

type UserMigrationItem struct {
Id int `json:"id" gorm:"primaryKey"`
BatchId int `json:"batch_id" gorm:"index:idx_batch_source_user,unique"`
SourceUserId int `json:"source_user_id" gorm:"index:idx_batch_source_user,unique"`
TargetUserId int `json:"target_user_id" gorm:"default:0;index"`
SourceUsername string `json:"source_username" gorm:"type:varchar(128);not null"`
SourceEmail string `json:"source_email" gorm:"type:varchar(255);default:''"`
SourceQuota int `json:"source_quota" gorm:"default:0"`
ConflictFlags string `json:"conflict_flags" gorm:"type:text"`
ResolutionStrategy string `json:"resolution_strategy" gorm:"type:varchar(32);default:''"`
Status string `json:"status" gorm:"type:varchar(32);not null;default:'pending_scan';index"`
ErrorMessage string `json:"error_message" gorm:"type:text"`
Snapshot string `json:"snapshot" gorm:"type:text"`
MigrationToken string `json:"migration_token" gorm:"type:varchar(64);default:'';index"`
MigratedAt int64 `json:"migrated_at" gorm:"bigint;default:0"`
CreatedAt int64 `json:"created_at" gorm:"bigint"`
UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
}

type MigrationQuotaGrant struct {
Id int `json:"id" gorm:"primaryKey"`
MigrationToken string `json:"migration_token" gorm:"type:varchar(64);not null;uniqueIndex"`
CreatedAt int64 `json:"created_at" gorm:"bigint"`
}

func (b *UserMigrationBatch) BeforeCreate(tx *gorm.DB) error {
now := common.GetTimestamp()
if b.Status == "" {
b.Status = UserMigrationBatchStatusDraft
}
b.CreatedAt = now
b.UpdatedAt = now
return nil
}

func (b *UserMigrationBatch) BeforeUpdate(tx *gorm.DB) error {
b.UpdatedAt = common.GetTimestamp()
return nil
}

func (i *UserMigrationItem) BeforeCreate(tx *gorm.DB) error {
now := common.GetTimestamp()
if i.Status == "" {
i.Status = UserMigrationItemStatusPendingScan
}
i.CreatedAt = now
i.UpdatedAt = now
return nil
}

func (i *UserMigrationItem) BeforeUpdate(tx *gorm.DB) error {
i.UpdatedAt = common.GetTimestamp()
return nil
}

func (g *MigrationQuotaGrant) BeforeCreate(tx *gorm.DB) error {
g.CreatedAt = common.GetTimestamp()
return nil
}

func CreateUserMigrationBatch(batch *UserMigrationBatch) error {
return DB.Create(batch).Error
}

func UpsertUserMigrationItem(item *UserMigrationItem) error {
var existing UserMigrationItem
result := DB.Where("batch_id = ? AND source_user_id = ?", item.BatchId, item.SourceUserId).Limit(1).Find(&existing)
if result.Error != nil {
return result.Error
}
if result.RowsAffected > 0 {
updates := map[string]any{
"source_username": item.SourceUsername,
"source_email": item.SourceEmail,
"source_quota": item.SourceQuota,
"conflict_flags": item.ConflictFlags,
"status": item.Status,
"snapshot": item.Snapshot,
}
if existing.ResolutionStrategy == "" && item.ResolutionStrategy != "" {
updates["resolution_strategy"] = item.ResolutionStrategy
}
if existing.TargetUserId == 0 && item.TargetUserId != 0 {
updates["target_user_id"] = item.TargetUserId
}
item.Id = existing.Id
return DB.Model(&existing).Updates(updates).Error
}
return DB.Create(item).Error
}

func ListUserMigrationItemsByBatch(batchId int) ([]*UserMigrationItem, error) {
var items []*UserMigrationItem
err := DB.Where("batch_id = ?", batchId).Order("id asc").Find(&items).Error
return items, err
}

func ListExecutableUserMigrationItems(batchId int) ([]*UserMigrationItem, error) {
var items []*UserMigrationItem
err := DB.Where("batch_id = ? AND status IN ?", batchId, []string{
UserMigrationItemStatusReady,
UserMigrationItemStatusMigrating,
UserMigrationItemStatusFailed,
}).Order("id asc").Find(&items).Error
return items, err
}

func MarkUserMigrationItemFailed(id int, errMsg string) error {
return DB.Model(&UserMigrationItem{}).Where("id = ?", id).Updates(map[string]any{
"status": UserMigrationItemStatusFailed,
"error_message": errMsg,
}).Error
}

func MarkUserMigrationItemSkipped(id int) error {
return DB.Model(&UserMigrationItem{}).Where("id = ?", id).
Update("status", UserMigrationItemStatusSkipped).Error
}

func RefreshUserMigrationBatchStats(batchId int) error {
var total, conflicts, pendingScan int64
DB.Model(&UserMigrationItem{}).Where("batch_id = ?", batchId).Count(&total)
DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusConflict).Count(&conflicts)
DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusPendingScan).Count(&pendingScan)

status := UserMigrationBatchStatusScanned
if conflicts == 0 && pendingScan == 0 && total > 0 {
status = UserMigrationBatchStatusReady
}

return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
"total_users": int(total),
"conflict_users": int(conflicts),
"status": status,
}).Error
}

func RefreshUserMigrationBatchExecutionStats(batchId int) error {
var migrated, failed, pending int64
DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusMigrated).Count(&migrated)
DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status = ?", batchId, UserMigrationItemStatusFailed).Count(&failed)
DB.Model(&UserMigrationItem{}).Where("batch_id = ? AND status IN ?", batchId, []string{
UserMigrationItemStatusReady,
UserMigrationItemStatusMigrating,
}).Count(&pending)

status := UserMigrationBatchStatusCompleted
if failed > 0 {
status = UserMigrationBatchStatusFailed
}
if pending > 0 {
status = UserMigrationBatchStatusRunning
}

return DB.Model(&UserMigrationBatch{}).Where("id = ?", batchId).Updates(map[string]any{
"migrated_users": int(migrated),
"failed_users": int(failed),
"status": status,
}).Error
}

func TryLockBatchForExecution(batchId int) (bool, error) {
result := DB.Model(&UserMigrationBatch{}).
Where("id = ? AND status IN ?", batchId, []string{
UserMigrationBatchStatusReady,
UserMigrationBatchStatusFailed,
}).
Update("status", UserMigrationBatchStatusRunning)
return result.RowsAffected == 1, result.Error
}

func TryLockBatchForScan(batchId int) (bool, error) {
result := DB.Model(&UserMigrationBatch{}).
Where("id = ? AND status IN ?", batchId, []string{
UserMigrationBatchStatusDraft,
UserMigrationBatchStatusScanned,
UserMigrationBatchStatusReady,
}).
Update("status", UserMigrationBatchStatusScanned)
return result.RowsAffected == 1, result.Error
}

func AtomicGrantQuotaIfNotExists(tx *gorm.DB, userId, delta int, token string) (bool, error) {
result := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&MigrationQuotaGrant{MigrationToken: token})
if result.Error != nil {
return false, result.Error
}
if result.RowsAffected == 0 {
return true, nil
}

return false, tx.Model(&User{}).Where("id = ?", userId).
Update("quota", gorm.Expr("quota + ?", delta)).Error
}

+ 150
- 0
model/user_migration_test.go Просмотреть файл

@@ -0,0 +1,150 @@
package model

import (
"testing"

"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)

func setupUserMigrationTestDB(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(
&User{},
&CustomOAuthProvider{},
&UserOAuthBinding{},
&UserMigrationBatch{},
&UserMigrationItem{},
&MigrationQuotaGrant{},
))

return db
}

func TestUserMigrationModels_AutoMigrate(t *testing.T) {
db := setupUserMigrationTestDB(t)

require.True(t, db.Migrator().HasTable(&UserMigrationBatch{}))
require.True(t, db.Migrator().HasTable(&UserMigrationItem{}))
require.True(t, db.Migrator().HasTable(&MigrationQuotaGrant{}))
}

func TestCreateUserMigrationBatch_DefaultStatus(t *testing.T) {
db := setupUserMigrationTestDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

batch := &UserMigrationBatch{
Name: "ov-to-cn-wave-1",
SourceRegion: "overseas",
TargetRegion: "cn",
OperatorId: 1,
}

require.NoError(t, CreateUserMigrationBatch(batch))
require.Equal(t, UserMigrationBatchStatusDraft, batch.Status)
require.Greater(t, batch.Id, 0)
}

func TestUpsertMigrationItem_PreservesHumanSetStrategy(t *testing.T) {
db := setupUserMigrationTestDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

batch := &UserMigrationBatch{
Name: "b1",
SourceRegion: "ov",
TargetRegion: "cn",
OperatorId: 1,
}
require.NoError(t, CreateUserMigrationBatch(batch))

first := &UserMigrationItem{
BatchId: batch.Id,
SourceUserId: 10000001,
SourceUsername: "alice",
SourceEmail: "alice@example.com",
SourceQuota: 12345,
ResolutionStrategy: "",
Status: UserMigrationItemStatusConflict,
}
require.NoError(t, UpsertUserMigrationItem(first))

require.NoError(t, DB.Model(&UserMigrationItem{}).
Where("batch_id = ? AND source_user_id = ?", batch.Id, 10000001).
Update("resolution_strategy", UserMigrationStrategyMergeExisting).Error)

second := &UserMigrationItem{
BatchId: batch.Id,
SourceUserId: 10000001,
SourceUsername: "alice",
SourceEmail: "alice@example.com",
SourceQuota: 99999,
ResolutionStrategy: UserMigrationStrategyCreateNew,
Status: UserMigrationItemStatusReady,
}
require.NoError(t, UpsertUserMigrationItem(second))

items, err := ListUserMigrationItemsByBatch(batch.Id)
require.NoError(t, err)
require.Len(t, items, 1)
require.Equal(t, UserMigrationStrategyMergeExisting, items[0].ResolutionStrategy)
require.Equal(t, UserMigrationItemStatusReady, items[0].Status)
require.Equal(t, 99999, items[0].SourceQuota)
}

func TestAtomicGrantQuotaIfNotExists_Idempotent(t *testing.T) {
db := setupUserMigrationTestDB(t)
orig := DB
DB = db
defer func() {
DB = orig
}()

require.NoError(t, db.Create(&User{
Id: 1,
Username: "cn-user",
Password: "hash",
AffCode: "AFF1",
Quota: 500,
}).Error)

var alreadyDone bool
err := db.Transaction(func(tx *gorm.DB) error {
var innerErr error
alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
return innerErr
})
require.NoError(t, err)
require.False(t, alreadyDone)

var user User
require.NoError(t, db.First(&user, 1).Error)
require.Equal(t, 700, user.Quota)

err = db.Transaction(func(tx *gorm.DB) error {
var innerErr error
alreadyDone, innerErr = AtomicGrantQuotaIfNotExists(tx, 1, 200, "token-abc")
return innerErr
})
require.NoError(t, err)
require.True(t, alreadyDone)

require.NoError(t, db.First(&user, 1).Error)
require.Equal(t, 700, user.Quota)
}

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