diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 4e19b498ba..e7626696bf 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -64,6 +64,8 @@ var schedulerNeutralExtraKeys = map[string]struct{}{ "session_window_utilization": {}, } +const postgresParameterBatchSize = 50000 + // NewAccountRepository 创建账户仓储实例。 // 这是对外暴露的构造函数，返回接口类型以便于依赖注入。 func NewAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AccountRepository { @@ -1662,17 +1664,23 @@ func notExpiredPredicate(now time.Time) dbpredicate.Account { func (r *accountRepository) loadProxies(ctx context.Context, proxyIDs []int64) (map[int64]*service.Proxy, error) { proxyMap := make(map[int64]*service.Proxy) + proxyIDs = uniquePositiveInt64s(proxyIDs) if len(proxyIDs) == 0 { return proxyMap, nil } - proxies, err := r.client.Proxy.Query().Where(dbproxy.IDIn(proxyIDs...)).All(ctx) - if err != nil { - return nil, err - } - - for _, p := range proxies { - proxyMap[p.ID] = proxyEntityToService(p) + for start := 0; start < len(proxyIDs); start += postgresParameterBatchSize { + end := start + postgresParameterBatchSize + if end > len(proxyIDs) { + end = len(proxyIDs) + } + proxies, err := r.client.Proxy.Query().Where(dbproxy.IDIn(proxyIDs[start:end]...)).All(ctx) + if err != nil { + return nil, err + } + for _, p := range proxies { + proxyMap[p.ID] = proxyEntityToService(p) + } } return proxyMap, nil } @@ -1682,38 +1690,94 @@ func (r *accountRepository) loadAccountGroups(ctx context.Context, accountIDs [] groupIDsByAccount := make(map[int64][]int64) accountGroupsByAccount := make(map[int64][]service.AccountGroup) + accountIDs = uniquePositiveInt64s(accountIDs) if len(accountIDs) == 0 { return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil } - entries, err := r.client.AccountGroup.Query(). - Where(dbaccountgroup.AccountIDIn(accountIDs...)). - WithGroup(). - Order(dbaccountgroup.ByAccountID(), dbaccountgroup.ByPriority()). - All(ctx) - if err != nil { - return nil, nil, nil, err - } - - for _, ag := range entries { - groupSvc := groupEntityToService(ag.Edges.Group) - agSvc := service.AccountGroup{ - AccountID: ag.AccountID, - GroupID: ag.GroupID, - Priority: ag.Priority, - CreatedAt: ag.CreatedAt, - Group: groupSvc, + for start := 0; start < len(accountIDs); start += postgresParameterBatchSize { + end := start + postgresParameterBatchSize + if end > len(accountIDs) { + end = len(accountIDs) } - accountGroupsByAccount[ag.AccountID] = append(accountGroupsByAccount[ag.AccountID], agSvc) - groupIDsByAccount[ag.AccountID] = append(groupIDsByAccount[ag.AccountID], ag.GroupID) - if groupSvc != nil { - groupsByAccount[ag.AccountID] = append(groupsByAccount[ag.AccountID], groupSvc) + entries, err := r.client.AccountGroup.Query(). + Where(dbaccountgroup.AccountIDIn(accountIDs[start:end]...)). + Order(dbaccountgroup.ByAccountID(), dbaccountgroup.ByPriority()). + All(ctx) + if err != nil { + return nil, nil, nil, err + } + groupIDs := make([]int64, 0, len(entries)) + for _, ag := range entries { + groupIDs = append(groupIDs, ag.GroupID) + } + groupMap, err := r.loadGroups(ctx, groupIDs) + if err != nil { + return nil, nil, nil, err + } + + for _, ag := range entries { + groupSvc := groupMap[ag.GroupID] + agSvc := service.AccountGroup{ + AccountID: ag.AccountID, + GroupID: ag.GroupID, + Priority: ag.Priority, + CreatedAt: ag.CreatedAt, + Group: groupSvc, + } + accountGroupsByAccount[ag.AccountID] = append(accountGroupsByAccount[ag.AccountID], agSvc) + groupIDsByAccount[ag.AccountID] = append(groupIDsByAccount[ag.AccountID], ag.GroupID) + if groupSvc != nil { + groupsByAccount[ag.AccountID] = append(groupsByAccount[ag.AccountID], groupSvc) + } } } return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil } +func (r *accountRepository) loadGroups(ctx context.Context, groupIDs []int64) (map[int64]*service.Group, error) { + groupMap := make(map[int64]*service.Group) + groupIDs = uniquePositiveInt64s(groupIDs) + if len(groupIDs) == 0 { + return groupMap, nil + } + + for start := 0; start < len(groupIDs); start += postgresParameterBatchSize { + end := start + postgresParameterBatchSize + if end > len(groupIDs) { + end = len(groupIDs) + } + groups, err := r.client.Group.Query().Where(dbgroup.IDIn(groupIDs[start:end]...)).All(ctx) + if err != nil { + return nil, err + } + for _, g := range groups { + groupMap[g.ID] = groupEntityToService(g) + } + } + return groupMap, nil +} + +func uniquePositiveInt64s(ids []int64) []int64 { + if len(ids) == 0 { + return nil + } + out := make([]int64, 0, len(ids)) + seen := make(map[int64]struct{}, len(ids)) + for _, id := range ids { + if id <= 0 { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + return out +} + func (r *accountRepository) loadAccountGroupIDs(ctx context.Context, accountID int64) ([]int64, error) { entries, err := r.client.AccountGroup. Query(). diff --git a/backend/internal/repository/account_repo_test.go b/backend/internal/repository/account_repo_test.go new file mode 100644 index 0000000000..5de74364df --- /dev/null +++ b/backend/internal/repository/account_repo_test.go @@ -0,0 +1,158 @@ +package repository + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + "io" + "sync" + "testing" + + dbent "github.com/Wei-Shaw/sub2api/ent" + _ "github.com/Wei-Shaw/sub2api/ent/runtime" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" + + "entgo.io/ent/dialect" + entsql "entgo.io/ent/dialect/sql" +) + +const parameterLimitTestDriverName = "sub2api_param_limit_test" + +var registerParameterLimitTestDriverOnce sync.Once + +func TestAccountsToService_LargeActiveAccountSetDoesNotExceedPostgresParameterLimit(t *testing.T) { + repo := newParameterLimitAccountRepo(t) + + accounts := make([]*dbent.Account, 0, 65536) + for i := range 65536 { + accounts = append(accounts, &dbent.Account{ + ID: int64(i + 1), + Name: "large-active", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: map[string]any{}, + Extra: map[string]any{}, + Status: service.StatusActive, + Schedulable: true, + }) + } + + got, err := repo.accountsToService(context.Background(), accounts) + require.NoError(t, err) + require.Len(t, got, len(accounts)) +} + +func newParameterLimitAccountRepo(t *testing.T) *accountRepository { + t.Helper() + + registerParameterLimitTestDriverOnce.Do(func() { + sql.Register(parameterLimitTestDriverName, parameterLimitDriver{}) + }) + + db, err := sql.Open(parameterLimitTestDriverName, "") + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + drv := entsql.OpenDB(dialect.Postgres, db) + client := dbent.NewClient(dbent.Driver(drv)) + t.Cleanup(func() { _ = client.Close() }) + + return newAccountRepositoryWithSQL(client, nil, nil) +} + +type parameterLimitDriver struct{} + +func (parameterLimitDriver) Open(string) (driver.Conn, error) { + return parameterLimitConn{}, nil +} + +type parameterLimitConn struct{} + +func (parameterLimitConn) Prepare(query string) (driver.Stmt, error) { + return parameterLimitStmt{query: query}, nil +} + +func (parameterLimitConn) Close() error { + return nil +} + +func (parameterLimitConn) Begin() (driver.Tx, error) { + return parameterLimitTx{}, nil +} + +func (parameterLimitConn) QueryContext(_ context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + return queryWithParameterLimit(query, args) +} + +type parameterLimitStmt struct { + query string +} + +func (s parameterLimitStmt) Close() error { + return nil +} + +func (s parameterLimitStmt) NumInput() int { + return -1 +} + +func (s parameterLimitStmt) Exec(args []driver.Value) (driver.Result, error) { + return driver.RowsAffected(0), parameterLimitError(len(args)) +} + +func (s parameterLimitStmt) Query(args []driver.Value) (driver.Rows, error) { + namedArgs := make([]driver.NamedValue, len(args)) + for i, arg := range args { + namedArgs[i] = driver.NamedValue{Ordinal: i + 1, Value: arg} + } + return queryWithParameterLimit(s.query, namedArgs) +} + +type parameterLimitTx struct{} + +func (parameterLimitTx) Commit() error { + return nil +} + +func (parameterLimitTx) Rollback() error { + return nil +} + +func queryWithParameterLimit(query string, args []driver.NamedValue) (driver.Rows, error) { + if err := parameterLimitError(len(args)); err != nil { + return nil, err + } + return parameterLimitRows{columns: columnsForParameterLimitQuery(query)}, nil +} + +func parameterLimitError(paramCount int) error { + if paramCount <= 65535 { + return nil + } + return fmt.Errorf("pq: got %d parameters but PostgreSQL only supports 65535 parameters", paramCount) +} + +func columnsForParameterLimitQuery(query string) []string { + if query == "" { + return nil + } + return []string{"account_id", "group_id", "priority", "created_at"} +} + +type parameterLimitRows struct { + columns []string +} + +func (r parameterLimitRows) Columns() []string { + return r.columns +} + +func (parameterLimitRows) Close() error { + return nil +} + +func (parameterLimitRows) Next([]driver.Value) error { + return io.EOF +}