fix(group): SIMPLE 模式启动补齐默认分组
This commit is contained in:
@@ -65,5 +65,18 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
|
|||||||
|
|
||||||
// 创建 Ent 客户端,绑定到已配置的数据库驱动。
|
// 创建 Ent 客户端,绑定到已配置的数据库驱动。
|
||||||
client := ent.NewClient(ent.Driver(drv))
|
client := ent.NewClient(ent.Driver(drv))
|
||||||
|
|
||||||
|
// SIMPLE 模式:启动时补齐各平台默认分组。
|
||||||
|
// - anthropic/openai/gemini: 确保存在 <platform>-default
|
||||||
|
// - antigravity: 仅要求存在 >=2 个未软删除分组(用于 claude/gemini 混合调度场景)
|
||||||
|
if cfg.RunMode == config.RunModeSimple {
|
||||||
|
seedCtx, seedCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||||
|
defer seedCancel()
|
||||||
|
if err := ensureSimpleModeDefaultGroups(seedCtx, client); err != nil {
|
||||||
|
_ = client.Close()
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
return client, drv.DB(), nil
|
return client, drv.DB(), nil
|
||||||
}
|
}
|
||||||
|
|||||||
82
backend/internal/repository/simple_mode_default_groups.go
Normal file
82
backend/internal/repository/simple_mode_default_groups.go
Normal file
@@ -0,0 +1,82 @@
|
|||||||
|
package repository
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||||
|
"github.com/Wei-Shaw/sub2api/ent/group"
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||||
|
)
|
||||||
|
|
||||||
|
func ensureSimpleModeDefaultGroups(ctx context.Context, client *dbent.Client) error {
|
||||||
|
if client == nil {
|
||||||
|
return fmt.Errorf("nil ent client")
|
||||||
|
}
|
||||||
|
|
||||||
|
requiredByPlatform := map[string]int{
|
||||||
|
service.PlatformAnthropic: 1,
|
||||||
|
service.PlatformOpenAI: 1,
|
||||||
|
service.PlatformGemini: 1,
|
||||||
|
service.PlatformAntigravity: 2,
|
||||||
|
}
|
||||||
|
|
||||||
|
for platform, minCount := range requiredByPlatform {
|
||||||
|
count, err := client.Group.Query().
|
||||||
|
Where(group.PlatformEQ(platform), group.DeletedAtIsNil()).
|
||||||
|
Count(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("count groups for platform %s: %w", platform, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if platform == service.PlatformAntigravity {
|
||||||
|
if count < minCount {
|
||||||
|
for i := count; i < minCount; i++ {
|
||||||
|
name := fmt.Sprintf("%s-default-%d", platform, i+1)
|
||||||
|
if err := createGroupIfNotExists(ctx, client, name, platform); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Non-antigravity platforms: ensure <platform>-default exists.
|
||||||
|
name := platform + "-default"
|
||||||
|
if err := createGroupIfNotExists(ctx, client, name, platform); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func createGroupIfNotExists(ctx context.Context, client *dbent.Client, name, platform string) error {
|
||||||
|
exists, err := client.Group.Query().
|
||||||
|
Where(group.NameEQ(name), group.DeletedAtIsNil()).
|
||||||
|
Exist(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("check group exists %s: %w", name, err)
|
||||||
|
}
|
||||||
|
if exists {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = client.Group.Create().
|
||||||
|
SetName(name).
|
||||||
|
SetDescription("Auto-created default group").
|
||||||
|
SetPlatform(platform).
|
||||||
|
SetStatus(service.StatusActive).
|
||||||
|
SetSubscriptionType(service.SubscriptionTypeStandard).
|
||||||
|
SetRateMultiplier(1.0).
|
||||||
|
SetIsExclusive(false).
|
||||||
|
Save(ctx)
|
||||||
|
if err != nil {
|
||||||
|
if dbent.IsConstraintError(err) {
|
||||||
|
// Concurrent server startups may race on creation; treat as success.
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("create default group %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,84 @@
|
|||||||
|
//go:build integration
|
||||||
|
|
||||||
|
package repository
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Wei-Shaw/sub2api/ent/group"
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestEnsureSimpleModeDefaultGroups_CreatesMissingDefaults(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
tx := testEntTx(t)
|
||||||
|
client := tx.Client()
|
||||||
|
|
||||||
|
seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
|
||||||
|
|
||||||
|
assertGroupExists := func(name string) {
|
||||||
|
exists, err := client.Group.Query().Where(group.NameEQ(name), group.DeletedAtIsNil()).Exist(seedCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.True(t, exists, "expected group %s to exist", name)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertGroupExists(service.PlatformAnthropic + "-default")
|
||||||
|
assertGroupExists(service.PlatformOpenAI + "-default")
|
||||||
|
assertGroupExists(service.PlatformGemini + "-default")
|
||||||
|
assertGroupExists(service.PlatformAntigravity + "-default-1")
|
||||||
|
assertGroupExists(service.PlatformAntigravity + "-default-2")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEnsureSimpleModeDefaultGroups_IgnoresSoftDeletedGroups(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
tx := testEntTx(t)
|
||||||
|
client := tx.Client()
|
||||||
|
|
||||||
|
seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
// Create and then soft-delete an anthropic default group.
|
||||||
|
g, err := client.Group.Create().
|
||||||
|
SetName(service.PlatformAnthropic + "-default").
|
||||||
|
SetPlatform(service.PlatformAnthropic).
|
||||||
|
SetStatus(service.StatusActive).
|
||||||
|
SetSubscriptionType(service.SubscriptionTypeStandard).
|
||||||
|
SetRateMultiplier(1.0).
|
||||||
|
SetIsExclusive(false).
|
||||||
|
Save(seedCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
_, err = client.Group.Delete().Where(group.IDEQ(g.ID)).Exec(seedCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
|
||||||
|
|
||||||
|
// New active one should exist.
|
||||||
|
count, err := client.Group.Query().Where(group.NameEQ(service.PlatformAnthropic+"-default"), group.DeletedAtIsNil()).Count(seedCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, 1, count)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEnsureSimpleModeDefaultGroups_AntigravityNeedsTwoGroupsOnlyByCount(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
tx := testEntTx(t)
|
||||||
|
client := tx.Client()
|
||||||
|
|
||||||
|
seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
mustCreateGroup(t, client, &service.Group{Name: "ag-custom-1-" + time.Now().Format(time.RFC3339Nano), Platform: service.PlatformAntigravity})
|
||||||
|
mustCreateGroup(t, client, &service.Group{Name: "ag-custom-2-" + time.Now().Format(time.RFC3339Nano), Platform: service.PlatformAntigravity})
|
||||||
|
|
||||||
|
require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
|
||||||
|
|
||||||
|
count, err := client.Group.Query().Where(group.PlatformEQ(service.PlatformAntigravity), group.DeletedAtIsNil()).Count(seedCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.GreaterOrEqual(t, count, 2)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user