fix(groups): merge main and resolve model allowlist conflicts

Resolve three conflicts from main's #6680 (supplementUnmappedOpenAIModels):

1. gateway_handler.go — keep PR's allowlist refactoring (ModelAllowlistEnabled,
   writeAllowlistedModelsList, modelListingSource); main's OpenAI branch in
   writeModelsList was auto-merged.
2. group_models_list.go (modify/delete) — delete file per PR intent; relocate
   supplementUnmappedOpenAIModels to group_model_allowlist.go where gateway_service
   and openai_codex_models_service can still call it.
3. openai_codex_models_service.go — combine main's supplementUnmappedOpenAIModels
   wrapping with PR's ModelAllowlistEnabled() guard.

Also fix two auto-merged test files that still referenced the old
GroupModelsListConfig / ModelsListConfig types.
This commit is contained in:
shaw committed 2026-09-07 19:44:03 +08:00
commit 2031c1d5ad
92 files changed
+3056 -299

No files matched your search

+79 -16
View File
@@ -2,9 +2,14 @@ package repository
import (
"context"
"database/sql"
"database/sql/driver"
"errors"
"fmt"
"io"
"os/exec"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -12,16 +17,39 @@ import (
// PgDumper implements service.DBDumper using pg_dump/psql
type PgDumper struct {
cfg *config.DatabaseConfig
cfg *config.DatabaseConfig
db *sql.DB
commandContext func(context.Context, string, ...string) *exec.Cmd
}
// NewPgDumper creates a new PgDumper
func NewPgDumper(cfg *config.Config) service.DBDumper {
return &PgDumper{cfg: &cfg.Database}
func NewPgDumper(cfg *config.Config, db *sql.DB) service.DBDumper {
return &PgDumper{
cfg: &cfg.Database,
db: db,
commandContext: exec.CommandContext,
}
}
// Dump executes pg_dump and returns a streaming reader of the output
func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
if d.db == nil {
return nil, errors.New("acquire backup migration lock: nil sql db")
}
lockConn, err := d.db.Conn(ctx)
if err != nil {
return nil, fmt.Errorf("acquire backup migration lock connection: %w", err)
}
if err := pgAdvisoryLock(ctx, lockConn); err != nil {
// Lock acquisition can fail after PostgreSQL accepted the request. Discard
// the session because ownership is ambiguous and advisory locks are session-scoped.
discardSQLConnection(lockConn)
return nil, fmt.Errorf("acquire backup migration lock: %w", err)
}
releaseLock := func() error {
return releaseBackupMigrationLock(lockConn)
}
args := []string{
"-h", d.cfg.Host,
"-p", fmt.Sprintf("%d", d.cfg.Port),
@@ -33,7 +61,11 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
"--if-exists",
}
cmd := exec.CommandContext(ctx, "pg_dump", args...)
commandContext := d.commandContext
if commandContext == nil {
commandContext = exec.CommandContext
}
cmd := commandContext(ctx, "pg_dump", args...)
if d.cfg.Password != "" {
cmd.Env = append(cmd.Environ(), "PGPASSWORD="+d.cfg.Password)
}
@@ -43,15 +75,34 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, fmt.Errorf("create stdout pipe: %w", err)
return nil, errors.Join(fmt.Errorf("create stdout pipe: %w", err), releaseLock())
}
if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("start pg_dump: %w", err)
_ = stdout.Close()
return nil, errors.Join(fmt.Errorf("start pg_dump: %w", err), releaseLock())
}
// 返回一个 ReadCloser:读 stdout,关闭时等待进程退出
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd, release: releaseLock}, nil
}
func releaseBackupMigrationLock(conn *sql.Conn) error {
unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := pgAdvisoryUnlock(unlockCtx, conn); err != nil {
// A failed unlock must not return a possibly lock-owning session to the pool.
discardSQLConnection(conn)
return fmt.Errorf("release backup migration lock: %w", err)
}
if err := conn.Close(); err != nil {
return fmt.Errorf("close backup migration lock connection: %w", err)
}
return nil
}
func discardSQLConnection(conn *sql.Conn) {
_ = conn.Raw(func(any) error { return driver.ErrBadConn })
_ = conn.Close()
}
// Restore executes psql to restore from a streaming reader
@@ -84,15 +135,27 @@ func (d *PgDumper) Restore(ctx context.Context, data io.Reader) error {
// cmdReadCloser wraps a command stdout pipe and waits for the process on Close
type cmdReadCloser struct {
io.ReadCloser
cmd *exec.Cmd
cmd *exec.Cmd
release func() error
closeOnce sync.Once
closeErr error
}
func (c *cmdReadCloser) Close() error {
// Close the pipe first
_ = c.ReadCloser.Close()
// Wait for the process to exit
if err := c.cmd.Wait(); err != nil {
return fmt.Errorf("pg_dump exited with error: %w", err)
}
return nil
c.closeOnce.Do(func() {
var closeErrs []error
if err := c.ReadCloser.Close(); err != nil {
closeErrs = append(closeErrs, fmt.Errorf("close pg_dump stdout: %w", err))
}
if err := c.cmd.Wait(); err != nil {
closeErrs = append(closeErrs, fmt.Errorf("pg_dump exited with error: %w", err))
}
if c.release != nil {
if err := c.release(); err != nil {
closeErrs = append(closeErrs, err)
}
}
c.closeErr = errors.Join(closeErrs...)
})
return c.closeErr
}
@@ -0,0 +1,157 @@
package repository
import (
"context"
"errors"
"io"
"os/exec"
"regexp"
"testing"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
func newTestPgDumper(t *testing.T, commandContext func(context.Context, string, ...string) *exec.Cmd) (*PgDumper, sqlmock.Sqlmock) {
t.Helper()
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
return &PgDumper{
cfg: &config.DatabaseConfig{
Host: "db.example.test",
Port: 5432,
User: "sub2api",
Password: "secret",
DBName: "sub2api",
SSLMode: "require",
},
db: db,
commandContext: commandContext,
}, mock
}
func expectBackupMigrationLock(mock sqlmock.Sqlmock) {
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
WithArgs(migrationsAdvisoryLockID).
WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(true))
}
func expectBackupMigrationUnlock(mock sqlmock.Sqlmock) {
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
WithArgs(migrationsAdvisoryLockID).
WillReturnResult(sqlmock.NewResult(0, 1))
}
func TestPgDumperHoldsMigrationLockThroughReaderClose(t *testing.T) {
var mock sqlmock.Sqlmock
commandCreated := false
dumper, createdMock := newTestPgDumper(t, func(ctx context.Context, name string, args ...string) *exec.Cmd {
commandCreated = true
require.Equal(t, "pg_dump", name)
require.Contains(t, args, "--clean")
require.NoError(t, mock.ExpectationsWereMet(), "migration lock must be acquired before pg_dump is created")
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
})
mock = createdMock
expectBackupMigrationLock(mock)
reader, err := dumper.Dump(context.Background())
require.NoError(t, err)
require.True(t, commandCreated)
data, err := io.ReadAll(reader)
require.NoError(t, err)
require.Equal(t, "backup-data", string(data))
expectBackupMigrationUnlock(mock)
require.Error(t, mock.ExpectationsWereMet(), "migration lock was released before the reader closed")
require.NoError(t, reader.Close())
require.NoError(t, mock.ExpectationsWereMet())
require.NoError(t, reader.Close(), "reader close must be idempotent")
}
func TestPgDumperReleasesMigrationLockWhenStdoutPipeSetupFails(t *testing.T) {
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
cmd := exec.CommandContext(ctx, "sh", "-c", "true")
cmd.Stdout = io.Discard
return cmd
})
expectBackupMigrationLock(mock)
expectBackupMigrationUnlock(mock)
reader, err := dumper.Dump(context.Background())
require.Nil(t, reader)
require.ErrorContains(t, err, "create stdout pipe")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestPgDumperReleasesMigrationLockWhenProcessStartFails(t *testing.T) {
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
return exec.CommandContext(ctx, "/path/that/does/not/exist/pg_dump")
})
expectBackupMigrationLock(mock)
expectBackupMigrationUnlock(mock)
reader, err := dumper.Dump(context.Background())
require.Nil(t, reader)
require.ErrorContains(t, err, "start pg_dump")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestPgDumperReleasesMigrationLockWhenProcessFails(t *testing.T) {
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
return exec.CommandContext(ctx, "sh", "-c", "printf partial-backup; exit 7")
})
expectBackupMigrationLock(mock)
reader, err := dumper.Dump(context.Background())
require.NoError(t, err)
data, err := io.ReadAll(reader)
require.NoError(t, err)
require.Equal(t, "partial-backup", string(data))
expectBackupMigrationUnlock(mock)
require.ErrorContains(t, reader.Close(), "pg_dump exited with error")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestPgDumperReportsUnlockFailureAndDiscardsConnection(t *testing.T) {
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
})
expectBackupMigrationLock(mock)
reader, err := dumper.Dump(context.Background())
require.NoError(t, err)
_, err = io.ReadAll(reader)
require.NoError(t, err)
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
WithArgs(migrationsAdvisoryLockID).
WillReturnError(errors.New("unlock unavailable"))
require.ErrorContains(t, reader.Close(), "release backup migration lock")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestPgDumperDoesNotStartProcessWhenMigrationLockFails(t *testing.T) {
commandCreated := false
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
commandCreated = true
return exec.CommandContext(ctx, "sh", "-c", "true")
})
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
WithArgs(migrationsAdvisoryLockID).
WillReturnError(errors.New("database unavailable"))
reader, err := dumper.Dump(context.Background())
require.Nil(t, reader)
require.ErrorContains(t, err, "acquire backup migration lock")
require.False(t, commandCreated)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestPgDumperRejectsNilDatabase(t *testing.T) {
dumper := &PgDumper{cfg: &config.DatabaseConfig{}}
reader, err := dumper.Dump(context.Background())
require.Nil(t, reader)
require.ErrorContains(t, err, "nil sql db")
}
+19 -7
View File
@@ -10,11 +10,22 @@ import (
)
const (
redeemRateLimitKeyPrefix = "redeem:ratelimit:"
redeemRateLimitKeyPrefix = "redeem:ratelimit:v2:"
redeemLockKeyPrefix = "redeem:lock:"
redeemRateLimitDuration = 24 * time.Hour
redeemRateLimitWindow = 10 * time.Minute
)
var incrementRedeemAttemptScript = redis.NewScript(`
local current = redis.call('INCR', KEYS[1])
local ttl = redis.call('PTTL', KEYS[1])
if current == 1 or ttl == -1 then
redis.call('PEXPIRE', KEYS[1], ARGV[1])
end
return current
`)
// redeemRateLimitKey generates the Redis key for redeem attempt rate limiting.
func redeemRateLimitKey(userID int64) string {
return fmt.Sprintf("%s%d", redeemRateLimitKeyPrefix, userID)
@@ -44,11 +55,12 @@ func (c *redeemCache) GetRedeemAttemptCount(ctx context.Context, userID int64) (
func (c *redeemCache) IncrementRedeemAttemptCount(ctx context.Context, userID int64) error {
key := redeemRateLimitKey(userID)
pipe := c.rdb.Pipeline()
pipe.Incr(ctx, key)
pipe.Expire(ctx, key, redeemRateLimitDuration)
_, err := pipe.Exec(ctx)
return err
return incrementRedeemAttemptScript.Run(
ctx,
c.rdb,
[]string{key},
redeemRateLimitWindow.Milliseconds(),
).Err()
}
func (c *redeemCache) AcquireRedeemLock(ctx context.Context, code string, ttl time.Duration) (bool, error) {
@@ -39,7 +39,35 @@ func (s *RedeemCacheSuite) TestIncrementAndGetRedeemAttemptCount() {
ttl, err := s.rdb.TTL(s.ctx, key).Result()
require.NoError(s.T(), err, "TTL")
s.AssertTTLWithin(ttl, 1*time.Second, redeemRateLimitDuration)
s.AssertTTLWithin(ttl, redeemRateLimitWindow-time.Second, redeemRateLimitWindow)
}
func (s *RedeemCacheSuite) TestIncrementDoesNotRefreshFixedWindow() {
userID := int64(3)
key := redeemRateLimitKey(userID)
shortTTL := 5 * time.Minute
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
require.NoError(s.T(), s.rdb.PExpire(s.ctx, key, shortTTL).Err())
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
ttl, err := s.rdb.PTTL(s.ctx, key).Result()
require.NoError(s.T(), err)
s.AssertTTLWithin(ttl, shortTTL-time.Second, shortTTL)
}
func (s *RedeemCacheSuite) TestIncrementRepairsMissingTTL() {
userID := int64(4)
key := redeemRateLimitKey(userID)
require.NoError(s.T(), s.rdb.Set(s.ctx, key, 5, 0).Err())
ttl, err := s.rdb.PTTL(s.ctx, key).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), time.Duration(-1), ttl)
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
count, err := s.cache.GetRedeemAttemptCount(s.ctx, userID)
require.NoError(s.T(), err)
require.Equal(s.T(), 6, count)
ttl, err = s.rdb.PTTL(s.ctx, key).Result()
require.NoError(s.T(), err)
s.AssertTTLWithin(ttl, redeemRateLimitWindow-time.Second, redeemRateLimitWindow)
}
func (s *RedeemCacheSuite) TestMultipleIncrements() {
@@ -18,22 +18,22 @@ func TestRedeemRateLimitKey(t *testing.T) {
{
name: "normal_user_id",
userID: 123,
expected: "redeem:ratelimit:123",
expected: "redeem:ratelimit:v2:123",
},
{
name: "zero_user_id",
userID: 0,
expected: "redeem:ratelimit:0",
expected: "redeem:ratelimit:v2:0",
},
{
name: "negative_user_id",
userID: -1,
expected: "redeem:ratelimit:-1",
expected: "redeem:ratelimit:v2:-1",
},
{
name: "max_int64",
userID: math.MaxInt64,
expected: "redeem:ratelimit:9223372036854775807",
expected: "redeem:ratelimit:v2:9223372036854775807",
},
}