feat(api-key): add independent quota and expiration support

This feature allows API Keys to have their own quota limits and expiration
times, independent of the user's balance.

Backend:
- Add quota, quota_used, expires_at fields to api_key schema
- Implement IsExpired() and IsQuotaExhausted() checks in middleware
- Add ResetQuota and ClearExpiration API endpoints
- Integrate quota billing in gateway handlers (OpenAI, Anthropic, Gemini)
- Include quota/expiration fields in auth cache for performance
- Expiration check returns 403, quota exhausted returns 429

Frontend:
- Add quota and expiration inputs to key create/edit dialog
- Add quick-select buttons for expiration (+7, +30, +90 days)
- Add reset quota confirmation dialog
- Add expires_at column to keys list
- Add i18n translations for new features (en/zh)

Migration:
- Add 045_add_api_key_quota.sql for new columns
This commit is contained in:
bayma888
2026-02-03 19:01:49 +08:00
parent bb3df5785a
commit 6146be1474
32 changed files with 1804 additions and 172 deletions

View File

@@ -1,4 +1,4 @@
FROM golang:1.25.5-alpine
FROM golang:1.25.6-alpine
WORKDIR /app
@@ -15,7 +15,7 @@ RUN go mod download
COPY . .
# 构建应用
RUN go build -o main cmd/server/main.go
RUN go build -o main ./cmd/server/
# 暴露端口
EXPOSE 8080

View File

@@ -173,8 +173,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
userAttributeService := service.NewUserAttributeService(userAttributeDefinitionRepository, userAttributeValueRepository)
userAttributeHandler := admin.NewUserAttributeHandler(userAttributeService)
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler)
gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, configConfig)
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, configConfig)
gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, configConfig)
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, configConfig)
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo)
totpHandler := handler.NewTotpHandler(totpService)
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler)

View File

@@ -40,6 +40,12 @@ type APIKey struct {
IPWhitelist []string `json:"ip_whitelist,omitempty"`
// Blocked IPs/CIDRs
IPBlacklist []string `json:"ip_blacklist,omitempty"`
// Quota limit in USD for this API key (0 = unlimited)
Quota float64 `json:"quota,omitempty"`
// Used quota amount in USD
QuotaUsed float64 `json:"quota_used,omitempty"`
// Expiration time for this API key (null = never expires)
ExpiresAt *time.Time `json:"expires_at,omitempty"`
// Edges holds the relations/edges for other nodes in the graph.
// The values are being populated by the APIKeyQuery when eager-loading is set.
Edges APIKeyEdges `json:"edges"`
@@ -97,11 +103,13 @@ func (*APIKey) scanValues(columns []string) ([]any, error) {
switch columns[i] {
case apikey.FieldIPWhitelist, apikey.FieldIPBlacklist:
values[i] = new([]byte)
case apikey.FieldQuota, apikey.FieldQuotaUsed:
values[i] = new(sql.NullFloat64)
case apikey.FieldID, apikey.FieldUserID, apikey.FieldGroupID:
values[i] = new(sql.NullInt64)
case apikey.FieldKey, apikey.FieldName, apikey.FieldStatus:
values[i] = new(sql.NullString)
case apikey.FieldCreatedAt, apikey.FieldUpdatedAt, apikey.FieldDeletedAt:
case apikey.FieldCreatedAt, apikey.FieldUpdatedAt, apikey.FieldDeletedAt, apikey.FieldExpiresAt:
values[i] = new(sql.NullTime)
default:
values[i] = new(sql.UnknownType)
@@ -190,6 +198,25 @@ func (_m *APIKey) assignValues(columns []string, values []any) error {
return fmt.Errorf("unmarshal field ip_blacklist: %w", err)
}
}
case apikey.FieldQuota:
if value, ok := values[i].(*sql.NullFloat64); !ok {
return fmt.Errorf("unexpected type %T for field quota", values[i])
} else if value.Valid {
_m.Quota = value.Float64
}
case apikey.FieldQuotaUsed:
if value, ok := values[i].(*sql.NullFloat64); !ok {
return fmt.Errorf("unexpected type %T for field quota_used", values[i])
} else if value.Valid {
_m.QuotaUsed = value.Float64
}
case apikey.FieldExpiresAt:
if value, ok := values[i].(*sql.NullTime); !ok {
return fmt.Errorf("unexpected type %T for field expires_at", values[i])
} else if value.Valid {
_m.ExpiresAt = new(time.Time)
*_m.ExpiresAt = value.Time
}
default:
_m.selectValues.Set(columns[i], values[i])
}
@@ -274,6 +301,17 @@ func (_m *APIKey) String() string {
builder.WriteString(", ")
builder.WriteString("ip_blacklist=")
builder.WriteString(fmt.Sprintf("%v", _m.IPBlacklist))
builder.WriteString(", ")
builder.WriteString("quota=")
builder.WriteString(fmt.Sprintf("%v", _m.Quota))
builder.WriteString(", ")
builder.WriteString("quota_used=")
builder.WriteString(fmt.Sprintf("%v", _m.QuotaUsed))
builder.WriteString(", ")
if v := _m.ExpiresAt; v != nil {
builder.WriteString("expires_at=")
builder.WriteString(v.Format(time.ANSIC))
}
builder.WriteByte(')')
return builder.String()
}

View File

@@ -35,6 +35,12 @@ const (
FieldIPWhitelist = "ip_whitelist"
// FieldIPBlacklist holds the string denoting the ip_blacklist field in the database.
FieldIPBlacklist = "ip_blacklist"
// FieldQuota holds the string denoting the quota field in the database.
FieldQuota = "quota"
// FieldQuotaUsed holds the string denoting the quota_used field in the database.
FieldQuotaUsed = "quota_used"
// FieldExpiresAt holds the string denoting the expires_at field in the database.
FieldExpiresAt = "expires_at"
// EdgeUser holds the string denoting the user edge name in mutations.
EdgeUser = "user"
// EdgeGroup holds the string denoting the group edge name in mutations.
@@ -79,6 +85,9 @@ var Columns = []string{
FieldStatus,
FieldIPWhitelist,
FieldIPBlacklist,
FieldQuota,
FieldQuotaUsed,
FieldExpiresAt,
}
// ValidColumn reports if the column name is valid (part of the table columns).
@@ -113,6 +122,10 @@ var (
DefaultStatus string
// StatusValidator is a validator for the "status" field. It is called by the builders before save.
StatusValidator func(string) error
// DefaultQuota holds the default value on creation for the "quota" field.
DefaultQuota float64
// DefaultQuotaUsed holds the default value on creation for the "quota_used" field.
DefaultQuotaUsed float64
)
// OrderOption defines the ordering options for the APIKey queries.
@@ -163,6 +176,21 @@ func ByStatus(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldStatus, opts...).ToFunc()
}
// ByQuota orders the results by the quota field.
func ByQuota(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldQuota, opts...).ToFunc()
}
// ByQuotaUsed orders the results by the quota_used field.
func ByQuotaUsed(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldQuotaUsed, opts...).ToFunc()
}
// ByExpiresAt orders the results by the expires_at field.
func ByExpiresAt(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldExpiresAt, opts...).ToFunc()
}
// ByUserField orders the results by user field.
func ByUserField(field string, opts ...sql.OrderTermOption) OrderOption {
return func(s *sql.Selector) {

View File

@@ -95,6 +95,21 @@ func Status(v string) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldStatus, v))
}
// Quota applies equality check predicate on the "quota" field. It's identical to QuotaEQ.
func Quota(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldQuota, v))
}
// QuotaUsed applies equality check predicate on the "quota_used" field. It's identical to QuotaUsedEQ.
func QuotaUsed(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldQuotaUsed, v))
}
// ExpiresAt applies equality check predicate on the "expires_at" field. It's identical to ExpiresAtEQ.
func ExpiresAt(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldExpiresAt, v))
}
// CreatedAtEQ applies the EQ predicate on the "created_at" field.
func CreatedAtEQ(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldCreatedAt, v))
@@ -490,6 +505,136 @@ func IPBlacklistNotNil() predicate.APIKey {
return predicate.APIKey(sql.FieldNotNull(FieldIPBlacklist))
}
// QuotaEQ applies the EQ predicate on the "quota" field.
func QuotaEQ(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldQuota, v))
}
// QuotaNEQ applies the NEQ predicate on the "quota" field.
func QuotaNEQ(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldNEQ(FieldQuota, v))
}
// QuotaIn applies the In predicate on the "quota" field.
func QuotaIn(vs ...float64) predicate.APIKey {
return predicate.APIKey(sql.FieldIn(FieldQuota, vs...))
}
// QuotaNotIn applies the NotIn predicate on the "quota" field.
func QuotaNotIn(vs ...float64) predicate.APIKey {
return predicate.APIKey(sql.FieldNotIn(FieldQuota, vs...))
}
// QuotaGT applies the GT predicate on the "quota" field.
func QuotaGT(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldGT(FieldQuota, v))
}
// QuotaGTE applies the GTE predicate on the "quota" field.
func QuotaGTE(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldGTE(FieldQuota, v))
}
// QuotaLT applies the LT predicate on the "quota" field.
func QuotaLT(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldLT(FieldQuota, v))
}
// QuotaLTE applies the LTE predicate on the "quota" field.
func QuotaLTE(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldLTE(FieldQuota, v))
}
// QuotaUsedEQ applies the EQ predicate on the "quota_used" field.
func QuotaUsedEQ(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldQuotaUsed, v))
}
// QuotaUsedNEQ applies the NEQ predicate on the "quota_used" field.
func QuotaUsedNEQ(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldNEQ(FieldQuotaUsed, v))
}
// QuotaUsedIn applies the In predicate on the "quota_used" field.
func QuotaUsedIn(vs ...float64) predicate.APIKey {
return predicate.APIKey(sql.FieldIn(FieldQuotaUsed, vs...))
}
// QuotaUsedNotIn applies the NotIn predicate on the "quota_used" field.
func QuotaUsedNotIn(vs ...float64) predicate.APIKey {
return predicate.APIKey(sql.FieldNotIn(FieldQuotaUsed, vs...))
}
// QuotaUsedGT applies the GT predicate on the "quota_used" field.
func QuotaUsedGT(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldGT(FieldQuotaUsed, v))
}
// QuotaUsedGTE applies the GTE predicate on the "quota_used" field.
func QuotaUsedGTE(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldGTE(FieldQuotaUsed, v))
}
// QuotaUsedLT applies the LT predicate on the "quota_used" field.
func QuotaUsedLT(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldLT(FieldQuotaUsed, v))
}
// QuotaUsedLTE applies the LTE predicate on the "quota_used" field.
func QuotaUsedLTE(v float64) predicate.APIKey {
return predicate.APIKey(sql.FieldLTE(FieldQuotaUsed, v))
}
// ExpiresAtEQ applies the EQ predicate on the "expires_at" field.
func ExpiresAtEQ(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldEQ(FieldExpiresAt, v))
}
// ExpiresAtNEQ applies the NEQ predicate on the "expires_at" field.
func ExpiresAtNEQ(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldNEQ(FieldExpiresAt, v))
}
// ExpiresAtIn applies the In predicate on the "expires_at" field.
func ExpiresAtIn(vs ...time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldIn(FieldExpiresAt, vs...))
}
// ExpiresAtNotIn applies the NotIn predicate on the "expires_at" field.
func ExpiresAtNotIn(vs ...time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldNotIn(FieldExpiresAt, vs...))
}
// ExpiresAtGT applies the GT predicate on the "expires_at" field.
func ExpiresAtGT(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldGT(FieldExpiresAt, v))
}
// ExpiresAtGTE applies the GTE predicate on the "expires_at" field.
func ExpiresAtGTE(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldGTE(FieldExpiresAt, v))
}
// ExpiresAtLT applies the LT predicate on the "expires_at" field.
func ExpiresAtLT(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldLT(FieldExpiresAt, v))
}
// ExpiresAtLTE applies the LTE predicate on the "expires_at" field.
func ExpiresAtLTE(v time.Time) predicate.APIKey {
return predicate.APIKey(sql.FieldLTE(FieldExpiresAt, v))
}
// ExpiresAtIsNil applies the IsNil predicate on the "expires_at" field.
func ExpiresAtIsNil() predicate.APIKey {
return predicate.APIKey(sql.FieldIsNull(FieldExpiresAt))
}
// ExpiresAtNotNil applies the NotNil predicate on the "expires_at" field.
func ExpiresAtNotNil() predicate.APIKey {
return predicate.APIKey(sql.FieldNotNull(FieldExpiresAt))
}
// HasUser applies the HasEdge predicate on the "user" edge.
func HasUser() predicate.APIKey {
return predicate.APIKey(func(s *sql.Selector) {

View File

@@ -125,6 +125,48 @@ func (_c *APIKeyCreate) SetIPBlacklist(v []string) *APIKeyCreate {
return _c
}
// SetQuota sets the "quota" field.
func (_c *APIKeyCreate) SetQuota(v float64) *APIKeyCreate {
_c.mutation.SetQuota(v)
return _c
}
// SetNillableQuota sets the "quota" field if the given value is not nil.
func (_c *APIKeyCreate) SetNillableQuota(v *float64) *APIKeyCreate {
if v != nil {
_c.SetQuota(*v)
}
return _c
}
// SetQuotaUsed sets the "quota_used" field.
func (_c *APIKeyCreate) SetQuotaUsed(v float64) *APIKeyCreate {
_c.mutation.SetQuotaUsed(v)
return _c
}
// SetNillableQuotaUsed sets the "quota_used" field if the given value is not nil.
func (_c *APIKeyCreate) SetNillableQuotaUsed(v *float64) *APIKeyCreate {
if v != nil {
_c.SetQuotaUsed(*v)
}
return _c
}
// SetExpiresAt sets the "expires_at" field.
func (_c *APIKeyCreate) SetExpiresAt(v time.Time) *APIKeyCreate {
_c.mutation.SetExpiresAt(v)
return _c
}
// SetNillableExpiresAt sets the "expires_at" field if the given value is not nil.
func (_c *APIKeyCreate) SetNillableExpiresAt(v *time.Time) *APIKeyCreate {
if v != nil {
_c.SetExpiresAt(*v)
}
return _c
}
// SetUser sets the "user" edge to the User entity.
func (_c *APIKeyCreate) SetUser(v *User) *APIKeyCreate {
return _c.SetUserID(v.ID)
@@ -205,6 +247,14 @@ func (_c *APIKeyCreate) defaults() error {
v := apikey.DefaultStatus
_c.mutation.SetStatus(v)
}
if _, ok := _c.mutation.Quota(); !ok {
v := apikey.DefaultQuota
_c.mutation.SetQuota(v)
}
if _, ok := _c.mutation.QuotaUsed(); !ok {
v := apikey.DefaultQuotaUsed
_c.mutation.SetQuotaUsed(v)
}
return nil
}
@@ -243,6 +293,12 @@ func (_c *APIKeyCreate) check() error {
return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "APIKey.status": %w`, err)}
}
}
if _, ok := _c.mutation.Quota(); !ok {
return &ValidationError{Name: "quota", err: errors.New(`ent: missing required field "APIKey.quota"`)}
}
if _, ok := _c.mutation.QuotaUsed(); !ok {
return &ValidationError{Name: "quota_used", err: errors.New(`ent: missing required field "APIKey.quota_used"`)}
}
if len(_c.mutation.UserIDs()) == 0 {
return &ValidationError{Name: "user", err: errors.New(`ent: missing required edge "APIKey.user"`)}
}
@@ -305,6 +361,18 @@ func (_c *APIKeyCreate) createSpec() (*APIKey, *sqlgraph.CreateSpec) {
_spec.SetField(apikey.FieldIPBlacklist, field.TypeJSON, value)
_node.IPBlacklist = value
}
if value, ok := _c.mutation.Quota(); ok {
_spec.SetField(apikey.FieldQuota, field.TypeFloat64, value)
_node.Quota = value
}
if value, ok := _c.mutation.QuotaUsed(); ok {
_spec.SetField(apikey.FieldQuotaUsed, field.TypeFloat64, value)
_node.QuotaUsed = value
}
if value, ok := _c.mutation.ExpiresAt(); ok {
_spec.SetField(apikey.FieldExpiresAt, field.TypeTime, value)
_node.ExpiresAt = &value
}
if nodes := _c.mutation.UserIDs(); len(nodes) > 0 {
edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.M2O,
@@ -539,6 +607,60 @@ func (u *APIKeyUpsert) ClearIPBlacklist() *APIKeyUpsert {
return u
}
// SetQuota sets the "quota" field.
func (u *APIKeyUpsert) SetQuota(v float64) *APIKeyUpsert {
u.Set(apikey.FieldQuota, v)
return u
}
// UpdateQuota sets the "quota" field to the value that was provided on create.
func (u *APIKeyUpsert) UpdateQuota() *APIKeyUpsert {
u.SetExcluded(apikey.FieldQuota)
return u
}
// AddQuota adds v to the "quota" field.
func (u *APIKeyUpsert) AddQuota(v float64) *APIKeyUpsert {
u.Add(apikey.FieldQuota, v)
return u
}
// SetQuotaUsed sets the "quota_used" field.
func (u *APIKeyUpsert) SetQuotaUsed(v float64) *APIKeyUpsert {
u.Set(apikey.FieldQuotaUsed, v)
return u
}
// UpdateQuotaUsed sets the "quota_used" field to the value that was provided on create.
func (u *APIKeyUpsert) UpdateQuotaUsed() *APIKeyUpsert {
u.SetExcluded(apikey.FieldQuotaUsed)
return u
}
// AddQuotaUsed adds v to the "quota_used" field.
func (u *APIKeyUpsert) AddQuotaUsed(v float64) *APIKeyUpsert {
u.Add(apikey.FieldQuotaUsed, v)
return u
}
// SetExpiresAt sets the "expires_at" field.
func (u *APIKeyUpsert) SetExpiresAt(v time.Time) *APIKeyUpsert {
u.Set(apikey.FieldExpiresAt, v)
return u
}
// UpdateExpiresAt sets the "expires_at" field to the value that was provided on create.
func (u *APIKeyUpsert) UpdateExpiresAt() *APIKeyUpsert {
u.SetExcluded(apikey.FieldExpiresAt)
return u
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (u *APIKeyUpsert) ClearExpiresAt() *APIKeyUpsert {
u.SetNull(apikey.FieldExpiresAt)
return u
}
// UpdateNewValues updates the mutable fields using the new values that were set on create.
// Using this option is equivalent to using:
//
@@ -738,6 +860,69 @@ func (u *APIKeyUpsertOne) ClearIPBlacklist() *APIKeyUpsertOne {
})
}
// SetQuota sets the "quota" field.
func (u *APIKeyUpsertOne) SetQuota(v float64) *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.SetQuota(v)
})
}
// AddQuota adds v to the "quota" field.
func (u *APIKeyUpsertOne) AddQuota(v float64) *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.AddQuota(v)
})
}
// UpdateQuota sets the "quota" field to the value that was provided on create.
func (u *APIKeyUpsertOne) UpdateQuota() *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateQuota()
})
}
// SetQuotaUsed sets the "quota_used" field.
func (u *APIKeyUpsertOne) SetQuotaUsed(v float64) *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.SetQuotaUsed(v)
})
}
// AddQuotaUsed adds v to the "quota_used" field.
func (u *APIKeyUpsertOne) AddQuotaUsed(v float64) *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.AddQuotaUsed(v)
})
}
// UpdateQuotaUsed sets the "quota_used" field to the value that was provided on create.
func (u *APIKeyUpsertOne) UpdateQuotaUsed() *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateQuotaUsed()
})
}
// SetExpiresAt sets the "expires_at" field.
func (u *APIKeyUpsertOne) SetExpiresAt(v time.Time) *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.SetExpiresAt(v)
})
}
// UpdateExpiresAt sets the "expires_at" field to the value that was provided on create.
func (u *APIKeyUpsertOne) UpdateExpiresAt() *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateExpiresAt()
})
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (u *APIKeyUpsertOne) ClearExpiresAt() *APIKeyUpsertOne {
return u.Update(func(s *APIKeyUpsert) {
s.ClearExpiresAt()
})
}
// Exec executes the query.
func (u *APIKeyUpsertOne) Exec(ctx context.Context) error {
if len(u.create.conflict) == 0 {
@@ -1103,6 +1288,69 @@ func (u *APIKeyUpsertBulk) ClearIPBlacklist() *APIKeyUpsertBulk {
})
}
// SetQuota sets the "quota" field.
func (u *APIKeyUpsertBulk) SetQuota(v float64) *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.SetQuota(v)
})
}
// AddQuota adds v to the "quota" field.
func (u *APIKeyUpsertBulk) AddQuota(v float64) *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.AddQuota(v)
})
}
// UpdateQuota sets the "quota" field to the value that was provided on create.
func (u *APIKeyUpsertBulk) UpdateQuota() *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateQuota()
})
}
// SetQuotaUsed sets the "quota_used" field.
func (u *APIKeyUpsertBulk) SetQuotaUsed(v float64) *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.SetQuotaUsed(v)
})
}
// AddQuotaUsed adds v to the "quota_used" field.
func (u *APIKeyUpsertBulk) AddQuotaUsed(v float64) *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.AddQuotaUsed(v)
})
}
// UpdateQuotaUsed sets the "quota_used" field to the value that was provided on create.
func (u *APIKeyUpsertBulk) UpdateQuotaUsed() *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateQuotaUsed()
})
}
// SetExpiresAt sets the "expires_at" field.
func (u *APIKeyUpsertBulk) SetExpiresAt(v time.Time) *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.SetExpiresAt(v)
})
}
// UpdateExpiresAt sets the "expires_at" field to the value that was provided on create.
func (u *APIKeyUpsertBulk) UpdateExpiresAt() *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.UpdateExpiresAt()
})
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (u *APIKeyUpsertBulk) ClearExpiresAt() *APIKeyUpsertBulk {
return u.Update(func(s *APIKeyUpsert) {
s.ClearExpiresAt()
})
}
// Exec executes the query.
func (u *APIKeyUpsertBulk) Exec(ctx context.Context) error {
if u.create.err != nil {

View File

@@ -170,6 +170,68 @@ func (_u *APIKeyUpdate) ClearIPBlacklist() *APIKeyUpdate {
return _u
}
// SetQuota sets the "quota" field.
func (_u *APIKeyUpdate) SetQuota(v float64) *APIKeyUpdate {
_u.mutation.ResetQuota()
_u.mutation.SetQuota(v)
return _u
}
// SetNillableQuota sets the "quota" field if the given value is not nil.
func (_u *APIKeyUpdate) SetNillableQuota(v *float64) *APIKeyUpdate {
if v != nil {
_u.SetQuota(*v)
}
return _u
}
// AddQuota adds value to the "quota" field.
func (_u *APIKeyUpdate) AddQuota(v float64) *APIKeyUpdate {
_u.mutation.AddQuota(v)
return _u
}
// SetQuotaUsed sets the "quota_used" field.
func (_u *APIKeyUpdate) SetQuotaUsed(v float64) *APIKeyUpdate {
_u.mutation.ResetQuotaUsed()
_u.mutation.SetQuotaUsed(v)
return _u
}
// SetNillableQuotaUsed sets the "quota_used" field if the given value is not nil.
func (_u *APIKeyUpdate) SetNillableQuotaUsed(v *float64) *APIKeyUpdate {
if v != nil {
_u.SetQuotaUsed(*v)
}
return _u
}
// AddQuotaUsed adds value to the "quota_used" field.
func (_u *APIKeyUpdate) AddQuotaUsed(v float64) *APIKeyUpdate {
_u.mutation.AddQuotaUsed(v)
return _u
}
// SetExpiresAt sets the "expires_at" field.
func (_u *APIKeyUpdate) SetExpiresAt(v time.Time) *APIKeyUpdate {
_u.mutation.SetExpiresAt(v)
return _u
}
// SetNillableExpiresAt sets the "expires_at" field if the given value is not nil.
func (_u *APIKeyUpdate) SetNillableExpiresAt(v *time.Time) *APIKeyUpdate {
if v != nil {
_u.SetExpiresAt(*v)
}
return _u
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (_u *APIKeyUpdate) ClearExpiresAt() *APIKeyUpdate {
_u.mutation.ClearExpiresAt()
return _u
}
// SetUser sets the "user" edge to the User entity.
func (_u *APIKeyUpdate) SetUser(v *User) *APIKeyUpdate {
return _u.SetUserID(v.ID)
@@ -350,6 +412,24 @@ func (_u *APIKeyUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if _u.mutation.IPBlacklistCleared() {
_spec.ClearField(apikey.FieldIPBlacklist, field.TypeJSON)
}
if value, ok := _u.mutation.Quota(); ok {
_spec.SetField(apikey.FieldQuota, field.TypeFloat64, value)
}
if value, ok := _u.mutation.AddedQuota(); ok {
_spec.AddField(apikey.FieldQuota, field.TypeFloat64, value)
}
if value, ok := _u.mutation.QuotaUsed(); ok {
_spec.SetField(apikey.FieldQuotaUsed, field.TypeFloat64, value)
}
if value, ok := _u.mutation.AddedQuotaUsed(); ok {
_spec.AddField(apikey.FieldQuotaUsed, field.TypeFloat64, value)
}
if value, ok := _u.mutation.ExpiresAt(); ok {
_spec.SetField(apikey.FieldExpiresAt, field.TypeTime, value)
}
if _u.mutation.ExpiresAtCleared() {
_spec.ClearField(apikey.FieldExpiresAt, field.TypeTime)
}
if _u.mutation.UserCleared() {
edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.M2O,
@@ -611,6 +691,68 @@ func (_u *APIKeyUpdateOne) ClearIPBlacklist() *APIKeyUpdateOne {
return _u
}
// SetQuota sets the "quota" field.
func (_u *APIKeyUpdateOne) SetQuota(v float64) *APIKeyUpdateOne {
_u.mutation.ResetQuota()
_u.mutation.SetQuota(v)
return _u
}
// SetNillableQuota sets the "quota" field if the given value is not nil.
func (_u *APIKeyUpdateOne) SetNillableQuota(v *float64) *APIKeyUpdateOne {
if v != nil {
_u.SetQuota(*v)
}
return _u
}
// AddQuota adds value to the "quota" field.
func (_u *APIKeyUpdateOne) AddQuota(v float64) *APIKeyUpdateOne {
_u.mutation.AddQuota(v)
return _u
}
// SetQuotaUsed sets the "quota_used" field.
func (_u *APIKeyUpdateOne) SetQuotaUsed(v float64) *APIKeyUpdateOne {
_u.mutation.ResetQuotaUsed()
_u.mutation.SetQuotaUsed(v)
return _u
}
// SetNillableQuotaUsed sets the "quota_used" field if the given value is not nil.
func (_u *APIKeyUpdateOne) SetNillableQuotaUsed(v *float64) *APIKeyUpdateOne {
if v != nil {
_u.SetQuotaUsed(*v)
}
return _u
}
// AddQuotaUsed adds value to the "quota_used" field.
func (_u *APIKeyUpdateOne) AddQuotaUsed(v float64) *APIKeyUpdateOne {
_u.mutation.AddQuotaUsed(v)
return _u
}
// SetExpiresAt sets the "expires_at" field.
func (_u *APIKeyUpdateOne) SetExpiresAt(v time.Time) *APIKeyUpdateOne {
_u.mutation.SetExpiresAt(v)
return _u
}
// SetNillableExpiresAt sets the "expires_at" field if the given value is not nil.
func (_u *APIKeyUpdateOne) SetNillableExpiresAt(v *time.Time) *APIKeyUpdateOne {
if v != nil {
_u.SetExpiresAt(*v)
}
return _u
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (_u *APIKeyUpdateOne) ClearExpiresAt() *APIKeyUpdateOne {
_u.mutation.ClearExpiresAt()
return _u
}
// SetUser sets the "user" edge to the User entity.
func (_u *APIKeyUpdateOne) SetUser(v *User) *APIKeyUpdateOne {
return _u.SetUserID(v.ID)
@@ -821,6 +963,24 @@ func (_u *APIKeyUpdateOne) sqlSave(ctx context.Context) (_node *APIKey, err erro
if _u.mutation.IPBlacklistCleared() {
_spec.ClearField(apikey.FieldIPBlacklist, field.TypeJSON)
}
if value, ok := _u.mutation.Quota(); ok {
_spec.SetField(apikey.FieldQuota, field.TypeFloat64, value)
}
if value, ok := _u.mutation.AddedQuota(); ok {
_spec.AddField(apikey.FieldQuota, field.TypeFloat64, value)
}
if value, ok := _u.mutation.QuotaUsed(); ok {
_spec.SetField(apikey.FieldQuotaUsed, field.TypeFloat64, value)
}
if value, ok := _u.mutation.AddedQuotaUsed(); ok {
_spec.AddField(apikey.FieldQuotaUsed, field.TypeFloat64, value)
}
if value, ok := _u.mutation.ExpiresAt(); ok {
_spec.SetField(apikey.FieldExpiresAt, field.TypeTime, value)
}
if _u.mutation.ExpiresAtCleared() {
_spec.ClearField(apikey.FieldExpiresAt, field.TypeTime)
}
if _u.mutation.UserCleared() {
edge := &sqlgraph.EdgeSpec{
Rel: sqlgraph.M2O,

View File

@@ -20,6 +20,9 @@ var (
{Name: "status", Type: field.TypeString, Size: 20, Default: "active"},
{Name: "ip_whitelist", Type: field.TypeJSON, Nullable: true},
{Name: "ip_blacklist", Type: field.TypeJSON, Nullable: true},
{Name: "quota", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
{Name: "quota_used", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
{Name: "expires_at", Type: field.TypeTime, Nullable: true},
{Name: "group_id", Type: field.TypeInt64, Nullable: true},
{Name: "user_id", Type: field.TypeInt64},
}
@@ -31,13 +34,13 @@ var (
ForeignKeys: []*schema.ForeignKey{
{
Symbol: "api_keys_groups_api_keys",
Columns: []*schema.Column{APIKeysColumns[9]},
Columns: []*schema.Column{APIKeysColumns[12]},
RefColumns: []*schema.Column{GroupsColumns[0]},
OnDelete: schema.SetNull,
},
{
Symbol: "api_keys_users_api_keys",
Columns: []*schema.Column{APIKeysColumns[10]},
Columns: []*schema.Column{APIKeysColumns[13]},
RefColumns: []*schema.Column{UsersColumns[0]},
OnDelete: schema.NoAction,
},
@@ -46,12 +49,12 @@ var (
{
Name: "apikey_user_id",
Unique: false,
Columns: []*schema.Column{APIKeysColumns[10]},
Columns: []*schema.Column{APIKeysColumns[13]},
},
{
Name: "apikey_group_id",
Unique: false,
Columns: []*schema.Column{APIKeysColumns[9]},
Columns: []*schema.Column{APIKeysColumns[12]},
},
{
Name: "apikey_status",
@@ -63,6 +66,16 @@ var (
Unique: false,
Columns: []*schema.Column{APIKeysColumns[3]},
},
{
Name: "apikey_quota_quota_used",
Unique: false,
Columns: []*schema.Column{APIKeysColumns[9], APIKeysColumns[10]},
},
{
Name: "apikey_expires_at",
Unique: false,
Columns: []*schema.Column{APIKeysColumns[11]},
},
},
}
// AccountsColumns holds the columns for the "accounts" table.

View File

@@ -79,6 +79,11 @@ type APIKeyMutation struct {
appendip_whitelist []string
ip_blacklist *[]string
appendip_blacklist []string
quota *float64
addquota *float64
quota_used *float64
addquota_used *float64
expires_at *time.Time
clearedFields map[string]struct{}
user *int64
cleareduser bool
@@ -634,6 +639,167 @@ func (m *APIKeyMutation) ResetIPBlacklist() {
delete(m.clearedFields, apikey.FieldIPBlacklist)
}
// SetQuota sets the "quota" field.
func (m *APIKeyMutation) SetQuota(f float64) {
m.quota = &f
m.addquota = nil
}
// Quota returns the value of the "quota" field in the mutation.
func (m *APIKeyMutation) Quota() (r float64, exists bool) {
v := m.quota
if v == nil {
return
}
return *v, true
}
// OldQuota returns the old "quota" field's value of the APIKey entity.
// If the APIKey object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *APIKeyMutation) OldQuota(ctx context.Context) (v float64, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldQuota is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldQuota requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldQuota: %w", err)
}
return oldValue.Quota, nil
}
// AddQuota adds f to the "quota" field.
func (m *APIKeyMutation) AddQuota(f float64) {
if m.addquota != nil {
*m.addquota += f
} else {
m.addquota = &f
}
}
// AddedQuota returns the value that was added to the "quota" field in this mutation.
func (m *APIKeyMutation) AddedQuota() (r float64, exists bool) {
v := m.addquota
if v == nil {
return
}
return *v, true
}
// ResetQuota resets all changes to the "quota" field.
func (m *APIKeyMutation) ResetQuota() {
m.quota = nil
m.addquota = nil
}
// SetQuotaUsed sets the "quota_used" field.
func (m *APIKeyMutation) SetQuotaUsed(f float64) {
m.quota_used = &f
m.addquota_used = nil
}
// QuotaUsed returns the value of the "quota_used" field in the mutation.
func (m *APIKeyMutation) QuotaUsed() (r float64, exists bool) {
v := m.quota_used
if v == nil {
return
}
return *v, true
}
// OldQuotaUsed returns the old "quota_used" field's value of the APIKey entity.
// If the APIKey object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *APIKeyMutation) OldQuotaUsed(ctx context.Context) (v float64, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldQuotaUsed is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldQuotaUsed requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldQuotaUsed: %w", err)
}
return oldValue.QuotaUsed, nil
}
// AddQuotaUsed adds f to the "quota_used" field.
func (m *APIKeyMutation) AddQuotaUsed(f float64) {
if m.addquota_used != nil {
*m.addquota_used += f
} else {
m.addquota_used = &f
}
}
// AddedQuotaUsed returns the value that was added to the "quota_used" field in this mutation.
func (m *APIKeyMutation) AddedQuotaUsed() (r float64, exists bool) {
v := m.addquota_used
if v == nil {
return
}
return *v, true
}
// ResetQuotaUsed resets all changes to the "quota_used" field.
func (m *APIKeyMutation) ResetQuotaUsed() {
m.quota_used = nil
m.addquota_used = nil
}
// SetExpiresAt sets the "expires_at" field.
func (m *APIKeyMutation) SetExpiresAt(t time.Time) {
m.expires_at = &t
}
// ExpiresAt returns the value of the "expires_at" field in the mutation.
func (m *APIKeyMutation) ExpiresAt() (r time.Time, exists bool) {
v := m.expires_at
if v == nil {
return
}
return *v, true
}
// OldExpiresAt returns the old "expires_at" field's value of the APIKey entity.
// If the APIKey object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *APIKeyMutation) OldExpiresAt(ctx context.Context) (v *time.Time, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldExpiresAt is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldExpiresAt requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldExpiresAt: %w", err)
}
return oldValue.ExpiresAt, nil
}
// ClearExpiresAt clears the value of the "expires_at" field.
func (m *APIKeyMutation) ClearExpiresAt() {
m.expires_at = nil
m.clearedFields[apikey.FieldExpiresAt] = struct{}{}
}
// ExpiresAtCleared returns if the "expires_at" field was cleared in this mutation.
func (m *APIKeyMutation) ExpiresAtCleared() bool {
_, ok := m.clearedFields[apikey.FieldExpiresAt]
return ok
}
// ResetExpiresAt resets all changes to the "expires_at" field.
func (m *APIKeyMutation) ResetExpiresAt() {
m.expires_at = nil
delete(m.clearedFields, apikey.FieldExpiresAt)
}
// ClearUser clears the "user" edge to the User entity.
func (m *APIKeyMutation) ClearUser() {
m.cleareduser = true
@@ -776,7 +942,7 @@ func (m *APIKeyMutation) Type() string {
// order to get all numeric fields that were incremented/decremented, call
// AddedFields().
func (m *APIKeyMutation) Fields() []string {
fields := make([]string, 0, 10)
fields := make([]string, 0, 13)
if m.created_at != nil {
fields = append(fields, apikey.FieldCreatedAt)
}
@@ -807,6 +973,15 @@ func (m *APIKeyMutation) Fields() []string {
if m.ip_blacklist != nil {
fields = append(fields, apikey.FieldIPBlacklist)
}
if m.quota != nil {
fields = append(fields, apikey.FieldQuota)
}
if m.quota_used != nil {
fields = append(fields, apikey.FieldQuotaUsed)
}
if m.expires_at != nil {
fields = append(fields, apikey.FieldExpiresAt)
}
return fields
}
@@ -835,6 +1010,12 @@ func (m *APIKeyMutation) Field(name string) (ent.Value, bool) {
return m.IPWhitelist()
case apikey.FieldIPBlacklist:
return m.IPBlacklist()
case apikey.FieldQuota:
return m.Quota()
case apikey.FieldQuotaUsed:
return m.QuotaUsed()
case apikey.FieldExpiresAt:
return m.ExpiresAt()
}
return nil, false
}
@@ -864,6 +1045,12 @@ func (m *APIKeyMutation) OldField(ctx context.Context, name string) (ent.Value,
return m.OldIPWhitelist(ctx)
case apikey.FieldIPBlacklist:
return m.OldIPBlacklist(ctx)
case apikey.FieldQuota:
return m.OldQuota(ctx)
case apikey.FieldQuotaUsed:
return m.OldQuotaUsed(ctx)
case apikey.FieldExpiresAt:
return m.OldExpiresAt(ctx)
}
return nil, fmt.Errorf("unknown APIKey field %s", name)
}
@@ -943,6 +1130,27 @@ func (m *APIKeyMutation) SetField(name string, value ent.Value) error {
}
m.SetIPBlacklist(v)
return nil
case apikey.FieldQuota:
v, ok := value.(float64)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetQuota(v)
return nil
case apikey.FieldQuotaUsed:
v, ok := value.(float64)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetQuotaUsed(v)
return nil
case apikey.FieldExpiresAt:
v, ok := value.(time.Time)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetExpiresAt(v)
return nil
}
return fmt.Errorf("unknown APIKey field %s", name)
}
@@ -951,6 +1159,12 @@ func (m *APIKeyMutation) SetField(name string, value ent.Value) error {
// this mutation.
func (m *APIKeyMutation) AddedFields() []string {
var fields []string
if m.addquota != nil {
fields = append(fields, apikey.FieldQuota)
}
if m.addquota_used != nil {
fields = append(fields, apikey.FieldQuotaUsed)
}
return fields
}
@@ -959,6 +1173,10 @@ func (m *APIKeyMutation) AddedFields() []string {
// was not set, or was not defined in the schema.
func (m *APIKeyMutation) AddedField(name string) (ent.Value, bool) {
switch name {
case apikey.FieldQuota:
return m.AddedQuota()
case apikey.FieldQuotaUsed:
return m.AddedQuotaUsed()
}
return nil, false
}
@@ -968,6 +1186,20 @@ func (m *APIKeyMutation) AddedField(name string) (ent.Value, bool) {
// type.
func (m *APIKeyMutation) AddField(name string, value ent.Value) error {
switch name {
case apikey.FieldQuota:
v, ok := value.(float64)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.AddQuota(v)
return nil
case apikey.FieldQuotaUsed:
v, ok := value.(float64)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.AddQuotaUsed(v)
return nil
}
return fmt.Errorf("unknown APIKey numeric field %s", name)
}
@@ -988,6 +1220,9 @@ func (m *APIKeyMutation) ClearedFields() []string {
if m.FieldCleared(apikey.FieldIPBlacklist) {
fields = append(fields, apikey.FieldIPBlacklist)
}
if m.FieldCleared(apikey.FieldExpiresAt) {
fields = append(fields, apikey.FieldExpiresAt)
}
return fields
}
@@ -1014,6 +1249,9 @@ func (m *APIKeyMutation) ClearField(name string) error {
case apikey.FieldIPBlacklist:
m.ClearIPBlacklist()
return nil
case apikey.FieldExpiresAt:
m.ClearExpiresAt()
return nil
}
return fmt.Errorf("unknown APIKey nullable field %s", name)
}
@@ -1052,6 +1290,15 @@ func (m *APIKeyMutation) ResetField(name string) error {
case apikey.FieldIPBlacklist:
m.ResetIPBlacklist()
return nil
case apikey.FieldQuota:
m.ResetQuota()
return nil
case apikey.FieldQuotaUsed:
m.ResetQuotaUsed()
return nil
case apikey.FieldExpiresAt:
m.ResetExpiresAt()
return nil
}
return fmt.Errorf("unknown APIKey field %s", name)
}

View File

@@ -91,6 +91,14 @@ func init() {
apikey.DefaultStatus = apikeyDescStatus.Default.(string)
// apikey.StatusValidator is a validator for the "status" field. It is called by the builders before save.
apikey.StatusValidator = apikeyDescStatus.Validators[0].(func(string) error)
// apikeyDescQuota is the schema descriptor for quota field.
apikeyDescQuota := apikeyFields[7].Descriptor()
// apikey.DefaultQuota holds the default value on creation for the quota field.
apikey.DefaultQuota = apikeyDescQuota.Default.(float64)
// apikeyDescQuotaUsed is the schema descriptor for quota_used field.
apikeyDescQuotaUsed := apikeyFields[8].Descriptor()
// apikey.DefaultQuotaUsed holds the default value on creation for the quota_used field.
apikey.DefaultQuotaUsed = apikeyDescQuotaUsed.Default.(float64)
accountMixin := schema.Account{}.Mixin()
accountMixinHooks1 := accountMixin[1].Hooks()
account.Hooks[0] = accountMixinHooks1[0]

View File

@@ -5,6 +5,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/domain"
"entgo.io/ent"
"entgo.io/ent/dialect"
"entgo.io/ent/dialect/entsql"
"entgo.io/ent/schema"
"entgo.io/ent/schema/edge"
@@ -52,6 +53,23 @@ func (APIKey) Fields() []ent.Field {
field.JSON("ip_blacklist", []string{}).
Optional().
Comment("Blocked IPs/CIDRs"),
// ========== Quota fields ==========
// Quota limit in USD (0 = unlimited)
field.Float("quota").
SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}).
Default(0).
Comment("Quota limit in USD for this API key (0 = unlimited)"),
// Used quota amount
field.Float("quota_used").
SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}).
Default(0).
Comment("Used quota amount in USD"),
// Expiration time (nil = never expires)
field.Time("expires_at").
Optional().
Nillable().
Comment("Expiration time for this API key (null = never expires)"),
}
}
@@ -77,5 +95,8 @@ func (APIKey) Indexes() []ent.Index {
index.Fields("group_id"),
index.Fields("status"),
index.Fields("deleted_at"),
// Index for quota queries
index.Fields("quota", "quota_used"),
index.Fields("expires_at"),
}
}

View File

@@ -3,6 +3,7 @@ package handler
import (
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
@@ -27,11 +28,13 @@ func NewAPIKeyHandler(apiKeyService *service.APIKeyService) *APIKeyHandler {
// CreateAPIKeyRequest represents the create API key request payload
type CreateAPIKeyRequest struct {
Name string `json:"name" binding:"required"`
GroupID *int64 `json:"group_id"` // nullable
CustomKey *string `json:"custom_key"` // 可选的自定义key
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
Name string `json:"name" binding:"required"`
GroupID *int64 `json:"group_id"` // nullable
CustomKey *string `json:"custom_key"` // 可选的自定义key
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
Quota *float64 `json:"quota"` // 配额限制 (USD)
ExpiresInDays *int `json:"expires_in_days"` // 过期天数
}
// UpdateAPIKeyRequest represents the update API key request payload
@@ -41,6 +44,9 @@ type UpdateAPIKeyRequest struct {
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
Quota *float64 `json:"quota"` // 配额限制 (USD), 0=无限制
ExpiresAt *string `json:"expires_at"` // 过期时间 (ISO 8601)
ResetQuota *bool `json:"reset_quota"` // 重置已用配额
}
// List handles listing user's API keys with pagination
@@ -114,11 +120,15 @@ func (h *APIKeyHandler) Create(c *gin.Context) {
}
svcReq := service.CreateAPIKeyRequest{
Name: req.Name,
GroupID: req.GroupID,
CustomKey: req.CustomKey,
IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist,
Name: req.Name,
GroupID: req.GroupID,
CustomKey: req.CustomKey,
IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist,
ExpiresInDays: req.ExpiresInDays,
}
if req.Quota != nil {
svcReq.Quota = *req.Quota
}
key, err := h.apiKeyService.Create(c.Request.Context(), subject.UserID, svcReq)
if err != nil {
@@ -153,6 +163,8 @@ func (h *APIKeyHandler) Update(c *gin.Context) {
svcReq := service.UpdateAPIKeyRequest{
IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist,
Quota: req.Quota,
ResetQuota: req.ResetQuota,
}
if req.Name != "" {
svcReq.Name = &req.Name
@@ -161,6 +173,21 @@ func (h *APIKeyHandler) Update(c *gin.Context) {
if req.Status != "" {
svcReq.Status = &req.Status
}
// Parse expires_at if provided
if req.ExpiresAt != nil {
if *req.ExpiresAt == "" {
// Empty string means clear expiration
svcReq.ExpiresAt = nil
svcReq.ClearExpiration = true
} else {
t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
if err != nil {
response.BadRequest(c, "Invalid expires_at format: "+err.Error())
return
}
svcReq.ExpiresAt = &t
}
}
key, err := h.apiKeyService.Update(c.Request.Context(), keyID, subject.UserID, svcReq)
if err != nil {

View File

@@ -76,6 +76,9 @@ func APIKeyFromService(k *service.APIKey) *APIKey {
Status: k.Status,
IPWhitelist: k.IPWhitelist,
IPBlacklist: k.IPBlacklist,
Quota: k.Quota,
QuotaUsed: k.QuotaUsed,
ExpiresAt: k.ExpiresAt,
CreatedAt: k.CreatedAt,
UpdatedAt: k.UpdatedAt,
User: UserFromServiceShallow(k.User),

View File

@@ -40,6 +40,9 @@ type APIKey struct {
Status string `json:"status"`
IPWhitelist []string `json:"ip_whitelist"`
IPBlacklist []string `json:"ip_blacklist"`
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
QuotaUsed float64 `json:"quota_used"` // Used quota amount in USD
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires)
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`

View File

@@ -31,6 +31,7 @@ type GatewayHandler struct {
userService *service.UserService
billingCacheService *service.BillingCacheService
usageService *service.UsageService
apiKeyService *service.APIKeyService
concurrencyHelper *ConcurrencyHelper
maxAccountSwitches int
maxAccountSwitchesGemini int
@@ -45,6 +46,7 @@ func NewGatewayHandler(
concurrencyService *service.ConcurrencyService,
billingCacheService *service.BillingCacheService,
usageService *service.UsageService,
apiKeyService *service.APIKeyService,
cfg *config.Config,
) *GatewayHandler {
pingInterval := time.Duration(0)
@@ -66,6 +68,7 @@ func NewGatewayHandler(
userService: userService,
billingCacheService: billingCacheService,
usageService: usageService,
apiKeyService: apiKeyService,
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude, pingInterval),
maxAccountSwitches: maxAccountSwitches,
maxAccountSwitchesGemini: maxAccountSwitchesGemini,
@@ -316,13 +319,14 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: clientIP,
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
}); err != nil {
log.Printf("Record usage failed: %v", err)
}
@@ -452,13 +456,14 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: clientIP,
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
}); err != nil {
log.Printf("Record usage failed: %v", err)
}

View File

@@ -381,6 +381,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
IPAddress: ip,
LongContextThreshold: 200000, // Gemini 200K 阈值
LongContextMultiplier: 2.0, // 超出部分双倍计费
APIKeyService: h.apiKeyService,
}); err != nil {
log.Printf("Record usage failed: %v", err)
}

