完善 API Key 能力范围可视化和维护

2026-06-07 19:01:32 +08:00
parent dc14866210
commit f47132a653
19 changed files with 1165 additions and 49 deletions
@@ -0,0 +1,14 @@
+package store
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestNormalizeAPIKeyScopes(t *testing.T) {
+	got := normalizeAPIKeyScopes([]string{" Chat ", "AUDIO", "", "chat", "*", "all", " text_to_speech "})
+	want := []string{"chat", "audio", "all", "text_to_speech"}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("normalize scopes = %#v, want %#v", got, want)
+	}
+}
@@ -22,9 +22,34 @@ type Store struct {
 	pool *pgxpool.Pool
 }

+func defaultAPIKeyScopes() []string {
+	return []string{"chat", "embedding", "rerank", "image", "video", "music", "audio"}
+}
+
+func normalizeAPIKeyScopes(scopes []string) []string {
+	out := make([]string, 0, len(scopes))
+	seen := make(map[string]struct{}, len(scopes))
+	for _, scope := range scopes {
+		scope = strings.ToLower(strings.TrimSpace(scope))
+		if scope == "" {
+			continue
+		}
+		if scope == "*" {
+			scope = "all"
+		}
+		if _, ok := seen[scope]; ok {
+			continue
+		}
+		seen[scope] = struct{}{}
+		out = append(out, scope)
+	}
+	return out
+}
+
 var (
 	ErrInvalidCredentials        = errors.New("invalid account or password")
 	ErrInvalidInvitation         = errors.New("invalid or expired invitation code")
+	ErrInvalidAPIKeyScopes       = errors.New("api key scopes must not be empty")
 	ErrAccessRuleResourceDenied  = errors.New("access rule resource is not available")
 	ErrInsufficientWalletBalance = errors.New("insufficient wallet balance")
 	ErrLocalUserRequired         = errors.New("local gateway user is required")
@@ -106,6 +131,10 @@ type CreateAPIKeyInput struct {
 	ExpiresAt string   `json:"expiresAt"`
 }

+type UpdateAPIKeyScopesInput struct {
+	Scopes []string `json:"scopes"`
+}
+
 type APIKey struct {
 	ID              string         `json:"id"`
 	GatewayTenantID string         `json:"gatewayTenantId,omitempty"`
@@ -1156,6 +1185,32 @@ ORDER BY priority ASC, group_key ASC`)
 	return items, rows.Err()
 }

+func (s *Store) ListUserGroupsBySubject(ctx context.Context, groupIDs []string, groupKeys []string) ([]UserGroup, error) {
+	rows, err := s.pool.Query(ctx, `
+SELECT `+userGroupColumns+`
+FROM gateway_user_groups
+WHERE status = 'active'
+  AND (
+    id::text = ANY($1::text[])
+    OR group_key = ANY($2::text[])
+  )
+ORDER BY priority ASC, group_key ASC`, groupIDs, groupKeys)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	items := make([]UserGroup, 0)
+	for rows.Next() {
+		item, err := scanUserGroup(rows)
+		if err != nil {
+			return nil, err
+		}
+		items = append(items, item)
+	}
+	return items, rows.Err()
+}
+
 func (s *Store) ListAPIKeys(ctx context.Context, user *auth.User) ([]APIKey, error) {
 	gatewayUserID := localGatewayUserID(user)
 	if gatewayUserID == "" {
@@ -1201,7 +1256,7 @@ func (s *Store) ListPlayableAPIKeys(ctx context.Context, user *auth.User) ([]Pla
 	}
 	created, err := s.CreateAPIKey(ctx, CreateAPIKeyInput{
 		Name:   "Playground API Key",
-		Scopes: []string{"chat", "image", "video"},
+		Scopes: defaultAPIKeyScopes(),
 	}, user)
 	if err != nil {
 		return nil, err
@@ -1256,8 +1311,9 @@ func (s *Store) CreateAPIKey(ctx context.Context, input CreateAPIKeyInput, user
 	}
 	scopes := input.Scopes
 	if len(scopes) == 0 {
-		scopes = []string{"chat", "image", "video"}
+		scopes = defaultAPIKeyScopes()
 	}
+	scopes = normalizeAPIKeyScopes(scopes)
 	secret, err := generateAPIKeySecret()
 	if err != nil {
 		return CreatedAPIKey{}, err
@@ -1317,6 +1373,32 @@ RETURNING id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text
 	return CreatedAPIKey{APIKey: item, Secret: secret}, nil
 }

+func (s *Store) UpdateAPIKeyScopes(ctx context.Context, apiKeyID string, input UpdateAPIKeyScopesInput, user *auth.User) (APIKey, error) {
+	gatewayUserID := localGatewayUserID(user)
+	if gatewayUserID == "" {
+		return APIKey{}, ErrLocalUserRequired
+	}
+	scopes := normalizeAPIKeyScopes(input.Scopes)
+	if len(scopes) == 0 {
+		return APIKey{}, ErrInvalidAPIKeyScopes
+	}
+	scopesJSON, err := json.Marshal(scopes)
+	if err != nil {
+		return APIKey{}, err
+	}
+	return scanAPIKey(s.pool.QueryRow(ctx, `
+UPDATE gateway_api_keys
+SET scopes = $3::jsonb, updated_at = now()
+WHERE id = $1::uuid AND gateway_user_id = $2::uuid AND deleted_at IS NULL
+RETURNING id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
+       COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
+       key_prefix, name, scopes, COALESCE(user_group_id::text, ''),
+       rate_limit_policy, quota_policy, status, COALESCE(expires_at::text, ''),
+       COALESCE(last_used_at::text, ''), created_at, updated_at`,
+		apiKeyID, gatewayUserID, string(scopesJSON),
+	))
+}
+
 func (s *Store) DisableAPIKey(ctx context.Context, apiKeyID string, user *auth.User) (APIKey, error) {
 	gatewayUserID := localGatewayUserID(user)
 	if gatewayUserID == "" {
@@ -87,6 +87,15 @@ type WalletBalanceAdjustmentInput struct {
 	Metadata       map[string]any `json:"metadata"`
 }

+type WalletRechargeInput struct {
+	GatewayUserID  string         `json:"gatewayUserId"`
+	Currency       string         `json:"currency"`
+	Amount         float64        `json:"amount"`
+	Reason         string         `json:"reason"`
+	IdempotencyKey string         `json:"idempotencyKey"`
+	Metadata       map[string]any `json:"metadata"`
+}
+
 type WalletAdjustmentResult struct {
 	Account     GatewayWalletAccount     `json:"account"`
 	Before      GatewayWalletAccount     `json:"before"`
@@ -668,6 +677,100 @@ RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COA
 	return WalletAdjustmentResult{Account: locked, Before: before, Transaction: transaction}, nil
 }

+func (s *Store) RechargeUserWalletBalance(ctx context.Context, input WalletRechargeInput) (WalletAdjustmentResult, error) {
+	var result WalletAdjustmentResult
+	err := s.InTx(ctx, func(tx Tx) error {
+		next, err := s.RechargeUserWalletBalanceTx(ctx, tx, input)
+		if err != nil {
+			return err
+		}
+		result = next
+		return nil
+	})
+	return result, err
+}
+
+func (s *Store) RechargeUserWalletBalanceTx(ctx context.Context, tx Tx, input WalletRechargeInput) (WalletAdjustmentResult, error) {
+	input.GatewayUserID = strings.TrimSpace(input.GatewayUserID)
+	if input.GatewayUserID == "" {
+		return WalletAdjustmentResult{}, ErrLocalUserRequired
+	}
+	amount := roundMoney(input.Amount)
+	if amount <= 0 {
+		return WalletAdjustmentResult{}, fmt.Errorf("wallet recharge amount must be positive")
+	}
+	account, err := s.ensureWalletAccount(ctx, tx, input.GatewayUserID, input.Currency)
+	if err != nil {
+		return WalletAdjustmentResult{}, err
+	}
+	locked, err := scanWalletAccount(tx.QueryRow(ctx, `
+SELECT id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
+       COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
+       currency, balance::float8, frozen_balance::float8, total_recharged::float8,
+       total_spent::float8, status, metadata, created_at, updated_at
+FROM gateway_wallet_accounts
+WHERE id = $1::uuid
+FOR UPDATE`, account.ID))
+	if err != nil {
+		return WalletAdjustmentResult{}, err
+	}
+	before := locked
+	nextBalance := roundMoney(locked.Balance + amount)
+	reason := strings.TrimSpace(input.Reason)
+	if reason == "" {
+		reason = "后台余额充值"
+	}
+	if _, err := tx.Exec(ctx, `
+UPDATE gateway_wallet_accounts
+SET balance = $2,
+    total_recharged = total_recharged + $3,
+    updated_at = now()
+WHERE id = $1::uuid`,
+		locked.ID,
+		nextBalance,
+		amount,
+	); err != nil {
+		return WalletAdjustmentResult{}, err
+	}
+	metadata := mergeObjects(input.Metadata, map[string]any{
+		"reason":          reason,
+		"previousBalance": roundMoney(before.Balance),
+		"rechargeAmount":  amount,
+		"targetBalance":   nextBalance,
+	})
+	metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata))
+	transaction, err := scanWalletTransaction(tx.QueryRow(ctx, `
+INSERT INTO gateway_wallet_transactions (
+  account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
+  amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
+)
+VALUES (
+  $1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'recharge',
+  $4, $5, $6, NULLIF($7, ''), 'gateway_user', $8, $9::jsonb
+)
+RETURNING id::text, account_id::text, COALESCE(gateway_tenant_id::text, ''), COALESCE(gateway_user_id::text, ''),
+          direction, transaction_type, amount::float8, balance_before::float8, balance_after::float8,
+          COALESCE(idempotency_key, ''), COALESCE(reference_type, ''), COALESCE(reference_id, ''),
+          metadata, created_at`,
+		locked.ID,
+		locked.GatewayTenantID,
+		locked.GatewayUserID,
+		amount,
+		roundMoney(before.Balance),
+		nextBalance,
+		strings.TrimSpace(input.IdempotencyKey),
+		locked.GatewayUserID,
+		string(metadataJSON),
+	))
+	if err != nil {
+		return WalletAdjustmentResult{}, err
+	}
+	locked.Balance = nextBalance
+	locked.TotalRecharged = roundMoney(locked.TotalRecharged + amount)
+	locked.UpdatedAt = time.Now()
+	return WalletAdjustmentResult{Account: locked, Before: before, Transaction: transaction}, nil
+}
+
 func (s *Store) ensureWalletAccount(ctx context.Context, q Tx, gatewayUserID string, currency string) (GatewayWalletAccount, error) {
 	currency = normalizeWalletCurrency(currency)
 	if _, err := q.Exec(ctx, `