fix account list parameter limit

2026-09-24 16:05:44 +08:00 · 2026-06-16 11:39:15 +08:00
parent 62fef6ebed
commit 8b698ff4c1
2 changed files with 250 additions and 28 deletions
@@ -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().
@@ -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
+}