View File

@@ -24,6 +24,7 @@ import (
type OpenAIGatewayHandler struct {
gatewayService *service.OpenAIGatewayService
billingCacheService *service.BillingCacheService
apiKeyService *service.APIKeyService
concurrencyHelper *ConcurrencyHelper
maxAccountSwitches int
}
@@ -33,6 +34,7 @@ func NewOpenAIGatewayHandler(
gatewayService *service.OpenAIGatewayService,
concurrencyService *service.ConcurrencyService,
billingCacheService *service.BillingCacheService,
apiKeyService *service.APIKeyService,
cfg *config.Config,
) *OpenAIGatewayHandler {
pingInterval := time.Duration(0)
@@ -46,6 +48,7 @@ func NewOpenAIGatewayHandler(
return &OpenAIGatewayHandler{
gatewayService: gatewayService,
billingCacheService: billingCacheService,
apiKeyService: apiKeyService,
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatComment, pingInterval),
maxAccountSwitches: maxAccountSwitches,
}
@@ -299,13 +302,14 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: ip,
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: usedAccount,
Subscription: subscription,
UserAgent: ua,
IPAddress: ip,
APIKeyService: h.apiKeyService,
}); err != nil {
log.Printf("Record usage failed: %v", err)
}

View File

@@ -33,7 +33,10 @@ func (r *apiKeyRepository) Create(ctx context.Context, key *service.APIKey) erro
SetKey(key.Key).
SetName(key.Name).
SetStatus(key.Status).
SetNillableGroupID(key.GroupID)
SetNillableGroupID(key.GroupID).
SetQuota(key.Quota).
SetQuotaUsed(key.QuotaUsed).
SetNillableExpiresAt(key.ExpiresAt)
if len(key.IPWhitelist) > 0 {
builder.SetIPWhitelist(key.IPWhitelist)
@@ -110,6 +113,9 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
apikey.FieldStatus,
apikey.FieldIPWhitelist,
apikey.FieldIPBlacklist,
apikey.FieldQuota,
apikey.FieldQuotaUsed,
apikey.FieldExpiresAt,
).
WithUser(func(q *dbent.UserQuery) {
q.Select(
@@ -161,6 +167,8 @@ func (r *apiKeyRepository) Update(ctx context.Context, key *service.APIKey) erro
Where(apikey.IDEQ(key.ID), apikey.DeletedAtIsNil()).
SetName(key.Name).
SetStatus(key.Status).
SetQuota(key.Quota).
SetQuotaUsed(key.QuotaUsed).
SetUpdatedAt(now)
if key.GroupID != nil {
builder.SetGroupID(*key.GroupID)
@@ -168,6 +176,13 @@ func (r *apiKeyRepository) Update(ctx context.Context, key *service.APIKey) erro
builder.ClearGroupID()
}
// Expiration time
if key.ExpiresAt != nil {
builder.SetExpiresAt(*key.ExpiresAt)
} else {
builder.ClearExpiresAt()
}
// IP 限制字段
if len(key.IPWhitelist) > 0 {
builder.SetIPWhitelist(key.IPWhitelist)
@@ -357,6 +372,38 @@ func (r *apiKeyRepository) ListKeysByGroupID(ctx context.Context, groupID int64)
return keys, nil
}
// IncrementQuotaUsed atomically increments the quota_used field and returns the new value
func (r *apiKeyRepository) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
// Use raw SQL for atomic increment to avoid race conditions
// First get current value
m, err := r.activeQuery().
Where(apikey.IDEQ(id)).
Select(apikey.FieldQuotaUsed).
Only(ctx)
if err != nil {
if dbent.IsNotFound(err) {
return 0, service.ErrAPIKeyNotFound
}
return 0, err
}
newValue := m.QuotaUsed + amount
// Update with new value
affected, err := r.client.APIKey.Update().
Where(apikey.IDEQ(id), apikey.DeletedAtIsNil()).
SetQuotaUsed(newValue).
Save(ctx)
if err != nil {
return 0, err
}
if affected == 0 {
return 0, service.ErrAPIKeyNotFound
}
return newValue, nil
}
func apiKeyEntityToService(m *dbent.APIKey) *service.APIKey {
if m == nil {
return nil
@@ -372,6 +419,9 @@ func apiKeyEntityToService(m *dbent.APIKey) *service.APIKey {
CreatedAt: m.CreatedAt,
UpdatedAt: m.UpdatedAt,
GroupID: m.GroupID,
Quota: m.Quota,
QuotaUsed: m.QuotaUsed,
ExpiresAt: m.ExpiresAt,
}
if m.Edges.User != nil {
out.User = userEntityToService(m.Edges.User)

View File

@@ -70,7 +70,27 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
// 检查API key是否激活
if !apiKey.IsActive() {
AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled")
// Provide more specific error message based on status
switch apiKey.Status {
case service.StatusAPIKeyQuotaExhausted:
AbortWithError(c, 429, "API_KEY_QUOTA_EXHAUSTED", "API key 额度已用完")
case service.StatusAPIKeyExpired:
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
default:
AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled")
}
return
}
// 检查API Key是否过期即使状态是active也要检查时间
if apiKey.IsExpired() {
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
return
}
// 检查API Key配额是否耗尽
if apiKey.IsQuotaExhausted() {
AbortWithError(c, 429, "API_KEY_QUOTA_EXHAUSTED", "API key 额度已用完")
return
}

View File

@@ -2,6 +2,14 @@ package service
import "time"
// API Key status constants
const (
StatusAPIKeyActive = "active"
StatusAPIKeyDisabled = "disabled"
StatusAPIKeyQuotaExhausted = "quota_exhausted"
StatusAPIKeyExpired = "expired"
)
type APIKey struct {
ID int64
UserID int64
@@ -15,8 +23,53 @@ type APIKey struct {
UpdatedAt time.Time
User *User
Group *Group
// Quota fields
Quota float64 // Quota limit in USD (0 = unlimited)
QuotaUsed float64 // Used quota amount
ExpiresAt *time.Time // Expiration time (nil = never expires)
}
func (k *APIKey) IsActive() bool {
return k.Status == StatusActive
}
// IsExpired checks if the API key has expired
func (k *APIKey) IsExpired() bool {
if k.ExpiresAt == nil {
return false
}
return time.Now().After(*k.ExpiresAt)
}
// IsQuotaExhausted checks if the API key quota is exhausted
func (k *APIKey) IsQuotaExhausted() bool {
if k.Quota <= 0 {
return false // unlimited
}
return k.QuotaUsed >= k.Quota
}
// GetQuotaRemaining returns remaining quota (-1 for unlimited)
func (k *APIKey) GetQuotaRemaining() float64 {
if k.Quota <= 0 {
return -1 // unlimited
}
remaining := k.Quota - k.QuotaUsed
if remaining < 0 {
return 0
}
return remaining
}
// GetDaysUntilExpiry returns days until expiry (-1 for never expires)
func (k *APIKey) GetDaysUntilExpiry() int {
if k.ExpiresAt == nil {
return -1 // never expires
}
duration := time.Until(*k.ExpiresAt)
if duration < 0 {
return 0
}
return int(duration.Hours() / 24)
}

View File

@@ -1,5 +1,7 @@
package service
import "time"
// APIKeyAuthSnapshot API Key 认证缓存快照(仅包含认证所需字段)
type APIKeyAuthSnapshot struct {
APIKeyID int64 `json:"api_key_id"`
@@ -10,6 +12,13 @@ type APIKeyAuthSnapshot struct {
IPBlacklist []string `json:"ip_blacklist,omitempty"`
User APIKeyAuthUserSnapshot `json:"user"`
Group *APIKeyAuthGroupSnapshot `json:"group,omitempty"`
// Quota fields for API Key independent quota feature
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
QuotaUsed float64 `json:"quota_used"` // Used quota amount
// Expiration field for API Key expiration feature
ExpiresAt *time.Time `json:"expires_at,omitempty"` // Expiration time (nil = never expires)
}
// APIKeyAuthUserSnapshot 用户快照

View File

@@ -213,6 +213,9 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot {
Status: apiKey.Status,
IPWhitelist: apiKey.IPWhitelist,
IPBlacklist: apiKey.IPBlacklist,
Quota: apiKey.Quota,
QuotaUsed: apiKey.QuotaUsed,
ExpiresAt: apiKey.ExpiresAt,
User: APIKeyAuthUserSnapshot{
ID: apiKey.User.ID,
Status: apiKey.User.Status,
@@ -256,6 +259,9 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
Status: snapshot.Status,
IPWhitelist: snapshot.IPWhitelist,
IPBlacklist: snapshot.IPBlacklist,
Quota: snapshot.Quota,
QuotaUsed: snapshot.QuotaUsed,
ExpiresAt: snapshot.ExpiresAt,
User: &User{
ID: snapshot.User.ID,
Status: snapshot.User.Status,

View File

@@ -17,13 +17,17 @@ import (
)
var (
ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found")
ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group")
ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists")
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found")
ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group")
ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists")
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
// ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期")
// ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted")
ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key 额度已用完")
)
const (
@@ -51,6 +55,9 @@ type APIKeyRepository interface {
CountByGroupID(ctx context.Context, groupID int64) (int64, error)
ListKeysByUserID(ctx context.Context, userID int64) ([]string, error)
ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error)
// Quota methods
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
}
// APIKeyCache defines cache operations for API key service
@@ -85,6 +92,10 @@ type CreateAPIKeyRequest struct {
CustomKey *string `json:"custom_key"` // 可选的自定义key
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
// Quota fields
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
}
// UpdateAPIKeyRequest 更新API Key请求
@@ -94,6 +105,12 @@ type UpdateAPIKeyRequest struct {
Status *string `json:"status"`
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单(空数组清空)
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单(空数组清空)
// Quota fields
Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited)
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
ClearExpiration bool `json:"-"` // Clear expiration (internal use)
ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
}
// APIKeyService API Key服务
@@ -289,6 +306,14 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK
Status: StatusActive,
IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist,
Quota: req.Quota,
QuotaUsed: 0,
}
// Set expiration time if specified
if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 {
expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays)
apiKey.ExpiresAt = &expiresAt
}
if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil {
@@ -436,6 +461,35 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
}
}
// Update quota fields
if req.Quota != nil {
apiKey.Quota = *req.Quota
// If quota is increased and status was quota_exhausted, reactivate
if apiKey.Status == StatusAPIKeyQuotaExhausted && *req.Quota > apiKey.QuotaUsed {
apiKey.Status = StatusActive
}
}
if req.ResetQuota != nil && *req.ResetQuota {
apiKey.QuotaUsed = 0
// If resetting quota and status was quota_exhausted, reactivate
if apiKey.Status == StatusAPIKeyQuotaExhausted {
apiKey.Status = StatusActive
}
}
if req.ClearExpiration {
apiKey.ExpiresAt = nil
// If clearing expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired {
apiKey.Status = StatusActive
}
} else if req.ExpiresAt != nil {
apiKey.ExpiresAt = req.ExpiresAt
// If extending expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired && time.Now().Before(*req.ExpiresAt) {
apiKey.Status = StatusActive
}
}
// 更新 IP 限制(空数组会清空设置)
apiKey.IPWhitelist = req.IPWhitelist
apiKey.IPBlacklist = req.IPBlacklist
@@ -572,3 +626,51 @@ func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword
}
return keys, nil
}
// CheckAPIKeyQuotaAndExpiry checks if the API key is valid for use (not expired, quota not exhausted)
// Returns nil if valid, error if invalid
func (s *APIKeyService) CheckAPIKeyQuotaAndExpiry(apiKey *APIKey) error {
// Check expiration
if apiKey.IsExpired() {
return ErrAPIKeyExpired
}
// Check quota
if apiKey.IsQuotaExhausted() {
return ErrAPIKeyQuotaExhausted
}
return nil
}
// UpdateQuotaUsed updates the quota_used field after a request
// Also checks if quota is exhausted and updates status accordingly
func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error {
if cost <= 0 {
return nil
}
// Use repository to atomically increment quota_used
newQuotaUsed, err := s.apiKeyRepo.IncrementQuotaUsed(ctx, apiKeyID, cost)
if err != nil {
return fmt.Errorf("increment quota used: %w", err)
}
// Check if quota is now exhausted and update status if needed
apiKey, err := s.apiKeyRepo.GetByID(ctx, apiKeyID)
if err != nil {
return nil // Don't fail the request, just log
}
// If quota is set and now exhausted, update status
if apiKey.Quota > 0 && newQuotaUsed >= apiKey.Quota {
apiKey.Status = StatusAPIKeyQuotaExhausted
if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil {
return nil // Don't fail the request
}
// Invalidate cache so next request sees the new status
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
}
return nil
}

View File

@@ -4489,13 +4489,19 @@ func (s *GatewayService) replaceToolNamesInResponseBody(body []byte, toolNameMap
// RecordUsageInput 记录使用量的输入参数
type RecordUsageInput struct {
Result *ForwardResult
APIKey *APIKey
User *User
Account *Account
Subscription *UserSubscription // 可选:订阅信息
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
Result *ForwardResult
APIKey *APIKey
User *User
Account *Account
Subscription *UserSubscription // 可选:订阅信息
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
APIKeyService APIKeyQuotaUpdater // 可选用于更新API Key配额
}
// APIKeyQuotaUpdater defines the interface for updating API Key quota
type APIKeyQuotaUpdater interface {
UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error
}
// RecordUsage 记录使用量并扣费(或更新订阅用量)
@@ -4635,6 +4641,13 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
}
}
// 更新 API Key 配额(如果设置了配额限制)
if shouldBill && cost.ActualCost > 0 && apiKey.Quota > 0 && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Update API key quota failed: %v", err)
}
}
// Schedule batch update for account last_used_at
s.deferredService.ScheduleLastUsedUpdate(account.ID)
@@ -4652,6 +4665,7 @@ type RecordUsageLongContextInput struct {
IPAddress string // 请求的客户端 IP 地址
LongContextThreshold int // 长上下文阈值(如 200000
LongContextMultiplier float64 // 超出阈值部分的倍率(如 2.0
APIKeyService *APIKeyService // API Key 配额服务(可选)
}
// RecordUsageWithLongContext 记录使用量并扣费,支持长上下文双倍计费(用于 Gemini
@@ -4788,6 +4802,12 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *
}
// 异步更新余额缓存
s.billingCacheService.QueueDeductBalance(user.ID, cost.ActualCost)
// API Key 独立配额扣费
if input.APIKeyService != nil && apiKey.Quota > 0 {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Add API key quota used failed: %v", err)
}
}
}
}

