Skip to content

GitLab

"frontend/src/i18n/git@web.lueluesay.top:chenxi/sub2api.git" did not exist on "4838ab74b3a040f54d9175c986c82415d4a1c9c1"
Commit 3d79773b authored by kyx236's avatar kyx236
Browse files

Merge branch 'main' of https://github.com/james-6-23/sub2api

parents 6aa8cbbf 742e73c9
Hide whitespace changes
Inline Side-by-side
backend/internal/repository/billing_cache_integration_test.go
View file @ 3d79773b
...@@ -278,6 +278,90 @@ func (s *BillingCacheSuite) TestSubscriptionCache() { ...@@ -278,6 +278,90 @@ func (s *BillingCacheSuite) TestSubscriptionCache() {
} }
} }
// TestDeductUserBalance_ErrorPropagation 验证 P2-12 修复：
// Redis 真实错误应传播，key 不存在（redis.Nil）应返回 nil。
func (s *BillingCacheSuite) TestDeductUserBalance_ErrorPropagation() {
tests := []struct {
name string
fn func(ctx context.Context, cache service.BillingCache)
expectErr bool
}{
{
name: "key_not_exists_returns_nil",
fn: func(ctx context.Context, cache service.BillingCache) {
// key 不存在时，Lua 脚本返回 0（redis.Nil），应返回 nil 而非错误
err := cache.DeductUserBalance(ctx, 99999, 1.0)
require.NoError(s.T(), err, "DeductUserBalance on non-existent key should return nil")
},
},
{
name: "existing_key_deducts_successfully",
fn: func(ctx context.Context, cache service.BillingCache) {
require.NoError(s.T(), cache.SetUserBalance(ctx, 200, 50.0))
err := cache.DeductUserBalance(ctx, 200, 10.0)
require.NoError(s.T(), err, "DeductUserBalance should succeed")
bal, err := cache.GetUserBalance(ctx, 200)
require.NoError(s.T(), err)
require.Equal(s.T(), 40.0, bal, "余额应为 40.0")
},
},
{
name: "cancelled_context_propagates_error",
fn: func(ctx context.Context, cache service.BillingCache) {
require.NoError(s.T(), cache.SetUserBalance(ctx, 201, 50.0))
cancelCtx, cancel := context.WithCancel(ctx)
cancel() // 立即取消
err := cache.DeductUserBalance(cancelCtx, 201, 10.0)
require.Error(s.T(), err, "cancelled context should propagate error")
},
},
}
for _, tt := range tests {
s.Run(tt.name, func() {
rdb := testRedis(s.T())
cache := NewBillingCache(rdb)
ctx := context.Background()
tt.fn(ctx, cache)
})
}
}
// TestUpdateSubscriptionUsage_ErrorPropagation 验证 P2-12 修复：
// Redis 真实错误应传播，key 不存在（redis.Nil）应返回 nil。
func (s *BillingCacheSuite) TestUpdateSubscriptionUsage_ErrorPropagation() {
s.Run("key_not_exists_returns_nil", func() {
rdb := testRedis(s.T())
cache := NewBillingCache(rdb)
ctx := context.Background()
err := cache.UpdateSubscriptionUsage(ctx, 88888, 77777, 1.0)
require.NoError(s.T(), err, "UpdateSubscriptionUsage on non-existent key should return nil")
})
s.Run("cancelled_context_propagates_error", func() {
rdb := testRedis(s.T())
cache := NewBillingCache(rdb)
ctx := context.Background()
data := &service.SubscriptionCacheData{
Status: "active",
ExpiresAt: time.Now().Add(1 * time.Hour),
Version: 1,
}
require.NoError(s.T(), cache.SetSubscriptionCache(ctx, 301, 401, data))
cancelCtx, cancel := context.WithCancel(ctx)
cancel()
err := cache.UpdateSubscriptionUsage(cancelCtx, 301, 401, 1.0)
require.Error(s.T(), err, "cancelled context should propagate error")
})
}
func TestBillingCacheSuite(t *testing.T) { func TestBillingCacheSuite(t *testing.T) {
suite.Run(t, new(BillingCacheSuite)) suite.Run(t, new(BillingCacheSuite))
} }
backend/internal/repository/billing_cache_jitter_test.go 0 → 100644
View file @ 3d79773b
package repository
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// --- Task 6.1 验证: math/rand/v2 迁移后 jitteredTTL 行为正确 ---
func TestJitteredTTL_WithinExpectedRange(t *testing.T) {
// jitteredTTL 使用减法抖动: billingCacheTTL - [0, billingCacheJitter)
// 所以结果应在 [billingCacheTTL - billingCacheJitter, billingCacheTTL] 范围内
lowerBound := billingCacheTTL - billingCacheJitter // 5min - 30s = 4min30s
upperBound := billingCacheTTL // 5min
for i := 0; i < 200; i++ {
ttl := jitteredTTL()
assert.GreaterOrEqual(t, int64(ttl), int64(lowerBound),
"TTL 不应低于 %v，实际得到 %v", lowerBound, ttl)
assert.LessOrEqual(t, int64(ttl), int64(upperBound),
"TTL 不应超过 %v（上界不变保证），实际得到 %v", upperBound, ttl)
}
}
func TestJitteredTTL_NeverExceedsBase(t *testing.T) {
// 关键安全性测试：jitteredTTL 使用减法抖动，确保永远不超过 billingCacheTTL
for i := 0; i < 500; i++ {
ttl := jitteredTTL()
assert.LessOrEqual(t, int64(ttl), int64(billingCacheTTL),
"jitteredTTL 不应超过基础 TTL（上界预期不被打破）")
}
}
func TestJitteredTTL_HasVariance(t *testing.T) {
// 验证抖动确实产生了不同的值
results := make(map[time.Duration]bool)
for i := 0; i < 100; i++ {
ttl := jitteredTTL()
results[ttl] = true
}
require.Greater(t, len(results), 1,
"jitteredTTL 应产生不同的值（抖动生效），但 100 次调用结果全部相同")
}
func TestJitteredTTL_AverageNearCenter(t *testing.T) {
// 验证平均值大约在抖动范围中间
var sum time.Duration
runs := 1000
for i := 0; i < runs; i++ {
sum += jitteredTTL()
}
avg := sum / time.Duration(runs)
expectedCenter := billingCacheTTL - billingCacheJitter/2 // 4min45s
// 允许 ±5s 的误差
tolerance := 5 * time.Second
assert.InDelta(t, float64(expectedCenter), float64(avg), float64(tolerance),
"平均 TTL 应接近抖动范围中心 %v", expectedCenter)
}
func TestBillingKeyGeneration(t *testing.T) {
t.Run("balance_key", func(t *testing.T) {
key := billingBalanceKey(12345)
assert.Equal(t, "billing:balance:12345", key)
})
t.Run("sub_key", func(t *testing.T) {
key := billingSubKey(100, 200)
assert.Equal(t, "billing:sub:100:200", key)
})
}
func BenchmarkJitteredTTL(b *testing.B) {
for i := 0; i < b.N; i++ {
_ = jitteredTTL()
}
}
backend/internal/repository/billing_cache_test.go
View file @ 3d79773b
...@@ -5,6 +5,7 @@ package repository ...@@ -5,6 +5,7 @@ package repository
import ( import (
"math" "math"
"testing" "testing"
"time"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
...@@ -85,3 +86,26 @@ func TestBillingSubKey(t *testing.T) { ...@@ -85,3 +86,26 @@ func TestBillingSubKey(t *testing.T) {
}) })
} }
} }
func TestJitteredTTL(t *testing.T) {
const (
minTTL = 4*time.Minute + 30*time.Second // 270s = 5min - 30s
maxTTL = 5*time.Minute + 30*time.Second // 330s = 5min + 30s
)
for i := 0; i < 200; i++ {
ttl := jitteredTTL()
require.GreaterOrEqual(t, ttl, minTTL, "jitteredTTL() 返回值低于下限: %v", ttl)
require.LessOrEqual(t, ttl, maxTTL, "jitteredTTL() 返回值超过上限: %v", ttl)
}
}
func TestJitteredTTL_HasVariation(t *testing.T) {
// 多次调用应该产生不同的值（验证抖动存在）
seen := make(map[time.Duration]struct{}, 50)
for i := 0; i < 50; i++ {
seen[jitteredTTL()] = struct{}{}
}
// 50 次调用中应该至少有 2 个不同的值
require.Greater(t, len(seen), 1, "jitteredTTL() 应产生不同的 TTL 值")
}
backend/internal/repository/claude_oauth_service.go
View file @ 3d79773b
...@@ -4,13 +4,14 @@ import ( ...@@ -4,13 +4,14 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"log"
"net/http" "net/http"
"net/url" "net/url"
"strings" "strings"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/oauth" "github.com/Wei-Shaw/sub2api/internal/pkg/oauth"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/service" "github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/util/logredact" "github.com/Wei-Shaw/sub2api/internal/util/logredact"
...@@ -28,11 +29,14 @@ func NewClaudeOAuthClient() service.ClaudeOAuthClient { ...@@ -28,11 +29,14 @@ func NewClaudeOAuthClient() service.ClaudeOAuthClient {
type claudeOAuthService struct { type claudeOAuthService struct {
baseURL string baseURL string
tokenURL string tokenURL string
clientFactory func(proxyURL string) *req.Client clientFactory func(proxyURL string) (*req.Client, error)
} }
func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey, proxyURL string) (string, error) { func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey, proxyURL string) (string, error) {
client := s.clientFactory(proxyURL) client, err := s.clientFactory(proxyURL)
if err != nil {
return "", fmt.Errorf("create HTTP client: %w", err)
}
var orgs []struct { var orgs []struct {
UUID string `json:"uuid"` UUID string `json:"uuid"`
...@@ -41,7 +45,7 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey ...@@ -41,7 +45,7 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey
} }
targetURL := s.baseURL + "/api/organizations" targetURL := s.baseURL + "/api/organizations"
log.Printf("[OAuth] Step 1: Getting organization UUID from %s", targetURL) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1: Getting organization UUID from %s", targetURL)
resp, err := client.R(). resp, err := client.R().
SetContext(ctx). SetContext(ctx).
...@@ -53,11 +57,11 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey ...@@ -53,11 +57,11 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey
Get(targetURL) Get(targetURL)
if err != nil { if err != nil {
log.Printf("[OAuth] Step 1 FAILED - Request error: %v", err) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1 FAILED - Request error: %v", err)
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
log.Printf("[OAuth] Step 1 Response - Status: %d", resp.StatusCode) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1 Response - Status: %d", resp.StatusCode)
if !resp.IsSuccessState() { if !resp.IsSuccessState() {
return "", fmt.Errorf("failed to get organizations: status %d, body: %s", resp.StatusCode, resp.String()) return "", fmt.Errorf("failed to get organizations: status %d, body: %s", resp.StatusCode, resp.String())
...@@ -69,26 +73,29 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey ...@@ -69,26 +73,29 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey
// 如果只有一个组织，直接使用 // 如果只有一个组织，直接使用
if len(orgs) == 1 { if len(orgs) == 1 {
log.Printf("[OAuth] Step 1 SUCCESS - Single org found, UUID: %s, Name: %s", orgs[0].UUID, orgs[0].Name) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1 SUCCESS - Single org found, UUID: %s, Name: %s", orgs[0].UUID, orgs[0].Name)
return orgs[0].UUID, nil return orgs[0].UUID, nil
} }
// 如果有多个组织，优先选择 raven_type 为 "team" 的组织 // 如果有多个组织，优先选择 raven_type 为 "team" 的组织
for _, org := range orgs { for _, org := range orgs {
if org.RavenType != nil && *org.RavenType == "team" { if org.RavenType != nil && *org.RavenType == "team" {
log.Printf("[OAuth] Step 1 SUCCESS - Selected team org, UUID: %s, Name: %s, RavenType: %s", logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1 SUCCESS - Selected team org, UUID: %s, Name: %s, RavenType: %s",
org.UUID, org.Name, *org.RavenType) org.UUID, org.Name, *org.RavenType)
return org.UUID, nil return org.UUID, nil
} }
} }
// 如果没有 team 类型的组织，使用第一个 // 如果没有 team 类型的组织，使用第一个
log.Printf("[OAuth] Step 1 SUCCESS - No team org found, using first org, UUID: %s, Name: %s", orgs[0].UUID, orgs[0].Name) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 1 SUCCESS - No team org found, using first org, UUID: %s, Name: %s", orgs[0].UUID, orgs[0].Name)
return orgs[0].UUID, nil return orgs[0].UUID, nil
} }
func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKey, orgUUID, scope, codeChallenge, state, proxyURL string) (string, error) { func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKey, orgUUID, scope, codeChallenge, state, proxyURL string) (string, error) {
client := s.clientFactory(proxyURL) client, err := s.clientFactory(proxyURL)
if err != nil {
return "", fmt.Errorf("create HTTP client: %w", err)
}
authURL := fmt.Sprintf("%s/v1/oauth/%s/authorize", s.baseURL, orgUUID) authURL := fmt.Sprintf("%s/v1/oauth/%s/authorize", s.baseURL, orgUUID)
...@@ -103,9 +110,9 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe ...@@ -103,9 +110,9 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
"code_challenge_method": "S256", "code_challenge_method": "S256",
} }
log.Printf("[OAuth] Step 2: Getting authorization code from %s", authURL) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 2: Getting authorization code from %s", authURL)
reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody)) reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody))
log.Printf("[OAuth] Step 2 Request Body: %s", string(reqBodyJSON)) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 2 Request Body: %s", string(reqBodyJSON))
var result struct { var result struct {
RedirectURI string `json:"redirect_uri"` RedirectURI string `json:"redirect_uri"`
...@@ -128,11 +135,11 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe ...@@ -128,11 +135,11 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
Post(authURL) Post(authURL)
if err != nil { if err != nil {
log.Printf("[OAuth] Step 2 FAILED - Request error: %v", err) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 2 FAILED - Request error: %v", err)
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
log.Printf("[OAuth] Step 2 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes())) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 2 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes()))
if !resp.IsSuccessState() { if !resp.IsSuccessState() {
return "", fmt.Errorf("failed to get authorization code: status %d, body: %s", resp.StatusCode, resp.String()) return "", fmt.Errorf("failed to get authorization code: status %d, body: %s", resp.StatusCode, resp.String())
...@@ -160,12 +167,15 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe ...@@ -160,12 +167,15 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
fullCode = authCode + "#" + responseState fullCode = authCode + "#" + responseState
} }
log.Printf("[OAuth] Step 2 SUCCESS - Got authorization code") logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 2 SUCCESS - Got authorization code")
return fullCode, nil return fullCode, nil
} }
func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string, isSetupToken bool) (*oauth.TokenResponse, error) { func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, codeVerifier, state, proxyURL string, isSetupToken bool) (*oauth.TokenResponse, error) {
client := s.clientFactory(proxyURL) client, err := s.clientFactory(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
// Parse code which may contain state in format "authCode#state" // Parse code which may contain state in format "authCode#state"
authCode := code authCode := code
...@@ -192,9 +202,9 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod ...@@ -192,9 +202,9 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
reqBody["expires_in"] = 31536000 // 365 * 24 * 60 * 60 seconds reqBody["expires_in"] = 31536000 // 365 * 24 * 60 * 60 seconds
} }
log.Printf("[OAuth] Step 3: Exchanging code for token at %s", s.tokenURL) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 3: Exchanging code for token at %s", s.tokenURL)
reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody)) reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody))
log.Printf("[OAuth] Step 3 Request Body: %s", string(reqBodyJSON)) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 3 Request Body: %s", string(reqBodyJSON))
var tokenResp oauth.TokenResponse var tokenResp oauth.TokenResponse
...@@ -208,22 +218,25 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod ...@@ -208,22 +218,25 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
Post(s.tokenURL) Post(s.tokenURL)
if err != nil { if err != nil {
log.Printf("[OAuth] Step 3 FAILED - Request error: %v", err) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 3 FAILED - Request error: %v", err)
return nil, fmt.Errorf("request failed: %w", err) return nil, fmt.Errorf("request failed: %w", err)
} }
log.Printf("[OAuth] Step 3 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes())) logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 3 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes()))
if !resp.IsSuccessState() { if !resp.IsSuccessState() {
return nil, fmt.Errorf("token exchange failed: status %d, body: %s", resp.StatusCode, resp.String()) return nil, fmt.Errorf("token exchange failed: status %d, body: %s", resp.StatusCode, resp.String())
} }
log.Printf("[OAuth] Step 3 SUCCESS - Got access token") logger.LegacyPrintf("repository.claude_oauth", "[OAuth] Step 3 SUCCESS - Got access token")
return &tokenResp, nil return &tokenResp, nil
} }
func (s *claudeOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*oauth.TokenResponse, error) { func (s *claudeOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*oauth.TokenResponse, error) {
client := s.clientFactory(proxyURL) client, err := s.clientFactory(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
reqBody := map[string]any{ reqBody := map[string]any{
"grant_type": "refresh_token", "grant_type": "refresh_token",
...@@ -253,16 +266,20 @@ func (s *claudeOAuthService) RefreshToken(ctx context.Context, refreshToken, pro ...@@ -253,16 +266,20 @@ func (s *claudeOAuthService) RefreshToken(ctx context.Context, refreshToken, pro
return &tokenResp, nil return &tokenResp, nil
} }
func createReqClient(proxyURL string) *req.Client { func createReqClient(proxyURL string) (*req.Client, error) {
// 禁用 CookieJar，确保每次授权都是干净的会话 // 禁用 CookieJar，确保每次授权都是干净的会话
client := req.C(). client := req.C().
SetTimeout(60 * time.Second). SetTimeout(60 * time.Second).
ImpersonateChrome(). ImpersonateChrome().
SetCookieJar(nil) // 禁用 CookieJar SetCookieJar(nil) // 禁用 CookieJar
if strings.TrimSpace(proxyURL) != "" { trimmed, _, err := proxyurl.Parse(proxyURL)
client.SetProxyURL(strings.TrimSpace(proxyURL)) if err != nil {
return nil, err
}
if trimmed != "" {
client.SetProxyURL(trimmed)
} }
return client return client, nil
} }
backend/internal/repository/claude_oauth_service_test.go
View file @ 3d79773b
...@@ -91,7 +91,7 @@ func (s *ClaudeOAuthServiceSuite) TestGetOrganizationUUID() { ...@@ -91,7 +91,7 @@ func (s *ClaudeOAuthServiceSuite) TestGetOrganizationUUID() {
require.True(s.T(), ok, "type assertion failed") require.True(s.T(), ok, "type assertion failed")
s.client = client s.client = client
s.client.baseURL = "http://in-process" s.client.baseURL = "http://in-process"
s.client.clientFactory = func(string) *req.Client { return newTestReqClient(rt) } s.client.clientFactory = func(string) (*req.Client, error) { return newTestReqClient(rt), nil }
got, err := s.client.GetOrganizationUUID(context.Background(), "sess", "") got, err := s.client.GetOrganizationUUID(context.Background(), "sess", "")
...@@ -169,7 +169,7 @@ func (s *ClaudeOAuthServiceSuite) TestGetAuthorizationCode() { ...@@ -169,7 +169,7 @@ func (s *ClaudeOAuthServiceSuite) TestGetAuthorizationCode() {
require.True(s.T(), ok, "type assertion failed") require.True(s.T(), ok, "type assertion failed")
s.client = client s.client = client
s.client.baseURL = "http://in-process" s.client.baseURL = "http://in-process"
s.client.clientFactory = func(string) *req.Client { return newTestReqClient(rt) } s.client.clientFactory = func(string) (*req.Client, error) { return newTestReqClient(rt), nil }
code, err := s.client.GetAuthorizationCode(context.Background(), "sess", "org-1", oauth.ScopeInference, "cc", "st", "") code, err := s.client.GetAuthorizationCode(context.Background(), "sess", "org-1", oauth.ScopeInference, "cc", "st", "")
...@@ -276,7 +276,7 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() { ...@@ -276,7 +276,7 @@ func (s *ClaudeOAuthServiceSuite) TestExchangeCodeForToken() {
require.True(s.T(), ok, "type assertion failed") require.True(s.T(), ok, "type assertion failed")
s.client = client s.client = client
s.client.tokenURL = "http://in-process/token" s.client.tokenURL = "http://in-process/token"
s.client.clientFactory = func(string) *req.Client { return newTestReqClient(rt) } s.client.clientFactory = func(string) (*req.Client, error) { return newTestReqClient(rt), nil }
resp, err := s.client.ExchangeCodeForToken(context.Background(), tt.code, "ver", "", "", tt.isSetupToken) resp, err := s.client.ExchangeCodeForToken(context.Background(), tt.code, "ver", "", "", tt.isSetupToken)
...@@ -372,7 +372,7 @@ func (s *ClaudeOAuthServiceSuite) TestRefreshToken() { ...@@ -372,7 +372,7 @@ func (s *ClaudeOAuthServiceSuite) TestRefreshToken() {
require.True(s.T(), ok, "type assertion failed") require.True(s.T(), ok, "type assertion failed")
s.client = client s.client = client
s.client.tokenURL = "http://in-process/token" s.client.tokenURL = "http://in-process/token"
s.client.clientFactory = func(string) *req.Client { return newTestReqClient(rt) } s.client.clientFactory = func(string) (*req.Client, error) { return newTestReqClient(rt), nil }
resp, err := s.client.RefreshToken(context.Background(), "rt", "") resp, err := s.client.RefreshToken(context.Background(), "rt", "")
......
backend/internal/repository/claude_usage_service.go
View file @ 3d79773b
...@@ -83,7 +83,7 @@ func (s *claudeUsageService) FetchUsageWithOptions(ctx context.Context, opts *se ...@@ -83,7 +83,7 @@ func (s *claudeUsageService) FetchUsageWithOptions(ctx context.Context, opts *se
AllowPrivateHosts: s.allowPrivateHosts, AllowPrivateHosts: s.allowPrivateHosts,
}) })
if err != nil { if err != nil {
client = &http.Client{Timeout: 30 * time.Second} return nil, fmt.Errorf("create http client failed: %w", err)
} }
resp, err = client.Do(req) resp, err = client.Do(req)
......
backend/internal/repository/claude_usage_service_test.go
View file @ 3d79773b
...@@ -50,7 +50,7 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_Success() { ...@@ -50,7 +50,7 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_Success() {
allowPrivateHosts: true, allowPrivateHosts: true,
} }
resp, err := s.fetcher.FetchUsage(context.Background(), "at", "://bad-proxy-url") resp, err := s.fetcher.FetchUsage(context.Background(), "at", "")
require.NoError(s.T(), err, "FetchUsage") require.NoError(s.T(), err, "FetchUsage")
require.Equal(s.T(), 12.5, resp.FiveHour.Utilization, "FiveHour utilization mismatch") require.Equal(s.T(), 12.5, resp.FiveHour.Utilization, "FiveHour utilization mismatch")
require.Equal(s.T(), 34.0, resp.SevenDay.Utilization, "SevenDay utilization mismatch") require.Equal(s.T(), 34.0, resp.SevenDay.Utilization, "SevenDay utilization mismatch")
...@@ -112,6 +112,17 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_ContextCancel() { ...@@ -112,6 +112,17 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_ContextCancel() {
require.Error(s.T(), err, "expected error for cancelled context") require.Error(s.T(), err, "expected error for cancelled context")
} }
func (s *ClaudeUsageServiceSuite) TestFetchUsage_InvalidProxyReturnsError() {
s.fetcher = &claudeUsageService{
usageURL: "http://example.com",
allowPrivateHosts: true,
}
_, err := s.fetcher.FetchUsage(context.Background(), "at", "://bad-proxy-url")
require.Error(s.T(), err)
require.ErrorContains(s.T(), err, "create http client failed")
}
func TestClaudeUsageServiceSuite(t *testing.T) { func TestClaudeUsageServiceSuite(t *testing.T) {
suite.Run(t, new(ClaudeUsageServiceSuite)) suite.Run(t, new(ClaudeUsageServiceSuite))
} }
backend/internal/repository/concurrency_cache.go
View file @ 3d79773b
...@@ -147,100 +147,6 @@ var ( ...@@ -147,100 +147,6 @@ var (
return 1 return 1
`) `)
// getAccountsLoadBatchScript - batch load query with expired slot cleanup
// ARGV[1] = slot TTL (seconds)
// ARGV[2..n] = accountID1, maxConcurrency1, accountID2, maxConcurrency2, ...
getAccountsLoadBatchScript = redis.NewScript(`
local result = {}
local slotTTL = tonumber(ARGV[1])
-- Get current server time
local timeResult = redis.call('TIME')
local nowSeconds = tonumber(timeResult[1])
local cutoffTime = nowSeconds - slotTTL
local i = 2
while i <= #ARGV do
local accountID = ARGV[i]
local maxConcurrency = tonumber(ARGV[i + 1])
local slotKey = 'concurrency:account:' .. accountID
-- Clean up expired slots before counting
redis.call('ZREMRANGEBYSCORE', slotKey, '-inf', cutoffTime)
local currentConcurrency = redis.call('ZCARD', slotKey)
local waitKey = 'wait:account:' .. accountID
local waitingCount = redis.call('GET', waitKey)
if waitingCount == false then
waitingCount = 0
else
waitingCount = tonumber(waitingCount)
end
local loadRate = 0
if maxConcurrency > 0 then
loadRate = math.floor((currentConcurrency + waitingCount) * 100 / maxConcurrency)
end
table.insert(result, accountID)
table.insert(result, currentConcurrency)
table.insert(result, waitingCount)
table.insert(result, loadRate)
i = i + 2
end
return result
`)
// getUsersLoadBatchScript - batch load query for users with expired slot cleanup
// ARGV[1] = slot TTL (seconds)
// ARGV[2..n] = userID1, maxConcurrency1, userID2, maxConcurrency2, ...
getUsersLoadBatchScript = redis.NewScript(`
local result = {}
local slotTTL = tonumber(ARGV[1])
-- Get current server time
local timeResult = redis.call('TIME')
local nowSeconds = tonumber(timeResult[1])
local cutoffTime = nowSeconds - slotTTL
local i = 2
while i <= #ARGV do
local userID = ARGV[i]
local maxConcurrency = tonumber(ARGV[i + 1])
local slotKey = 'concurrency:user:' .. userID
-- Clean up expired slots before counting
redis.call('ZREMRANGEBYSCORE', slotKey, '-inf', cutoffTime)
local currentConcurrency = redis.call('ZCARD', slotKey)
local waitKey = 'concurrency:wait:' .. userID
local waitingCount = redis.call('GET', waitKey)
if waitingCount == false then
waitingCount = 0
else
waitingCount = tonumber(waitingCount)
end
local loadRate = 0
if maxConcurrency > 0 then
loadRate = math.floor((currentConcurrency + waitingCount) * 100 / maxConcurrency)
end
table.insert(result, userID)
table.insert(result, currentConcurrency)
table.insert(result, waitingCount)
table.insert(result, loadRate)
i = i + 2
end
return result
`)
// cleanupExpiredSlotsScript - remove expired slots // cleanupExpiredSlotsScript - remove expired slots
// KEYS[1] = concurrency:account:{accountID} // KEYS[1] = concurrency:account:{accountID}
// ARGV[1] = TTL (seconds) // ARGV[1] = TTL (seconds)
...@@ -321,6 +227,43 @@ func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID ...@@ -321,6 +227,43 @@ func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID
return result, nil return result, nil
} }
func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error) {
if len(accountIDs) == 0 {
return map[int64]int{}, nil
}
now, err := c.rdb.Time(ctx).Result()
if err != nil {
return nil, fmt.Errorf("redis TIME: %w", err)
}
cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
pipe := c.rdb.Pipeline()
type accountCmd struct {
accountID int64
zcardCmd *redis.IntCmd
}
cmds := make([]accountCmd, 0, len(accountIDs))
for _, accountID := range accountIDs {
slotKey := accountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
cmds = append(cmds, accountCmd{
accountID: accountID,
zcardCmd: pipe.ZCard(ctx, slotKey),
})
}
if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
return nil, fmt.Errorf("pipeline exec: %w", err)
}
result := make(map[int64]int, len(accountIDs))
for _, cmd := range cmds {
result[cmd.accountID] = int(cmd.zcardCmd.Val())
}
return result, nil
}
// User slot operations // User slot operations
func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) { func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
...@@ -399,29 +342,53 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts [] ...@@ -399,29 +342,53 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []
return map[int64]*service.AccountLoadInfo{}, nil return map[int64]*service.AccountLoadInfo{}, nil
} }
args := []any{c.slotTTLSeconds} // 使用 Pipeline 替代 Lua 脚本，兼容 Redis Cluster（Lua 内动态拼 key 会 CROSSSLOT）。
for _, acc := range accounts { // 每个账号执行 3 个命令：ZREMRANGEBYSCORE（清理过期）、ZCARD（并发数）、GET（等待数）。
args = append(args, acc.ID, acc.MaxConcurrency) now, err := c.rdb.Time(ctx).Result()
}
result, err := getAccountsLoadBatchScript.Run(ctx, c.rdb, []string{}, args...).Slice()
if err != nil { if err != nil {
return nil, err return nil, fmt.Errorf("redis TIME: %w", err)
} }
cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
pipe := c.rdb.Pipeline()
loadMap := make(map[int64]*service.AccountLoadInfo) type accountCmds struct {
for i := 0; i < len(result); i += 4 { id int64
if i+3 >= len(result) { maxConcurrency int
break zcardCmd *redis.IntCmd
getCmd *redis.StringCmd
}
cmds := make([]accountCmds, 0, len(accounts))
for _, acc := range accounts {
slotKey := accountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
waitKey := accountWaitKeyPrefix + strconv.FormatInt(acc.ID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
ac := accountCmds{
id: acc.ID,
maxConcurrency: acc.MaxConcurrency,
zcardCmd: pipe.ZCard(ctx, slotKey),
getCmd: pipe.Get(ctx, waitKey),
} }
cmds = append(cmds, ac)
}
accountID, _ := strconv.ParseInt(fmt.Sprintf("%v", result[i]), 10, 64) if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
currentConcurrency, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+1])) return nil, fmt.Errorf("pipeline exec: %w", err)
waitingCount, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+2])) }
loadRate, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+3]))
loadMap[accountID] = &service.AccountLoadInfo{ loadMap := make(map[int64]*service.AccountLoadInfo, len(accounts))
AccountID: accountID, for _, ac := range cmds {
currentConcurrency := int(ac.zcardCmd.Val())
waitingCount := 0
if v, err := ac.getCmd.Int(); err == nil {
waitingCount = v
}
loadRate := 0
if ac.maxConcurrency > 0 {
loadRate = (currentConcurrency + waitingCount) * 100 / ac.maxConcurrency
}
loadMap[ac.id] = &service.AccountLoadInfo{
AccountID: ac.id,
CurrentConcurrency: currentConcurrency, CurrentConcurrency: currentConcurrency,
WaitingCount: waitingCount, WaitingCount: waitingCount,
LoadRate: loadRate, LoadRate: loadRate,
...@@ -436,29 +403,52 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic ...@@ -436,29 +403,52 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
return map[int64]*service.UserLoadInfo{}, nil return map[int64]*service.UserLoadInfo{}, nil
} }
args := []any{c.slotTTLSeconds} // 使用 Pipeline 替代 Lua 脚本，兼容 Redis Cluster。
for _, u := range users { now, err := c.rdb.Time(ctx).Result()
args = append(args, u.ID, u.MaxConcurrency)
}
result, err := getUsersLoadBatchScript.Run(ctx, c.rdb, []string{}, args...).Slice()
if err != nil { if err != nil {
return nil, err return nil, fmt.Errorf("redis TIME: %w", err)
} }
cutoffTime := now.Unix() - int64(c.slotTTLSeconds)
pipe := c.rdb.Pipeline()
loadMap := make(map[int64]*service.UserLoadInfo) type userCmds struct {
for i := 0; i < len(result); i += 4 { id int64
if i+3 >= len(result) { maxConcurrency int
break zcardCmd *redis.IntCmd
getCmd *redis.StringCmd
}
cmds := make([]userCmds, 0, len(users))
for _, u := range users {
slotKey := userSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
waitKey := waitQueueKeyPrefix + strconv.FormatInt(u.ID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
uc := userCmds{
id: u.ID,
maxConcurrency: u.MaxConcurrency,
zcardCmd: pipe.ZCard(ctx, slotKey),
getCmd: pipe.Get(ctx, waitKey),
} }
cmds = append(cmds, uc)
}
userID, _ := strconv.ParseInt(fmt.Sprintf("%v", result[i]), 10, 64) if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
currentConcurrency, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+1])) return nil, fmt.Errorf("pipeline exec: %w", err)
waitingCount, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+2])) }
loadRate, _ := strconv.Atoi(fmt.Sprintf("%v", result[i+3]))
loadMap[userID] = &service.UserLoadInfo{ loadMap := make(map[int64]*service.UserLoadInfo, len(users))
UserID: userID, for _, uc := range cmds {
currentConcurrency := int(uc.zcardCmd.Val())
waitingCount := 0
if v, err := uc.getCmd.Int(); err == nil {
waitingCount = v
}
loadRate := 0
if uc.maxConcurrency > 0 {
loadRate = (currentConcurrency + waitingCount) * 100 / uc.maxConcurrency
}
loadMap[uc.id] = &service.UserLoadInfo{
UserID: uc.id,
CurrentConcurrency: currentConcurrency, CurrentConcurrency: currentConcurrency,
WaitingCount: waitingCount, WaitingCount: waitingCount,
LoadRate: loadRate, LoadRate: loadRate,
......
backend/internal/repository/ent.go
View file @ 3d79773b
...@@ -5,6 +5,7 @@ package repository ...@@ -5,6 +5,7 @@ package repository
import ( import (
"context" "context"
"database/sql" "database/sql"
"fmt"
"time" "time"
"github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/ent"
...@@ -66,6 +67,18 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) { ...@@ -66,6 +67,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))
// 启动阶段：从配置或数据库中确保系统密钥可用。
if err := ensureBootstrapSecrets(migrationCtx, client, cfg); err != nil {
_ = client.Close()
return nil, nil, err
}
// 在密钥补齐后执行完整配置校验，避免空 jwt.secret 导致服务运行时失败。
if err := cfg.Validate(); err != nil {
_ = client.Close()
return nil, nil, fmt.Errorf("validate config after secret bootstrap: %w", err)
}
// SIMPLE 模式：启动时补齐各平台默认分组。 // SIMPLE 模式：启动时补齐各平台默认分组。
// - anthropic/openai/gemini: 确保存在 <platform>-default // - anthropic/openai/gemini: 确保存在 <platform>-default
// - antigravity: 仅要求存在 >=2 个未软删除分组（用于 claude/gemini 混合调度场景） // - antigravity: 仅要求存在 >=2 个未软删除分组（用于 claude/gemini 混合调度场景）
......
backend/internal/repository/gateway_cache_integration_test.go
View file @ 3d79773b
...@@ -104,7 +104,6 @@ func (s *GatewayCacheSuite) TestGetSessionAccountID_CorruptedValue() { ...@@ -104,7 +104,6 @@ func (s *GatewayCacheSuite) TestGetSessionAccountID_CorruptedValue() {
require.False(s.T(), errors.Is(err, redis.Nil), "expected parsing error, not redis.Nil") require.False(s.T(), errors.Is(err, redis.Nil), "expected parsing error, not redis.Nil")
} }
func TestGatewayCacheSuite(t *testing.T) { func TestGatewayCacheSuite(t *testing.T) {
suite.Run(t, new(GatewayCacheSuite)) suite.Run(t, new(GatewayCacheSuite))
} }
backend/internal/repository/gemini_drive_client.go 0 → 100644
View file @ 3d79773b
package repository
import "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
// NewGeminiDriveClient creates a concrete DriveClient for Google Drive API operations.
// Returned as geminicli.DriveClient interface for DI (Strategy A).
func NewGeminiDriveClient() geminicli.DriveClient {
return geminicli.NewDriveClient()
}
backend/internal/repository/gemini_oauth_client.go
View file @ 3d79773b
...@@ -26,7 +26,10 @@ func NewGeminiOAuthClient(cfg *config.Config) service.GeminiOAuthClient { ...@@ -26,7 +26,10 @@ func NewGeminiOAuthClient(cfg *config.Config) service.GeminiOAuthClient {
} }
func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) { func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) {
client := createGeminiReqClient(proxyURL) client, err := createGeminiReqClient(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
// Use different OAuth clients based on oauthType: // Use different OAuth clients based on oauthType:
// - code_assist: always use built-in Gemini CLI OAuth client (public) // - code_assist: always use built-in Gemini CLI OAuth client (public)
...@@ -72,7 +75,10 @@ func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, c ...@@ -72,7 +75,10 @@ func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, c
} }
func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) { func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
client := createGeminiReqClient(proxyURL) client, err := createGeminiReqClient(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
oauthCfgInput := geminicli.OAuthConfig{ oauthCfgInput := geminicli.OAuthConfig{
ClientID: c.cfg.Gemini.OAuth.ClientID, ClientID: c.cfg.Gemini.OAuth.ClientID,
...@@ -111,7 +117,7 @@ func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refresh ...@@ -111,7 +117,7 @@ func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refresh
return &tokenResp, nil return &tokenResp, nil
} }
func createGeminiReqClient(proxyURL string) *req.Client { func createGeminiReqClient(proxyURL string) (*req.Client, error) {
return getSharedReqClient(reqClientOptions{ return getSharedReqClient(reqClientOptions{
ProxyURL: proxyURL, ProxyURL: proxyURL,
Timeout: 60 * time.Second, Timeout: 60 * time.Second,
......
backend/internal/repository/geminicli_codeassist_client.go
View file @ 3d79773b
...@@ -26,7 +26,11 @@ func (c *geminiCliCodeAssistClient) LoadCodeAssist(ctx context.Context, accessTo ...@@ -26,7 +26,11 @@ func (c *geminiCliCodeAssistClient) LoadCodeAssist(ctx context.Context, accessTo
} }
var out geminicli.LoadCodeAssistResponse var out geminicli.LoadCodeAssistResponse
resp, err := createGeminiCliReqClient(proxyURL).R(). client, err := createGeminiCliReqClient(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
resp, err := client.R().
SetContext(ctx). SetContext(ctx).
SetHeader("Authorization", "Bearer "+accessToken). SetHeader("Authorization", "Bearer "+accessToken).
SetHeader("Content-Type", "application/json"). SetHeader("Content-Type", "application/json").
...@@ -66,7 +70,11 @@ func (c *geminiCliCodeAssistClient) OnboardUser(ctx context.Context, accessToken ...@@ -66,7 +70,11 @@ func (c *geminiCliCodeAssistClient) OnboardUser(ctx context.Context, accessToken
fmt.Printf("[CodeAssist] OnboardUser request body: %+v\n", reqBody) fmt.Printf("[CodeAssist] OnboardUser request body: %+v\n", reqBody)
var out geminicli.OnboardUserResponse var out geminicli.OnboardUserResponse
resp, err := createGeminiCliReqClient(proxyURL).R(). client, err := createGeminiCliReqClient(proxyURL)
if err != nil {
return nil, fmt.Errorf("create HTTP client: %w", err)
}
resp, err := client.R().
SetContext(ctx). SetContext(ctx).
SetHeader("Authorization", "Bearer "+accessToken). SetHeader("Authorization", "Bearer "+accessToken).
SetHeader("Content-Type", "application/json"). SetHeader("Content-Type", "application/json").
...@@ -98,7 +106,7 @@ func (c *geminiCliCodeAssistClient) OnboardUser(ctx context.Context, accessToken ...@@ -98,7 +106,7 @@ func (c *geminiCliCodeAssistClient) OnboardUser(ctx context.Context, accessToken
return &out, nil return &out, nil
} }
func createGeminiCliReqClient(proxyURL string) *req.Client { func createGeminiCliReqClient(proxyURL string) (*req.Client, error) {
return getSharedReqClient(reqClientOptions{ return getSharedReqClient(reqClientOptions{
ProxyURL: proxyURL, ProxyURL: proxyURL,
Timeout: 30 * time.Second, Timeout: 30 * time.Second,
......
backend/internal/repository/github_release_service.go
View file @ 3d79773b
...@@ -5,8 +5,10 @@ import ( ...@@ -5,8 +5,10 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"log/slog"
"net/http" "net/http"
"os" "os"
"strings"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient" "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
...@@ -18,14 +20,27 @@ type githubReleaseClient struct { ...@@ -18,14 +20,27 @@ type githubReleaseClient struct {
downloadHTTPClient *http.Client downloadHTTPClient *http.Client
} }
type githubReleaseClientError struct {
err error
}
// NewGitHubReleaseClient 创建 GitHub Release 客户端 // NewGitHubReleaseClient 创建 GitHub Release 客户端
// proxyURL 为空时直连 GitHub，支持 http/https/socks5/socks5h 协议 // proxyURL 为空时直连 GitHub，支持 http/https/socks5/socks5h 协议
func NewGitHubReleaseClient(proxyURL string) service.GitHubReleaseClient { // 代理配置失败时行为由 allowDirectOnProxyError 控制：
// - false（默认）：返回错误占位客户端，禁止回退到直连
// - true：回退到直连（仅限管理员显式开启）
func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) service.GitHubReleaseClient {
// 安全说明：httpclient.GetClient 的错误链（url.Parse / proxyutil）不含明文代理凭据，
// 但仍通过 slog 仅在服务端日志记录，不会暴露给 HTTP 响应。
sharedClient, err := httpclient.GetClient(httpclient.Options{ sharedClient, err := httpclient.GetClient(httpclient.Options{
Timeout: 30 * time.Second, Timeout: 30 * time.Second,
ProxyURL: proxyURL, ProxyURL: proxyURL,
}) })
if err != nil { if err != nil {
if strings.TrimSpace(proxyURL) != "" && !allowDirectOnProxyError {
slog.Warn("proxy client init failed, all requests will fail", "service", "github_release", "error", err)
return &githubReleaseClientError{err: fmt.Errorf("proxy client init failed and direct fallback is disabled; set security.proxy_fallback.allow_direct_on_error=true to allow fallback: %w", err)}
}
sharedClient = &http.Client{Timeout: 30 * time.Second} sharedClient = &http.Client{Timeout: 30 * time.Second}
} }
...@@ -35,6 +50,10 @@ func NewGitHubReleaseClient(proxyURL string) service.GitHubReleaseClient { ...@@ -35,6 +50,10 @@ func NewGitHubReleaseClient(proxyURL string) service.GitHubReleaseClient {
ProxyURL: proxyURL, ProxyURL: proxyURL,
}) })
if err != nil { if err != nil {
if strings.TrimSpace(proxyURL) != "" && !allowDirectOnProxyError {
slog.Warn("proxy download client init failed, all requests will fail", "service", "github_release", "error", err)
return &githubReleaseClientError{err: fmt.Errorf("proxy client init failed and direct fallback is disabled; set security.proxy_fallback.allow_direct_on_error=true to allow fallback: %w", err)}
}
downloadClient = &http.Client{Timeout: 10 * time.Minute} downloadClient = &http.Client{Timeout: 10 * time.Minute}
} }
...@@ -44,6 +63,18 @@ func NewGitHubReleaseClient(proxyURL string) service.GitHubReleaseClient { ...@@ -44,6 +63,18 @@ func NewGitHubReleaseClient(proxyURL string) service.GitHubReleaseClient {
} }
} }
func (c *githubReleaseClientError) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
return nil, c.err
}
func (c *githubReleaseClientError) DownloadFile(ctx context.Context, url, dest string, maxSize int64) error {
return c.err
}
func (c *githubReleaseClientError) FetchChecksumFile(ctx context.Context, url string) ([]byte, error) {
return nil, c.err
}
func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) { func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo) url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo)
......
backend/internal/repository/group_repo.go
View file @ 3d79773b
...@@ -4,11 +4,13 @@ import ( ...@@ -4,11 +4,13 @@ import (
"context" "context"
"database/sql" "database/sql"
"errors" "errors"
"log" "fmt"
"strings"
dbent "github.com/Wei-Shaw/sub2api/ent" dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/apikey" "github.com/Wei-Shaw/sub2api/ent/apikey"
"github.com/Wei-Shaw/sub2api/ent/group" "github.com/Wei-Shaw/sub2api/ent/group"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination" "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service" "github.com/Wei-Shaw/sub2api/internal/service"
"github.com/lib/pq" "github.com/lib/pq"
...@@ -47,12 +49,17 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er ...@@ -47,12 +49,17 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er
SetNillableImagePrice1k(groupIn.ImagePrice1K). SetNillableImagePrice1k(groupIn.ImagePrice1K).
SetNillableImagePrice2k(groupIn.ImagePrice2K). SetNillableImagePrice2k(groupIn.ImagePrice2K).
SetNillableImagePrice4k(groupIn.ImagePrice4K). SetNillableImagePrice4k(groupIn.ImagePrice4K).
SetNillableSoraImagePrice360(groupIn.SoraImagePrice360).
SetNillableSoraImagePrice540(groupIn.SoraImagePrice540).
SetNillableSoraVideoPricePerRequest(groupIn.SoraVideoPricePerRequest).
SetNillableSoraVideoPricePerRequestHd(groupIn.SoraVideoPricePerRequestHD).
SetDefaultValidityDays(groupIn.DefaultValidityDays). SetDefaultValidityDays(groupIn.DefaultValidityDays).
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
SetNillableFallbackGroupID(groupIn.FallbackGroupID). SetNillableFallbackGroupID(groupIn.FallbackGroupID).
SetNillableFallbackGroupIDOnInvalidRequest(groupIn.FallbackGroupIDOnInvalidRequest). SetNillableFallbackGroupIDOnInvalidRequest(groupIn.FallbackGroupIDOnInvalidRequest).
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled). SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
SetMcpXMLInject(groupIn.MCPXMLInject) SetMcpXMLInject(groupIn.MCPXMLInject).
SetSoraStorageQuotaBytes(groupIn.SoraStorageQuotaBytes)
// 设置模型路由配置 // 设置模型路由配置
if groupIn.ModelRouting != nil { if groupIn.ModelRouting != nil {
...@@ -68,7 +75,7 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er ...@@ -68,7 +75,7 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er
groupIn.CreatedAt = created.CreatedAt groupIn.CreatedAt = created.CreatedAt
groupIn.UpdatedAt = created.UpdatedAt groupIn.UpdatedAt = created.UpdatedAt
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue group create failed: group=%d err=%v", groupIn.ID, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group create failed: group=%d err=%v", groupIn.ID, err)
} }
} }
return translatePersistenceError(err, nil, service.ErrGroupExists) return translatePersistenceError(err, nil, service.ErrGroupExists)
...@@ -110,10 +117,47 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er ...@@ -110,10 +117,47 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
SetNillableImagePrice1k(groupIn.ImagePrice1K). SetNillableImagePrice1k(groupIn.ImagePrice1K).
SetNillableImagePrice2k(groupIn.ImagePrice2K). SetNillableImagePrice2k(groupIn.ImagePrice2K).
SetNillableImagePrice4k(groupIn.ImagePrice4K). SetNillableImagePrice4k(groupIn.ImagePrice4K).
SetNillableSoraImagePrice360(groupIn.SoraImagePrice360).
SetNillableSoraImagePrice540(groupIn.SoraImagePrice540).
SetNillableSoraVideoPricePerRequest(groupIn.SoraVideoPricePerRequest).
SetNillableSoraVideoPricePerRequestHd(groupIn.SoraVideoPricePerRequestHD).
SetDefaultValidityDays(groupIn.DefaultValidityDays). SetDefaultValidityDays(groupIn.DefaultValidityDays).
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled). SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
SetMcpXMLInject(groupIn.MCPXMLInject) SetMcpXMLInject(groupIn.MCPXMLInject).
SetSoraStorageQuotaBytes(groupIn.SoraStorageQuotaBytes)
// 显式处理可空字段：nil 需要 clear，非 nil 需要 set。
if groupIn.DailyLimitUSD != nil {
builder = builder.SetDailyLimitUsd(*groupIn.DailyLimitUSD)
} else {
builder = builder.ClearDailyLimitUsd()
}
if groupIn.WeeklyLimitUSD != nil {
builder = builder.SetWeeklyLimitUsd(*groupIn.WeeklyLimitUSD)
} else {
builder = builder.ClearWeeklyLimitUsd()
}
if groupIn.MonthlyLimitUSD != nil {
builder = builder.SetMonthlyLimitUsd(*groupIn.MonthlyLimitUSD)
} else {
builder = builder.ClearMonthlyLimitUsd()
}
if groupIn.ImagePrice1K != nil {
builder = builder.SetImagePrice1k(*groupIn.ImagePrice1K)
} else {
builder = builder.ClearImagePrice1k()
}
if groupIn.ImagePrice2K != nil {
builder = builder.SetImagePrice2k(*groupIn.ImagePrice2K)
} else {
builder = builder.ClearImagePrice2k()
}
if groupIn.ImagePrice4K != nil {
builder = builder.SetImagePrice4k(*groupIn.ImagePrice4K)
} else {
builder = builder.ClearImagePrice4k()
}
// 处理 FallbackGroupID：nil 时清除，否则设置 // 处理 FallbackGroupID：nil 时清除，否则设置
if groupIn.FallbackGroupID != nil { if groupIn.FallbackGroupID != nil {
...@@ -144,7 +188,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er ...@@ -144,7 +188,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
} }
groupIn.UpdatedAt = updated.UpdatedAt groupIn.UpdatedAt = updated.UpdatedAt
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue group update failed: group=%d err=%v", groupIn.ID, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group update failed: group=%d err=%v", groupIn.ID, err)
} }
return nil return nil
} }
...@@ -155,7 +199,7 @@ func (r *groupRepository) Delete(ctx context.Context, id int64) error { ...@@ -155,7 +199,7 @@ func (r *groupRepository) Delete(ctx context.Context, id int64) error {
return translatePersistenceError(err, service.ErrGroupNotFound, nil) return translatePersistenceError(err, service.ErrGroupNotFound, nil)
} }
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue group delete failed: group=%d err=%v", id, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group delete failed: group=%d err=%v", id, err)
} }
return nil return nil
} }
...@@ -183,7 +227,7 @@ func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination ...@@ -183,7 +227,7 @@ func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination
q = q.Where(group.IsExclusiveEQ(*isExclusive)) q = q.Where(group.IsExclusiveEQ(*isExclusive))
} }
total, err := q.Count(ctx) total, err := q.Clone().Count(ctx)
if err != nil { if err != nil {
return nil, nil, err return nil, nil, err
} }
...@@ -273,6 +317,54 @@ func (r *groupRepository) ExistsByName(ctx context.Context, name string) (bool, ...@@ -273,6 +317,54 @@ func (r *groupRepository) ExistsByName(ctx context.Context, name string) (bool,
return r.client.Group.Query().Where(group.NameEQ(name)).Exist(ctx) return r.client.Group.Query().Where(group.NameEQ(name)).Exist(ctx)
} }
// ExistsByIDs 批量检查分组是否存在（仅检查未软删除记录）。
// 返回结构：map[groupID]exists。
func (r *groupRepository) ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error) {
result := make(map[int64]bool, len(ids))
if len(ids) == 0 {
return result, nil
}
uniqueIDs := 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{}{}
uniqueIDs = append(uniqueIDs, id)
result[id] = false
}
if len(uniqueIDs) == 0 {
return result, nil
}
rows, err := r.sql.QueryContext(ctx, `
SELECT id
FROM groups
WHERE id = ANY($1) AND deleted_at IS NULL
`, pq.Array(uniqueIDs))
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
result[id] = true
}
if err := rows.Err(); err != nil {
return nil, err
}
return result, nil
}
func (r *groupRepository) GetAccountCount(ctx context.Context, groupID int64) (int64, error) { func (r *groupRepository) GetAccountCount(ctx context.Context, groupID int64) (int64, error) {
var count int64 var count int64
if err := scanSingleRow(ctx, r.sql, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{groupID}, &count); err != nil { if err := scanSingleRow(ctx, r.sql, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{groupID}, &count); err != nil {
...@@ -288,7 +380,7 @@ func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, grou ...@@ -288,7 +380,7 @@ func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, grou
} }
affected, _ := res.RowsAffected() affected, _ := res.RowsAffected()
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue group account clear failed: group=%d err=%v", groupID, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group account clear failed: group=%d err=%v", groupID, err)
} }
return affected, nil return affected, nil
} }
...@@ -398,7 +490,7 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64, ...@@ -398,7 +490,7 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
} }
} }
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue group cascade delete failed: group=%d err=%v", id, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group cascade delete failed: group=%d err=%v", id, err)
} }
return affectedUserIDs, nil return affectedUserIDs, nil
...@@ -492,7 +584,7 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64 ...@@ -492,7 +584,7 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
// 发送调度器事件 // 发送调度器事件
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil { if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
log.Printf("[SchedulerOutbox] enqueue bind accounts to group failed: group=%d err=%v", groupID, err) logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue bind accounts to group failed: group=%d err=%v", groupID, err)
} }
return nil return nil
...@@ -504,22 +596,72 @@ func (r *groupRepository) UpdateSortOrders(ctx context.Context, updates []servic ...@@ -504,22 +596,72 @@ func (r *groupRepository) UpdateSortOrders(ctx context.Context, updates []servic
return nil return nil
} }
// 使用事务批量更新 // 去重后保留最后一次排序值，避免重复 ID 造成 CASE 分支冲突。
tx, err := r.client.Tx(ctx) sortOrderByID := make(map[int64]int, len(updates))
if err != nil { groupIDs := make([]int64, 0, len(updates))
for _, u := range updates {
if u.ID <= 0 {
continue
}
if _, exists := sortOrderByID[u.ID]; !exists {
groupIDs = append(groupIDs, u.ID)
}
sortOrderByID[u.ID] = u.SortOrder
}
if len(groupIDs) == 0 {
return nil
}
// 与旧实现保持一致：任何不存在/已删除的分组都返回 not found，且不执行更新。
var existingCount int
if err := scanSingleRow(
ctx,
r.sql,
`SELECT COUNT(*) FROM groups WHERE deleted_at IS NULL AND id = ANY($1)`,
[]any{pq.Array(groupIDs)},
&existingCount,
); err != nil {
return err return err
} }
defer func() { _ = tx.Rollback() }() if existingCount != len(groupIDs) {
return service.ErrGroupNotFound
}
for _, u := range updates { args := make([]any, 0, len(groupIDs)*2+1)
if _, err := tx.Group.UpdateOneID(u.ID).SetSortOrder(u.SortOrder).Save(ctx); err != nil { caseClauses := make([]string, 0, len(groupIDs))
return translatePersistenceError(err, service.ErrGroupNotFound, nil) placeholder := 1
} for _, id := range groupIDs {
caseClauses = append(caseClauses, fmt.Sprintf("WHEN $%d THEN $%d", placeholder, placeholder+1))
args = append(args, id, sortOrderByID[id])
placeholder += 2
} }
args = append(args, pq.Array(groupIDs))
query := fmt.Sprintf(`
UPDATE groups
SET sort_order = CASE id
%s
ELSE sort_order
END
WHERE deleted_at IS NULL AND id = ANY($%d)
`, strings.Join(caseClauses, "\n\t\t\t"), placeholder)
if err := tx.Commit(); err != nil { result, err := r.sql.ExecContext(ctx, query, args...)
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return err return err
} }
if affected != int64(len(groupIDs)) {
return service.ErrGroupNotFound
}
for _, id := range groupIDs {
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group sort update failed: group=%d err=%v", id, err)
}
}
return nil return nil
} }
backend/internal/repository/group_repo_integration_test.go
View file @ 3d79773b
...@@ -352,6 +352,81 @@ func (s *GroupRepoSuite) TestListWithFilters_Search() { ...@@ -352,6 +352,81 @@ func (s *GroupRepoSuite) TestListWithFilters_Search() {
}) })
} }
func (s *GroupRepoSuite) TestUpdateSortOrders_BatchCaseWhen() {
g1 := &service.Group{
Name: "sort-g1",
Platform: service.PlatformAnthropic,
RateMultiplier: 1.0,
IsExclusive: false,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
}
g2 := &service.Group{
Name: "sort-g2",
Platform: service.PlatformAnthropic,
RateMultiplier: 1.0,
IsExclusive: false,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
}
g3 := &service.Group{
Name: "sort-g3",
Platform: service.PlatformAnthropic,
RateMultiplier: 1.0,
IsExclusive: false,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
}
s.Require().NoError(s.repo.Create(s.ctx, g1))
s.Require().NoError(s.repo.Create(s.ctx, g2))
s.Require().NoError(s.repo.Create(s.ctx, g3))
err := s.repo.UpdateSortOrders(s.ctx, []service.GroupSortOrderUpdate{
{ID: g1.ID, SortOrder: 30},
{ID: g2.ID, SortOrder: 10},
{ID: g3.ID, SortOrder: 20},
{ID: g2.ID, SortOrder: 15}, // 重复 ID 应以最后一次为准
})
s.Require().NoError(err)
got1, err := s.repo.GetByID(s.ctx, g1.ID)
s.Require().NoError(err)
got2, err := s.repo.GetByID(s.ctx, g2.ID)
s.Require().NoError(err)
got3, err := s.repo.GetByID(s.ctx, g3.ID)
s.Require().NoError(err)
s.Require().Equal(30, got1.SortOrder)
s.Require().Equal(15, got2.SortOrder)
s.Require().Equal(20, got3.SortOrder)
}
func (s *GroupRepoSuite) TestUpdateSortOrders_MissingGroupNoPartialUpdate() {
g1 := &service.Group{
Name: "sort-no-partial",
Platform: service.PlatformAnthropic,
RateMultiplier: 1.0,
IsExclusive: false,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
}
s.Require().NoError(s.repo.Create(s.ctx, g1))
before, err := s.repo.GetByID(s.ctx, g1.ID)
s.Require().NoError(err)
beforeSort := before.SortOrder
err = s.repo.UpdateSortOrders(s.ctx, []service.GroupSortOrderUpdate{
{ID: g1.ID, SortOrder: 99},
{ID: 99999999, SortOrder: 1},
})
s.Require().Error(err)
s.Require().ErrorIs(err, service.ErrGroupNotFound)
after, err := s.repo.GetByID(s.ctx, g1.ID)
s.Require().NoError(err)
s.Require().Equal(beforeSort, after.SortOrder)
}
func (s *GroupRepoSuite) TestListWithFilters_AccountCount() { func (s *GroupRepoSuite) TestListWithFilters_AccountCount() {
g1 := &service.Group{ g1 := &service.Group{
Name: "g1", Name: "g1",
......
backend/internal/repository/http_upstream.go
View file @ 3d79773b
...@@ -14,6 +14,7 @@ import ( ...@@ -14,6 +14,7 @@ import (
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil" "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service" "github.com/Wei-Shaw/sub2api/internal/service"
...@@ -235,7 +236,10 @@ func (s *httpUpstreamService) acquireClientWithTLS(proxyURL string, accountID in ...@@ -235,7 +236,10 @@ func (s *httpUpstreamService) acquireClientWithTLS(proxyURL string, accountID in
// TLS 指纹客户端使用独立的缓存键，与普通客户端隔离 // TLS 指纹客户端使用独立的缓存键，与普通客户端隔离
func (s *httpUpstreamService) getClientEntryWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) { func (s *httpUpstreamService) getClientEntryWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, markInFlight bool, enforceLimit bool) (*upstreamClientEntry, error) {
isolation := s.getIsolationMode() isolation := s.getIsolationMode()
proxyKey, parsedProxy := normalizeProxyURL(proxyURL) proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
if err != nil {
return nil, err
}
// TLS 指纹客户端使用独立的缓存键，加 "tls:" 前缀 // TLS 指纹客户端使用独立的缓存键，加 "tls:" 前缀
cacheKey := "tls:" + buildCacheKey(isolation, proxyKey, accountID) cacheKey := "tls:" + buildCacheKey(isolation, proxyKey, accountID)
poolKey := s.buildPoolKey(isolation, accountConcurrency) + ":tls" poolKey := s.buildPoolKey(isolation, accountConcurrency) + ":tls"
...@@ -373,9 +377,8 @@ func (s *httpUpstreamService) acquireClient(proxyURL string, accountID int64, ac ...@@ -373,9 +377,8 @@ func (s *httpUpstreamService) acquireClient(proxyURL string, accountID int64, ac
// - proxy: 按代理地址隔离，同一代理共享客户端 // - proxy: 按代理地址隔离，同一代理共享客户端
// - account: 按账户隔离，同一账户共享客户端（代理变更时重建） // - account: 按账户隔离，同一账户共享客户端（代理变更时重建）
// - account_proxy: 按账户+代理组合隔离，最细粒度 // - account_proxy: 按账户+代理组合隔离，最细粒度
func (s *httpUpstreamService) getOrCreateClient(proxyURL string, accountID int64, accountConcurrency int) *upstreamClientEntry { func (s *httpUpstreamService) getOrCreateClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
entry, _ := s.getClientEntry(proxyURL, accountID, accountConcurrency, false, false) return s.getClientEntry(proxyURL, accountID, accountConcurrency, false, false)
return entry
} }
// getClientEntry 获取或创建客户端条目 // getClientEntry 获取或创建客户端条目
...@@ -385,7 +388,10 @@ func (s *httpUpstreamService) getClientEntry(proxyURL string, accountID int64, a ...@@ -385,7 +388,10 @@ func (s *httpUpstreamService) getClientEntry(proxyURL string, accountID int64, a
// 获取隔离模式 // 获取隔离模式
isolation := s.getIsolationMode() isolation := s.getIsolationMode()
// 标准化代理 URL 并解析 // 标准化代理 URL 并解析
proxyKey, parsedProxy := normalizeProxyURL(proxyURL) proxyKey, parsedProxy, err := normalizeProxyURL(proxyURL)
if err != nil {
return nil, err
}
// 构建缓存键（根据隔离策略不同） // 构建缓存键（根据隔离策略不同）
cacheKey := buildCacheKey(isolation, proxyKey, accountID) cacheKey := buildCacheKey(isolation, proxyKey, accountID)
// 构建连接池配置键（用于检测配置变更） // 构建连接池配置键（用于检测配置变更）
...@@ -680,17 +686,18 @@ func buildCacheKey(isolation, proxyKey string, accountID int64) string { ...@@ -680,17 +686,18 @@ func buildCacheKey(isolation, proxyKey string, accountID int64) string {
// - raw: 原始代理 URL 字符串 // - raw: 原始代理 URL 字符串
// //
// 返回: // 返回:
// - string: 标准化的代理键（空或解析失败返回 "direct"） // - string: 标准化的代理键（空返回 "direct"）
// - *url.URL: 解析后的 URL（空或解析失败返回 nil） // - *url.URL: 解析后的 URL（空返回 nil）
func normalizeProxyURL(raw string) (string, *url.URL) { // - error: 非空代理 URL 解析失败时返回错误（禁止回退到直连）
proxyURL := strings.TrimSpace(raw) func normalizeProxyURL(raw string) (string, *url.URL, error) {
if proxyURL == "" { _, parsed, err := proxyurl.Parse(raw)
return directProxyKey, nil
}
parsed, err := url.Parse(proxyURL)
if err != nil { if err != nil {
return directProxyKey, nil return "", nil, err
}
if parsed == nil {
return directProxyKey, nil, nil
} }
// 规范化：小写 scheme/host，去除路径和查询参数
parsed.Scheme = strings.ToLower(parsed.Scheme) parsed.Scheme = strings.ToLower(parsed.Scheme)
parsed.Host = strings.ToLower(parsed.Host) parsed.Host = strings.ToLower(parsed.Host)
parsed.Path = "" parsed.Path = ""
...@@ -710,7 +717,7 @@ func normalizeProxyURL(raw string) (string, *url.URL) { ...@@ -710,7 +717,7 @@ func normalizeProxyURL(raw string) (string, *url.URL) {
parsed.Host = hostname parsed.Host = hostname
} }
} }
return parsed.String(), parsed return parsed.String(), parsed, nil
} }
// defaultPoolSettings 获取默认连接池配置 // defaultPoolSettings 获取默认连接池配置
......
backend/internal/repository/http_upstream_benchmark_test.go
View file @ 3d79773b
...@@ -59,7 +59,10 @@ func BenchmarkHTTPUpstreamProxyClient(b *testing.B) { ...@@ -59,7 +59,10 @@ func BenchmarkHTTPUpstreamProxyClient(b *testing.B) {
// 模拟优化后的行为，从缓存获取客户端 // 模拟优化后的行为，从缓存获取客户端
b.Run("复用", func(b *testing.B) { b.Run("复用", func(b *testing.B) {
// 预热：确保客户端已缓存 // 预热：确保客户端已缓存
entry := svc.getOrCreateClient(proxyURL, 1, 1) entry, err := svc.getOrCreateClient(proxyURL, 1, 1)
if err != nil {
b.Fatalf("getOrCreateClient: %v", err)
}
client := entry.client client := entry.client
b.ResetTimer() // 重置计时器，排除预热时间 b.ResetTimer() // 重置计时器，排除预热时间
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
......
backend/internal/repository/http_upstream_test.go
View file @ 3d79773b
...@@ -44,7 +44,7 @@ func (s *HTTPUpstreamSuite) newService() *httpUpstreamService { ...@@ -44,7 +44,7 @@ func (s *HTTPUpstreamSuite) newService() *httpUpstreamService {
// 验证未配置时使用 300 秒默认值 // 验证未配置时使用 300 秒默认值
func (s *HTTPUpstreamSuite) TestDefaultResponseHeaderTimeout() { func (s *HTTPUpstreamSuite) TestDefaultResponseHeaderTimeout() {
svc := s.newService() svc := s.newService()
entry := svc.getOrCreateClient("", 0, 0) entry := mustGetOrCreateClient(s.T(), svc, "", 0, 0)
transport, ok := entry.client.Transport.(*http.Transport) transport, ok := entry.client.Transport.(*http.Transport)
require.True(s.T(), ok, "expected *http.Transport") require.True(s.T(), ok, "expected *http.Transport")
require.Equal(s.T(), 300*time.Second, transport.ResponseHeaderTimeout, "ResponseHeaderTimeout mismatch") require.Equal(s.T(), 300*time.Second, transport.ResponseHeaderTimeout, "ResponseHeaderTimeout mismatch")
...@@ -55,25 +55,27 @@ func (s *HTTPUpstreamSuite) TestDefaultResponseHeaderTimeout() { ...@@ -55,25 +55,27 @@ func (s *HTTPUpstreamSuite) TestDefaultResponseHeaderTimeout() {
func (s *HTTPUpstreamSuite) TestCustomResponseHeaderTimeout() { func (s *HTTPUpstreamSuite) TestCustomResponseHeaderTimeout() {
s.cfg.Gateway = config.GatewayConfig{ResponseHeaderTimeout: 7} s.cfg.Gateway = config.GatewayConfig{ResponseHeaderTimeout: 7}
svc := s.newService() svc := s.newService()
entry := svc.getOrCreateClient("", 0, 0) entry := mustGetOrCreateClient(s.T(), svc, "", 0, 0)
transport, ok := entry.client.Transport.(*http.Transport) transport, ok := entry.client.Transport.(*http.Transport)
require.True(s.T(), ok, "expected *http.Transport") require.True(s.T(), ok, "expected *http.Transport")
require.Equal(s.T(), 7*time.Second, transport.ResponseHeaderTimeout, "ResponseHeaderTimeout mismatch") require.Equal(s.T(), 7*time.Second, transport.ResponseHeaderTimeout, "ResponseHeaderTimeout mismatch")
} }
// TestGetOrCreateClient_InvalidURLFallsBackToDirect 测试无效代理 URL 回退 // TestGetOrCreateClient_InvalidURLReturnsError 测试无效代理 URL 返回错误
// 验证解析失败时回退到直连模式 // 验证解析失败时拒绝回退到直连模式
func (s *HTTPUpstreamSuite) TestGetOrCreateClient_InvalidURLFallsBackToDirect() { func (s *HTTPUpstreamSuite) TestGetOrCreateClient_InvalidURLReturnsError() {
svc := s.newService() svc := s.newService()
entry := svc.getOrCreateClient("://bad-proxy-url", 1, 1) _, err := svc.getClientEntry("://bad-proxy-url", 1, 1, false, false)
require.Equal(s.T(), directProxyKey, entry.proxyKey, "expected direct proxy fallback") require.Error(s.T(), err, "expected error for invalid proxy URL")
} }
// TestNormalizeProxyURL_Canonicalizes 测试代理 URL 规范化 // TestNormalizeProxyURL_Canonicalizes 测试代理 URL 规范化
// 验证等价地址能够映射到同一缓存键 // 验证等价地址能够映射到同一缓存键
func (s *HTTPUpstreamSuite) TestNormalizeProxyURL_Canonicalizes() { func (s *HTTPUpstreamSuite) TestNormalizeProxyURL_Canonicalizes() {
key1, _ := normalizeProxyURL("http://proxy.local:8080") key1, _, err1 := normalizeProxyURL("http://proxy.local:8080")
key2, _ := normalizeProxyURL("http://proxy.local:8080/") require.NoError(s.T(), err1)
key2, _, err2 := normalizeProxyURL("http://proxy.local:8080/")
require.NoError(s.T(), err2)
require.Equal(s.T(), key1, key2, "expected normalized proxy keys to match") require.Equal(s.T(), key1, key2, "expected normalized proxy keys to match")
} }
...@@ -171,8 +173,8 @@ func (s *HTTPUpstreamSuite) TestAccountIsolation_DifferentAccounts() { ...@@ -171,8 +173,8 @@ func (s *HTTPUpstreamSuite) TestAccountIsolation_DifferentAccounts() {
s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount} s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount}
svc := s.newService() svc := s.newService()
// 同一代理，不同账户 // 同一代理，不同账户
entry1 := svc.getOrCreateClient("http://proxy.local:8080", 1, 3) entry1 := mustGetOrCreateClient(s.T(), svc, "http://proxy.local:8080", 1, 3)
entry2 := svc.getOrCreateClient("http://proxy.local:8080", 2, 3) entry2 := mustGetOrCreateClient(s.T(), svc, "http://proxy.local:8080", 2, 3)
require.NotSame(s.T(), entry1, entry2, "不同账号不应共享连接池") require.NotSame(s.T(), entry1, entry2, "不同账号不应共享连接池")
require.Equal(s.T(), 2, len(svc.clients), "账号隔离应缓存两个客户端") require.Equal(s.T(), 2, len(svc.clients), "账号隔离应缓存两个客户端")
} }
...@@ -183,8 +185,8 @@ func (s *HTTPUpstreamSuite) TestAccountProxyIsolation_DifferentProxy() { ...@@ -183,8 +185,8 @@ func (s *HTTPUpstreamSuite) TestAccountProxyIsolation_DifferentProxy() {
s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccountProxy} s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccountProxy}
svc := s.newService() svc := s.newService()
// 同一账户，不同代理 // 同一账户，不同代理
entry1 := svc.getOrCreateClient("http://proxy-a:8080", 1, 3) entry1 := mustGetOrCreateClient(s.T(), svc, "http://proxy-a:8080", 1, 3)
entry2 := svc.getOrCreateClient("http://proxy-b:8080", 1, 3) entry2 := mustGetOrCreateClient(s.T(), svc, "http://proxy-b:8080", 1, 3)
require.NotSame(s.T(), entry1, entry2, "账号+代理隔离应区分不同代理") require.NotSame(s.T(), entry1, entry2, "账号+代理隔离应区分不同代理")
require.Equal(s.T(), 2, len(svc.clients), "账号+代理隔离应缓存两个客户端") require.Equal(s.T(), 2, len(svc.clients), "账号+代理隔离应缓存两个客户端")
} }
...@@ -195,8 +197,8 @@ func (s *HTTPUpstreamSuite) TestAccountModeProxyChangeClearsPool() { ...@@ -195,8 +197,8 @@ func (s *HTTPUpstreamSuite) TestAccountModeProxyChangeClearsPool() {
s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount} s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount}
svc := s.newService() svc := s.newService()
// 同一账户，先后使用不同代理 // 同一账户，先后使用不同代理
entry1 := svc.getOrCreateClient("http://proxy-a:8080", 1, 3) entry1 := mustGetOrCreateClient(s.T(), svc, "http://proxy-a:8080", 1, 3)
entry2 := svc.getOrCreateClient("http://proxy-b:8080", 1, 3) entry2 := mustGetOrCreateClient(s.T(), svc, "http://proxy-b:8080", 1, 3)
require.NotSame(s.T(), entry1, entry2, "账号切换代理应创建新连接池") require.NotSame(s.T(), entry1, entry2, "账号切换代理应创建新连接池")
require.Equal(s.T(), 1, len(svc.clients), "账号模式下应仅保留一个连接池") require.Equal(s.T(), 1, len(svc.clients), "账号模式下应仅保留一个连接池")
require.False(s.T(), hasEntry(svc, entry1), "旧连接池应被清理") require.False(s.T(), hasEntry(svc, entry1), "旧连接池应被清理")
...@@ -208,7 +210,7 @@ func (s *HTTPUpstreamSuite) TestAccountConcurrencyOverridesPoolSettings() { ...@@ -208,7 +210,7 @@ func (s *HTTPUpstreamSuite) TestAccountConcurrencyOverridesPoolSettings() {
s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount} s.cfg.Gateway = config.GatewayConfig{ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount}
svc := s.newService() svc := s.newService()
// 账户并发数为 12 // 账户并发数为 12
entry := svc.getOrCreateClient("", 1, 12) entry := mustGetOrCreateClient(s.T(), svc, "", 1, 12)
transport, ok := entry.client.Transport.(*http.Transport) transport, ok := entry.client.Transport.(*http.Transport)
require.True(s.T(), ok, "expected *http.Transport") require.True(s.T(), ok, "expected *http.Transport")
// 连接池参数应与并发数一致 // 连接池参数应与并发数一致
...@@ -228,7 +230,7 @@ func (s *HTTPUpstreamSuite) TestAccountConcurrencyFallbackToDefault() { ...@@ -228,7 +230,7 @@ func (s *HTTPUpstreamSuite) TestAccountConcurrencyFallbackToDefault() {
} }
svc := s.newService() svc := s.newService()
// 账户并发数为 0，应使用全局配置 // 账户并发数为 0，应使用全局配置
entry := svc.getOrCreateClient("", 1, 0) entry := mustGetOrCreateClient(s.T(), svc, "", 1, 0)
transport, ok := entry.client.Transport.(*http.Transport) transport, ok := entry.client.Transport.(*http.Transport)
require.True(s.T(), ok, "expected *http.Transport") require.True(s.T(), ok, "expected *http.Transport")
require.Equal(s.T(), 66, transport.MaxConnsPerHost, "MaxConnsPerHost fallback mismatch") require.Equal(s.T(), 66, transport.MaxConnsPerHost, "MaxConnsPerHost fallback mismatch")
...@@ -245,12 +247,12 @@ func (s *HTTPUpstreamSuite) TestEvictOverLimitRemovesOldestIdle() { ...@@ -245,12 +247,12 @@ func (s *HTTPUpstreamSuite) TestEvictOverLimitRemovesOldestIdle() {
} }
svc := s.newService() svc := s.newService()
// 创建两个客户端，设置不同的最后使用时间 // 创建两个客户端，设置不同的最后使用时间
entry1 := svc.getOrCreateClient("http://proxy-a:8080", 1, 1) entry1 := mustGetOrCreateClient(s.T(), svc, "http://proxy-a:8080", 1, 1)
entry2 := svc.getOrCreateClient("http://proxy-b:8080", 2, 1) entry2 := mustGetOrCreateClient(s.T(), svc, "http://proxy-b:8080", 2, 1)
atomic.StoreInt64(&entry1.lastUsed, time.Now().Add(-2*time.Hour).UnixNano()) // 最久 atomic.StoreInt64(&entry1.lastUsed, time.Now().Add(-2*time.Hour).UnixNano()) // 最久
atomic.StoreInt64(&entry2.lastUsed, time.Now().Add(-time.Hour).UnixNano()) atomic.StoreInt64(&entry2.lastUsed, time.Now().Add(-time.Hour).UnixNano())
// 创建第三个客户端，触发淘汰 // 创建第三个客户端，触发淘汰
_ = svc.getOrCreateClient("http://proxy-c:8080", 3, 1) _ = mustGetOrCreateClient(s.T(), svc, "http://proxy-c:8080", 3, 1)
require.LessOrEqual(s.T(), len(svc.clients), 2, "应保持在缓存上限内") require.LessOrEqual(s.T(), len(svc.clients), 2, "应保持在缓存上限内")
require.False(s.T(), hasEntry(svc, entry1), "最久未使用的连接池应被清理") require.False(s.T(), hasEntry(svc, entry1), "最久未使用的连接池应被清理")
...@@ -264,12 +266,12 @@ func (s *HTTPUpstreamSuite) TestIdleTTLDoesNotEvictActive() { ...@@ -264,12 +266,12 @@ func (s *HTTPUpstreamSuite) TestIdleTTLDoesNotEvictActive() {
ClientIdleTTLSeconds: 1, // 1 秒空闲超时 ClientIdleTTLSeconds: 1, // 1 秒空闲超时
} }
svc := s.newService() svc := s.newService()
entry1 := svc.getOrCreateClient("", 1, 1) entry1 := mustGetOrCreateClient(s.T(), svc, "", 1, 1)
// 设置为很久之前使用，但有活跃请求 // 设置为很久之前使用，但有活跃请求
atomic.StoreInt64(&entry1.lastUsed, time.Now().Add(-2*time.Minute).UnixNano()) atomic.StoreInt64(&entry1.lastUsed, time.Now().Add(-2*time.Minute).UnixNano())
atomic.StoreInt64(&entry1.inFlight, 1) // 模拟有活跃请求 atomic.StoreInt64(&entry1.inFlight, 1) // 模拟有活跃请求
// 创建新客户端，触发淘汰检查 // 创建新客户端，触发淘汰检查
_ = svc.getOrCreateClient("", 2, 1) _, _ = svc.getOrCreateClient("", 2, 1)
require.True(s.T(), hasEntry(svc, entry1), "有活跃请求时不应回收") require.True(s.T(), hasEntry(svc, entry1), "有活跃请求时不应回收")
} }
...@@ -279,6 +281,14 @@ func TestHTTPUpstreamSuite(t *testing.T) { ...@@ -279,6 +281,14 @@ func TestHTTPUpstreamSuite(t *testing.T) {
suite.Run(t, new(HTTPUpstreamSuite)) suite.Run(t, new(HTTPUpstreamSuite))
} }
// mustGetOrCreateClient 测试辅助函数，调用 getOrCreateClient 并断言无错误
func mustGetOrCreateClient(t *testing.T, svc *httpUpstreamService, proxyURL string, accountID int64, concurrency int) *upstreamClientEntry {
t.Helper()
entry, err := svc.getOrCreateClient(proxyURL, accountID, concurrency)
require.NoError(t, err, "getOrCreateClient(%q, %d, %d)", proxyURL, accountID, concurrency)
return entry
}
// hasEntry 检查客户端是否存在于缓存中 // hasEntry 检查客户端是否存在于缓存中
// 辅助函数，用于验证淘汰逻辑 // 辅助函数，用于验证淘汰逻辑
func hasEntry(svc *httpUpstreamService, target *upstreamClientEntry) bool { func hasEntry(svc *httpUpstreamService, target *upstreamClientEntry) bool {
......
backend/internal/repository/idempotency_repo.go 0 → 100644
View file @ 3d79773b
package repository
import (
"context"
"database/sql"
"errors"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/service"
)
type idempotencyRepository struct {
sql sqlExecutor
}
func NewIdempotencyRepository(_ *dbent.Client, sqlDB *sql.DB) service.IdempotencyRepository {
return &idempotencyRepository{sql: sqlDB}
}
func (r *idempotencyRepository) CreateProcessing(ctx context.Context, record *service.IdempotencyRecord) (bool, error) {
if record == nil {
return false, nil
}
query := `
INSERT INTO idempotency_records (
scope, idempotency_key_hash, request_fingerprint, status, locked_until, expires_at
) VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (scope, idempotency_key_hash) DO NOTHING
RETURNING id, created_at, updated_at
`
var createdAt time.Time
var updatedAt time.Time
err := scanSingleRow(ctx, r.sql, query, []any{
record.Scope,
record.IdempotencyKeyHash,
record.RequestFingerprint,
record.Status,
record.LockedUntil,
record.ExpiresAt,
}, &record.ID, &createdAt, &updatedAt)
if errors.Is(err, sql.ErrNoRows) {
return false, nil
}
if err != nil {
return false, err
}
record.CreatedAt = createdAt
record.UpdatedAt = updatedAt
return true, nil
}
func (r *idempotencyRepository) GetByScopeAndKeyHash(ctx context.Context, scope, keyHash string) (*service.IdempotencyRecord, error) {
query := `
SELECT
id, scope, idempotency_key_hash, request_fingerprint, status, response_status,
response_body, error_reason, locked_until, expires_at, created_at, updated_at
FROM idempotency_records
WHERE scope = $1 AND idempotency_key_hash = $2
`
record := &service.IdempotencyRecord{}
var responseStatus sql.NullInt64
var responseBody sql.NullString
var errorReason sql.NullString
var lockedUntil sql.NullTime
err := scanSingleRow(ctx, r.sql, query, []any{scope, keyHash},
&record.ID,
&record.Scope,
&record.IdempotencyKeyHash,
&record.RequestFingerprint,
&record.Status,
&responseStatus,
&responseBody,
&errorReason,
&lockedUntil,
&record.ExpiresAt,
&record.CreatedAt,
&record.UpdatedAt,
)
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
if err != nil {
return nil, err
}
if responseStatus.Valid {
v := int(responseStatus.Int64)
record.ResponseStatus = &v
}
if responseBody.Valid {
v := responseBody.String
record.ResponseBody = &v
}
if errorReason.Valid {
v := errorReason.String
record.ErrorReason = &v
}
if lockedUntil.Valid {
v := lockedUntil.Time
record.LockedUntil = &v
}
return record, nil
}
func (r *idempotencyRepository) TryReclaim(
ctx context.Context,
id int64,
fromStatus string,
now, newLockedUntil, newExpiresAt time.Time,
) (bool, error) {
query := `
UPDATE idempotency_records
SET status = $2,
locked_until = $3,
error_reason = NULL,
updated_at = NOW(),
expires_at = $4
WHERE id = $1
AND status = $5
AND (locked_until IS NULL OR locked_until <= $6)
`
res, err := r.sql.ExecContext(ctx, query,
id,
service.IdempotencyStatusProcessing,
newLockedUntil,
newExpiresAt,
fromStatus,
now,
)
if err != nil {
return false, err
}
affected, err := res.RowsAffected()
if err != nil {
return false, err
}
return affected > 0, nil
}
func (r *idempotencyRepository) ExtendProcessingLock(
ctx context.Context,
id int64,
requestFingerprint string,
newLockedUntil,
newExpiresAt time.Time,
) (bool, error) {
query := `
UPDATE idempotency_records
SET locked_until = $2,
expires_at = $3,
updated_at = NOW()
WHERE id = $1
AND status = $4
AND request_fingerprint = $5
`
res, err := r.sql.ExecContext(
ctx,
query,
id,
newLockedUntil,
newExpiresAt,
service.IdempotencyStatusProcessing,
requestFingerprint,
)
if err != nil {
return false, err
}
affected, err := res.RowsAffected()
if err != nil {
return false, err
}
return affected > 0, nil
}
func (r *idempotencyRepository) MarkSucceeded(ctx context.Context, id int64, responseStatus int, responseBody string, expiresAt time.Time) error {
query := `
UPDATE idempotency_records
SET status = $2,
response_status = $3,
response_body = $4,
error_reason = NULL,
locked_until = NULL,
expires_at = $5,
updated_at = NOW()
WHERE id = $1
`
_, err := r.sql.ExecContext(ctx, query,
id,
service.IdempotencyStatusSucceeded,
responseStatus,
responseBody,
expiresAt,
)
return err
}
func (r *idempotencyRepository) MarkFailedRetryable(ctx context.Context, id int64, errorReason string, lockedUntil, expiresAt time.Time) error {
query := `
UPDATE idempotency_records
SET status = $2,
error_reason = $3,
locked_until = $4,
expires_at = $5,
updated_at = NOW()
WHERE id = $1
`
_, err := r.sql.ExecContext(ctx, query,
id,
service.IdempotencyStatusFailedRetryable,
errorReason,
lockedUntil,
expiresAt,
)
return err
}
func (r *idempotencyRepository) DeleteExpired(ctx context.Context, now time.Time, limit int) (int64, error) {
if limit <= 0 {
limit = 500
}
query := `
WITH victims AS (
SELECT id
FROM idempotency_records
WHERE expires_at <= $1
ORDER BY expires_at ASC
LIMIT $2
)
DELETE FROM idempotency_records
WHERE id IN (SELECT id FROM victims)
`
res, err := r.sql.ExecContext(ctx, query, now, limit)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment