mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-08 08:55:22 +00:00
fix: reduce token refresh retry amplification
This commit is contained in:
1 parent
b56c501d30
commit
74199b6a6e
7 files changed
+386
-14
No files matched your search
@@ -649,6 +649,58 @@ func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account,
|
||||
return r.accountsToService(ctx, accounts)
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]service.Account, error) {
|
||||
if r.sql == nil {
|
||||
return nil, errors.New("account repository SQL executor not configured")
|
||||
}
|
||||
rows, err := r.sql.QueryContext(ctx, `
|
||||
SELECT id
|
||||
FROM accounts
|
||||
WHERE deleted_at IS NULL
|
||||
AND status = 'active'
|
||||
AND type = 'oauth'
|
||||
AND platform IN ('anthropic', 'openai', 'gemini', 'antigravity')
|
||||
AND credentials ? 'refresh_token'
|
||||
AND btrim(credentials->>'refresh_token') <> ''
|
||||
AND NOT (
|
||||
temp_unschedulable_until > NOW()
|
||||
AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'
|
||||
)
|
||||
ORDER BY priority ASC, id ASC
|
||||
`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
|
||||
accounts, err := r.GetByIDs(ctx, ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]service.Account, 0, len(accounts))
|
||||
for _, account := range accounts {
|
||||
if account != nil {
|
||||
out = append(out, *account)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
|
||||
accounts, err := r.client.Account.Query().
|
||||
Where(
|
||||
@@ -1148,7 +1200,7 @@ func (r *accountRepository) SetOverloaded(ctx context.Context, id int64, until t
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
||||
_, err := r.sql.ExecContext(ctx, `
|
||||
result, err := r.sql.ExecContext(ctx, `
|
||||
UPDATE accounts
|
||||
SET temp_unschedulable_until = $1,
|
||||
temp_unschedulable_reason = $2,
|
||||
@@ -1160,6 +1212,13 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64,
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected <= 0 {
|
||||
return nil
|
||||
}
|
||||
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
|
||||
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v", id, err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sqlmock "github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAccountRepository_SetTempUnschedulable_NoRowsAffectedDoesNotWriteOutbox(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
until := time.Now().Add(10 * time.Minute)
|
||||
|
||||
err := repo.SetTempUnschedulable(context.Background(), 42, until, "retry")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, exec.execQueries, 1)
|
||||
require.Contains(t, exec.execQueries[0], "UPDATE accounts")
|
||||
require.NotContains(t, strings.Join(exec.execQueries, "\n"), "scheduler_outbox")
|
||||
}
|
||||
|
||||
func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
var capturedSQL string
|
||||
mock.ExpectQuery("SELECT id").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"})).
|
||||
WillDelayFor(0)
|
||||
|
||||
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
|
||||
|
||||
accounts, err := repo.ListOAuthRefreshCandidates(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, accounts)
|
||||
|
||||
normalized := normalizeSQLWhitespace(capturedSQL)
|
||||
require.Contains(t, normalized, "deleted_at IS NULL")
|
||||
require.Contains(t, normalized, "status = 'active'")
|
||||
require.Contains(t, normalized, "type = 'oauth'")
|
||||
require.Contains(t, normalized, "platform IN ('anthropic', 'openai', 'gemini', 'antigravity')")
|
||||
require.Contains(t, normalized, "credentials ? 'refresh_token'")
|
||||
require.Contains(t, normalized, "btrim(credentials->>'refresh_token') <> ''")
|
||||
require.Contains(t, normalized, "temp_unschedulable_until > NOW()")
|
||||
require.Contains(t, normalized, "temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'")
|
||||
require.Contains(t, normalized, "ORDER BY priority ASC, id ASC")
|
||||
require.NotContains(t, normalized, "credentials->>'expires_at'")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
type captureQuerySQL struct {
|
||||
db *sql.DB
|
||||
captured *string
|
||||
}
|
||||
|
||||
func (c captureQuerySQL) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
||||
return c.db.ExecContext(ctx, query, args...)
|
||||
}
|
||||
|
||||
func (c captureQuerySQL) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
|
||||
if c.captured != nil {
|
||||
*c.captured = query
|
||||
}
|
||||
return c.db.QueryContext(ctx, query, args...)
|
||||
}
|
||||
|
||||
func normalizeSQLWhitespace(sql string) string {
|
||||
return strings.Join(regexp.MustCompile(`\s+`).Split(strings.TrimSpace(sql), -1), " ")
|
||||
}
|
||||
|
||||
type rowsAffectedResult int64
|
||||
|
||||
func (r rowsAffectedResult) LastInsertId() (int64, error) { return 0, nil }
|
||||
func (r rowsAffectedResult) RowsAffected() (int64, error) { return int64(r), nil }
|
||||
|
||||
type recordingSQLExecutor struct {
|
||||
result sql.Result
|
||||
err error
|
||||
execQueries []string
|
||||
}
|
||||
|
||||
func (e *recordingSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
||||
e.execQueries = append(e.execQueries, query)
|
||||
if e.err != nil {
|
||||
return nil, e.err
|
||||
}
|
||||
return e.result, nil
|
||||
}
|
||||
|
||||
func (e *recordingSQLExecutor) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
Reference in new issue
Block a user