View File

@@ -1681,13 +1681,14 @@ func (s *OpenAIGatewayService) replaceModelInResponseBody(body []byte, fromModel
// OpenAIRecordUsageInput input for recording usage
type OpenAIRecordUsageInput struct {
Result *OpenAIForwardResult
APIKey *APIKey
User *User
Account *Account
Subscription *UserSubscription
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
Result *OpenAIForwardResult
APIKey *APIKey
User *User
Account *Account
Subscription *UserSubscription
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
APIKeyService APIKeyQuotaUpdater
}
// RecordUsage records usage and deducts balance
@@ -1799,6 +1800,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
}
}
// Update API key quota if applicable (only for balance mode with quota set)
if shouldBill && cost.ActualCost > 0 && apiKey.Quota > 0 && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Update API key quota failed: %v", err)
}
}
// Schedule batch update for account last_used_at
s.deferredService.ScheduleLastUsedUpdate(account.ID)

View File

@@ -0,0 +1,20 @@
-- Migration: Add quota fields to api_keys table
-- This migration adds independent quota and expiration support for API keys
-- Add quota limit field (0 = unlimited)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS quota DECIMAL(20, 8) NOT NULL DEFAULT 0;
-- Add used quota amount field
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS quota_used DECIMAL(20, 8) NOT NULL DEFAULT 0;
-- Add expiration time field (NULL = never expires)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS expires_at TIMESTAMPTZ;
-- Add indexes for efficient quota queries
CREATE INDEX IF NOT EXISTS idx_api_keys_quota_quota_used ON api_keys(quota, quota_used) WHERE deleted_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_api_keys_expires_at ON api_keys(expires_at) WHERE deleted_at IS NULL;
-- Comment on columns for documentation
COMMENT ON COLUMN api_keys.quota IS 'Quota limit in USD for this API key (0 = unlimited)';
COMMENT ON COLUMN api_keys.quota_used IS 'Used quota amount in USD';
COMMENT ON COLUMN api_keys.expires_at IS 'Expiration time for this API key (null = never expires)';

View File

@@ -1,9 +0,0 @@
//go:build tools
// +build tools
package tools
import (
_ "entgo.io/ent/cmd/ent"
_ "github.com/google/wire/cmd/wire"
)