superidou
5759c1862e
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
43 lines
1.2 KiB
Go
43 lines
1.2 KiB
Go
package apikey
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type UsageStore struct{ redis *redis.Client }
|
|
|
|
func NewUsageStore(client *redis.Client) *UsageStore { return &UsageStore{redis: client} }
|
|
|
|
func MonthlyTokenUsageKey(apiKeyID string, now time.Time) string {
|
|
return fmt.Sprintf("gateway:usage:api-key:%s:tokens:%s", apiKeyID, now.UTC().Format("200601"))
|
|
}
|
|
|
|
func (s *UsageStore) MonthlyTokens(ctx context.Context, apiKeyIDs []string, now time.Time) (map[string]int64, error) {
|
|
usage := make(map[string]int64, len(apiKeyIDs))
|
|
if len(apiKeyIDs) == 0 || s == nil || s.redis == nil {
|
|
return usage, nil
|
|
}
|
|
pipe := s.redis.Pipeline()
|
|
commands := make(map[string]*redis.StringCmd, len(apiKeyIDs))
|
|
for _, id := range apiKeyIDs {
|
|
commands[id] = pipe.Get(ctx, MonthlyTokenUsageKey(id, now))
|
|
}
|
|
_, err := pipe.Exec(ctx)
|
|
if err != nil && err != redis.Nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
|
}
|
|
for id, command := range commands {
|
|
value, commandErr := command.Int64()
|
|
if commandErr == nil {
|
|
usage[id] = max(value, 0)
|
|
} else if commandErr != redis.Nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrStore, commandErr)
|
|
}
|
|
}
|
|
return usage, nil
|
|
}
|