diff --git a/internal/key/key.go b/internal/key/key.go index 94eb537c..f90eb054 100644 --- a/internal/key/key.go +++ b/internal/key/key.go @@ -3,6 +3,7 @@ package key import ( "errors" "fmt" + "slices" "strings" "time" @@ -11,27 +12,41 @@ import ( const RevokedReasonExpired string = "expired" +type ExtendedBudgetLimitItem struct { + LimitInUsdOverTime float64 `json:"costLimitOverTime"` + Unit TimeUnit `json:"unit"` +} + +func (i *ExtendedBudgetLimitItem) ExtendedKey(k string) string { + return fmt.Sprintf("%s-%s", k, i.Unit) +} + +type ExtendedBudgetLimit struct { + Items []ExtendedBudgetLimitItem `json:"items"` +} + type UpdateKey struct { - Name string `json:"name"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - Revoked *bool `json:"revoked"` - RevokedReason string `json:"revokedReason"` - Key string `json:"key"` - SettingId string `json:"settingId"` - SettingIds []string `json:"settingIds"` - CostLimitInUsd *float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime *int `json:"rateLimitOverTime"` - RateLimitUnit *TimeUnit `json:"rateLimitUnit"` - RequestsLimit *int `json:"requestsLimit"` - AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"` - ShouldLogRequest *bool `json:"shouldLogRequest"` - ShouldLogResponse *bool `json:"shouldLogResponse"` - RotationEnabled *bool `json:"rotationEnabled"` - PolicyId *string `json:"policyId"` - IsKeyNotHashed *bool `json:"isKeyNotHashed"` + Name string `json:"name"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + Revoked *bool `json:"revoked"` + RevokedReason string `json:"revokedReason"` + Key string `json:"key"` + SettingId string `json:"settingId"` + SettingIds []string `json:"settingIds"` + CostLimitInUsd *float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime *int `json:"rateLimitOverTime"` + RateLimitUnit *TimeUnit `json:"rateLimitUnit"` + RequestsLimit *int `json:"requestsLimit"` + AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"` + ShouldLogRequest *bool `json:"shouldLogRequest"` + ShouldLogResponse *bool `json:"shouldLogResponse"` + RotationEnabled *bool `json:"rotationEnabled"` + PolicyId *string `json:"policyId"` + IsKeyNotHashed *bool `json:"isKeyNotHashed"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (uk *UpdateKey) Validate() error { @@ -125,7 +140,7 @@ func (uk *UpdateKey) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if *uk.RateLimitOverTime != 0 && *uk.RateLimitUnit != HourTimeUnit && *uk.RateLimitUnit != MinuteTimeUnit && *uk.RateLimitUnit != SecondTimeUnit && *uk.RateLimitUnit != DayTimeUnit { + if *uk.RateLimitOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -143,7 +158,7 @@ func (uk *UpdateKey) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if *uk.CostLimitInUsdOverTime != 0 && *uk.CostLimitInUsdUnit != DayTimeUnit && *uk.CostLimitInUsdUnit != HourTimeUnit && *uk.CostLimitInUsdUnit != MonthTimeUnit && *uk.CostLimitInUsdUnit != MinuteTimeUnit { + if *uk.CostLimitInUsdOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } @@ -157,28 +172,29 @@ type PathConfig struct { } type RequestKey struct { - Name string `json:"name"` - CreatedAt int64 `json:"createdAt"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - KeyId string `json:"keyId"` - Key string `json:"key"` - CostLimitInUsd float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime int `json:"rateLimitOverTime"` - RateLimitUnit TimeUnit `json:"rateLimitUnit"` - Ttl string `json:"ttl"` - KeyRing string `json:"keyRing"` - SettingId string `json:"settingId"` - AllowedPaths []PathConfig `json:"allowedPaths"` - SettingIds []string `json:"settingIds"` - ShouldLogRequest bool `json:"shouldLogRequest"` - ShouldLogResponse bool `json:"shouldLogResponse"` - RotationEnabled bool `json:"rotationEnabled"` - PolicyId string `json:"policyId"` - IsKeyNotHashed bool `json:"isKeyNotHashed"` - RequestsLimit int `json:"requestsLimit"` + Name string `json:"name"` + CreatedAt int64 `json:"createdAt"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + KeyId string `json:"keyId"` + Key string `json:"key"` + CostLimitInUsd float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime int `json:"rateLimitOverTime"` + RateLimitUnit TimeUnit `json:"rateLimitUnit"` + Ttl string `json:"ttl"` + KeyRing string `json:"keyRing"` + SettingId string `json:"settingId"` + AllowedPaths []PathConfig `json:"allowedPaths"` + SettingIds []string `json:"settingIds"` + ShouldLogRequest bool `json:"shouldLogRequest"` + ShouldLogResponse bool `json:"shouldLogResponse"` + RotationEnabled bool `json:"rotationEnabled"` + PolicyId string `json:"policyId"` + IsKeyNotHashed bool `json:"isKeyNotHashed"` + RequestsLimit int `json:"requestsLimit"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (rk *RequestKey) Validate() error { @@ -285,7 +301,7 @@ func (rk *RequestKey) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if rk.RateLimitUnit != HourTimeUnit && rk.RateLimitUnit != MinuteTimeUnit && rk.RateLimitUnit != SecondTimeUnit && rk.RateLimitUnit != DayTimeUnit { + if !slices.Contains(AllowedTimeUnits, rk.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -295,11 +311,40 @@ func (rk *RequestKey) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if rk.CostLimitInUsdUnit != DayTimeUnit && rk.CostLimitInUsdUnit != HourTimeUnit && rk.CostLimitInUsdUnit != MonthTimeUnit && rk.CostLimitInUsdUnit != MinuteTimeUnit { + if !slices.Contains(AllowedTimeUnits, rk.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } + if err := validateExtendedBudgetLimit(rk.ExtendedBudgetLimit); err != nil { + return err + } + return nil +} + +func validateExtendedBudgetLimit(ebl *ExtendedBudgetLimit) error { + if ebl == nil { + return nil + } + seenUnits := make(map[TimeUnit]struct{}, len(ebl.Items)) + for index, item := range ebl.Items { + if item.LimitInUsdOverTime < 0 { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].costLimitOverTime is invalid", index)) + } + + if len(item.Unit) == 0 { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is invalid", index)) + } + + if !slices.Contains(AllowedTimeUnits, item.Unit) { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit can not be identified", index)) + } + + if _, exists := seenUnits[item.Unit]; exists { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is duplicated", index)) + } + seenUnits[item.Unit] = struct{}{} + } return nil } @@ -310,34 +355,45 @@ const ( MinuteTimeUnit TimeUnit = "m" SecondTimeUnit TimeUnit = "s" DayTimeUnit TimeUnit = "d" + WeekTimeUnit TimeUnit = "w" MonthTimeUnit TimeUnit = "mo" ) +var AllowedTimeUnits = []TimeUnit{ + HourTimeUnit, + MinuteTimeUnit, + SecondTimeUnit, + DayTimeUnit, + WeekTimeUnit, + MonthTimeUnit, +} + type ResponseKey struct { - Name string `json:"name"` - CreatedAt int64 `json:"createdAt"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - KeyId string `json:"keyId"` - Revoked bool `json:"revoked"` - Key string `json:"key"` - RevokedReason string `json:"revokedReason"` - CostLimitInUsd float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime int `json:"rateLimitOverTime"` - RateLimitUnit TimeUnit `json:"rateLimitUnit"` - RequestsLimit int `json:"requestsLimit"` - Ttl string `json:"ttl"` - KeyRing string `json:"keyRing"` - SettingId string `json:"settingId"` - AllowedPaths []PathConfig `json:"allowedPaths"` - SettingIds []string `json:"settingIds"` - ShouldLogRequest bool `json:"shouldLogRequest"` - ShouldLogResponse bool `json:"shouldLogResponse"` - RotationEnabled bool `json:"rotationEnabled"` - PolicyId string `json:"policyId"` - IsKeyNotHashed bool `json:"isKeyNotHashed"` + Name string `json:"name"` + CreatedAt int64 `json:"createdAt"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + KeyId string `json:"keyId"` + Revoked bool `json:"revoked"` + Key string `json:"key"` + RevokedReason string `json:"revokedReason"` + CostLimitInUsd float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime int `json:"rateLimitOverTime"` + RateLimitUnit TimeUnit `json:"rateLimitUnit"` + RequestsLimit int `json:"requestsLimit"` + Ttl string `json:"ttl"` + KeyRing string `json:"keyRing"` + SettingId string `json:"settingId"` + AllowedPaths []PathConfig `json:"allowedPaths"` + SettingIds []string `json:"settingIds"` + ShouldLogRequest bool `json:"shouldLogRequest"` + ShouldLogResponse bool `json:"shouldLogResponse"` + RotationEnabled bool `json:"rotationEnabled"` + PolicyId string `json:"policyId"` + IsKeyNotHashed bool `json:"isKeyNotHashed"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (rk *ResponseKey) GetSettingIds() []string { diff --git a/internal/key/key_test.go b/internal/key/key_test.go new file mode 100644 index 00000000..bdacafaa --- /dev/null +++ b/internal/key/key_test.go @@ -0,0 +1,57 @@ +package key + +import "testing" + +func TestValidateExtendedBudgetLimit(t *testing.T) { + tests := []struct { + name string + ebl *ExtendedBudgetLimit + wantErr string + }{ + { + name: "nil limit is valid", + ebl: nil, + wantErr: "", + }, + { + name: "duplicate units are rejected", + ebl: &ExtendedBudgetLimit{ + Items: []ExtendedBudgetLimitItem{ + {LimitInUsdOverTime: 1, Unit: HourTimeUnit}, + {LimitInUsdOverTime: 2, Unit: HourTimeUnit}, + }, + }, + wantErr: "extendedBudgetLimit.items[1].unit is duplicated", + }, + { + name: "unique units are valid", + ebl: &ExtendedBudgetLimit{ + Items: []ExtendedBudgetLimitItem{ + {LimitInUsdOverTime: 1, Unit: HourTimeUnit}, + {LimitInUsdOverTime: 2, Unit: MinuteTimeUnit}, + }, + }, + wantErr: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateExtendedBudgetLimit(tt.ebl) + if tt.wantErr == "" { + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + return + } + + if err == nil { + t.Fatalf("expected error %q, got nil", tt.wantErr) + } + + if err.Error() != tt.wantErr { + t.Fatalf("expected error %q, got %q", tt.wantErr, err.Error()) + } + }) + } +} diff --git a/internal/message/consumer.go b/internal/message/consumer.go index 189043fd..17eaad32 100644 --- a/internal/message/consumer.go +++ b/internal/message/consumer.go @@ -18,7 +18,7 @@ type Consumer struct { } type recorder interface { - RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error + RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error RecordUserSpend(userId string, micros int64, costLimitUnit key.TimeUnit) error RecordEvent(e *event.Event) error RecordKeyRequestSpent(keyId string) error diff --git a/internal/message/handler.go b/internal/message/handler.go index 7c973ef3..c2576862 100644 --- a/internal/message/handler.go +++ b/internal/message/handler.go @@ -340,7 +340,7 @@ func (h *Handler) HandleEventWithRequestAndResponse(m Message) error { if e.Event.CostInUsd != 0 { micros := int64(e.Event.CostInUsd * 1000000) - err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit) + err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit, e.Key.ExtendedBudgetLimit) if err != nil { telemetry.Incr("bricksllm.message.handler.handle_event_with_request_and_response.record_key_spend_error", nil, 1) h.log.Debug("error when recording key spend", zap.Error(err)) diff --git a/internal/recorder/recorder.go b/internal/recorder/recorder.go index 5c8b73fe..3c3463bd 100644 --- a/internal/recorder/recorder.go +++ b/internal/recorder/recorder.go @@ -60,7 +60,7 @@ func (r *Recorder) RecordUserSpend(userId string, micros int64, costLimitUnit ke return nil } -func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error { +func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error { err := r.s.IncrementCounter(keyId, micros) if err != nil { return err @@ -73,6 +73,15 @@ func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key. } } + if extendedLimits != nil { + for _, item := range extendedLimits.Items { + err = r.c.IncrementCounter(item.ExtendedKey(keyId), item.Unit, int64(micros)) + if err != nil { + return err + } + } + } + return nil } diff --git a/internal/storage/postgresql/key.go b/internal/storage/postgresql/key.go index e92f114f..7ffa5f71 100644 --- a/internal/storage/postgresql/key.go +++ b/internal/storage/postgresql/key.go @@ -5,9 +5,10 @@ import ( "database/sql" "encoding/json" "fmt" - "github.com/bricks-cloud/bricksllm/internal/event" "strings" + "github.com/bricks-cloud/bricksllm/internal/event" + internal_errors "github.com/bricks-cloud/bricksllm/internal/errors" "github.com/bricks-cloud/bricksllm/internal/key" "github.com/lib/pq" @@ -57,7 +58,7 @@ func (s *Store) AlterKeysTable() error { END IF; END $$; - ALTER TABLE keys ADD COLUMN IF NOT EXISTS setting_id VARCHAR(255), ADD COLUMN IF NOT EXISTS allowed_paths JSONB, ADD COLUMN IF NOT EXISTS setting_ids VARCHAR(255)[] NOT NULL DEFAULT ARRAY[]::VARCHAR(255)[], ADD COLUMN IF NOT EXISTS should_log_request BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS should_log_response BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS rotation_enabled BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS policy_id VARCHAR(255) NOT NULL DEFAULT '', ADD COLUMN IF NOT EXISTS is_key_not_hashed BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS requests_limit INT NOT NULL DEFAULT 0; + ALTER TABLE keys ADD COLUMN IF NOT EXISTS setting_id VARCHAR(255), ADD COLUMN IF NOT EXISTS allowed_paths JSONB, ADD COLUMN IF NOT EXISTS setting_ids VARCHAR(255)[] NOT NULL DEFAULT ARRAY[]::VARCHAR(255)[], ADD COLUMN IF NOT EXISTS should_log_request BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS should_log_response BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS rotation_enabled BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS policy_id VARCHAR(255) NOT NULL DEFAULT '', ADD COLUMN IF NOT EXISTS is_key_not_hashed BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS requests_limit INT NOT NULL DEFAULT 0, ADD COLUMN IF NOT EXISTS extended_budget_limit JSONB; ` ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) @@ -191,6 +192,7 @@ func (s *Store) GetKeys(tags, keyIds []string, provider string) ([]*key.Response &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -317,6 +319,7 @@ func (s *Store) GetKeysV2(tags, keyIds []string, revoked *bool, limit, offset in &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -397,6 +400,7 @@ func (s *Store) GetKeyByHash(hash string) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ) if err != nil { @@ -462,6 +466,7 @@ func (s *Store) GetKey(keyId string) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -559,6 +564,7 @@ func (s *Store) GetSpentKeys(tags []string, order string, limit, offset int, val &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -625,6 +631,7 @@ func (s *Store) GetAllKeys() ([]*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -686,6 +693,7 @@ func (s *Store) GetUpdatedKeys(updatedAt int64) ([]*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } @@ -792,6 +800,12 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error counter++ } + if uk.ExtendedBudgetLimit != nil { + values = append(values, *uk.ExtendedBudgetLimit) + fields = append(fields, fmt.Sprintf("extended_budget_limit = $%d", counter)) + counter++ + } + if uk.ShouldLogRequest != nil { values = append(values, *uk.ShouldLogRequest) fields = append(fields, fmt.Sprintf("should_log_request = $%d", counter)) @@ -865,6 +879,7 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { if err == sql.ErrNoRows { return nil, internal_errors.NewNotFoundError(fmt.Sprintf("key not found for id: %s", id)) @@ -889,8 +904,8 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { query := ` - INSERT INTO keys (name, created_at, updated_at, tags, revoked, key_id, key, revoked_reason, cost_limit_in_usd, cost_limit_in_usd_over_time, cost_limit_in_usd_unit, rate_limit_over_time, rate_limit_unit, ttl, key_ring, setting_id, allowed_paths, setting_ids, should_log_request, should_log_response, rotation_enabled, policy_id, is_key_not_hashed, requests_limit) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24) + INSERT INTO keys (name, created_at, updated_at, tags, revoked, key_id, key, revoked_reason, cost_limit_in_usd, cost_limit_in_usd_over_time, cost_limit_in_usd_unit, rate_limit_over_time, rate_limit_unit, ttl, key_ring, setting_id, allowed_paths, setting_ids, should_log_request, should_log_response, rotation_enabled, policy_id, is_key_not_hashed, requests_limit, extended_budget_limit) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25) RETURNING *; ` @@ -924,6 +939,7 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { rk.PolicyId, rk.IsKeyNotHashed, rk.RequestsLimit, + rk.ExtendedBudgetLimit, } ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) @@ -958,6 +974,7 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &k.ExtendedBudgetLimit, ); err != nil { return nil, err } diff --git a/internal/storage/redis/cache.go b/internal/storage/redis/cache.go index a60f13b8..cd91a363 100644 --- a/internal/storage/redis/cache.go +++ b/internal/storage/redis/cache.go @@ -111,6 +111,8 @@ func getCounterTtl(rateLimitUnit key.TimeUnit) (time.Time, error) { return now.Truncate(60 * time.Minute).Add(time.Minute * 60).Add(-time.Millisecond), nil case key.DayTimeUnit: return now.Truncate(24 * time.Hour).Add(time.Hour * 24).Add(-time.Millisecond), nil + case key.WeekTimeUnit: + return now.Truncate(7 * 24 * time.Hour).Add(7 * 24 * time.Hour).Add(-time.Millisecond), nil case key.MonthTimeUnit: firstDayOfNextMonth := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC) return firstDayOfNextMonth.Add(-time.Millisecond), nil @@ -130,6 +132,8 @@ func getCounterTimeStamp(rateLimitUnit key.TimeUnit) (int64, error) { return int64(now.Minute()), nil case key.DayTimeUnit: return int64(now.Hour()), nil + case key.WeekTimeUnit: + return int64(now.Day()), nil case key.MonthTimeUnit: return int64(now.Day()), nil } diff --git a/internal/user/user.go b/internal/user/user.go index 0604c651..907bf1d1 100644 --- a/internal/user/user.go +++ b/internal/user/user.go @@ -2,6 +2,7 @@ package user import ( "fmt" + "slices" "strings" "time" @@ -113,7 +114,7 @@ func (u *User) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if u.RateLimitUnit != key.HourTimeUnit && u.RateLimitUnit != key.MinuteTimeUnit && u.RateLimitUnit != key.SecondTimeUnit && u.RateLimitUnit != key.DayTimeUnit { + if !slices.Contains(key.AllowedTimeUnits, u.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -123,7 +124,7 @@ func (u *User) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if u.CostLimitInUsdUnit != key.DayTimeUnit && u.CostLimitInUsdUnit != key.HourTimeUnit && u.CostLimitInUsdUnit != key.MonthTimeUnit && u.CostLimitInUsdUnit != key.MinuteTimeUnit { + if !slices.Contains(key.AllowedTimeUnits, u.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } @@ -228,7 +229,7 @@ func (uu *UpdateUser) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if *uu.RateLimitOverTime != 0 && *uu.RateLimitUnit != key.HourTimeUnit && *uu.RateLimitUnit != key.MinuteTimeUnit && *uu.RateLimitUnit != key.SecondTimeUnit && *uu.RateLimitUnit != key.DayTimeUnit { + if *uu.RateLimitOverTime != 0 && !slices.Contains(key.AllowedTimeUnits, *uu.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -246,7 +247,7 @@ func (uu *UpdateUser) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if *uu.CostLimitInUsdOverTime != 0 && *uu.CostLimitInUsdUnit != key.DayTimeUnit && *uu.CostLimitInUsdUnit != key.HourTimeUnit && *uu.CostLimitInUsdUnit != key.MonthTimeUnit && *uu.CostLimitInUsdUnit != key.MinuteTimeUnit { + if *uu.CostLimitInUsdOverTime != 0 && !slices.Contains(key.AllowedTimeUnits, *uu.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } diff --git a/internal/validator/validator.go b/internal/validator/validator.go index c3df819a..d4eeb761 100644 --- a/internal/validator/validator.go +++ b/internal/validator/validator.go @@ -75,6 +75,11 @@ func (v *Validator) Validate(k *key.ResponseKey, promptCost float64) error { return err } + err = v.validateExtendedCostLimitOverTime(k.KeyId, k.ExtendedBudgetLimit) + if err != nil { + return err + } + err = v.validateCostLimit(k.KeyId, k.CostLimitInUsd) if err != nil { return err @@ -128,6 +133,26 @@ func (v *Validator) validateCostLimitOverTime(keyId string, costLimitOverTime fl return nil } +func (v *Validator) validateExtendedCostLimitOverTime(keyId string, limits *key.ExtendedBudgetLimit) error { + if limits == nil { + return nil + } + if limits.Items == nil || len(limits.Items) == 0 { + return nil + } + + for _, item := range limits.Items { + if item.LimitInUsdOverTime == 0 { + continue + } + err := v.validateCostLimitOverTime(item.ExtendedKey(keyId), item.LimitInUsdOverTime, item.Unit) + if err != nil { + return err + } + } + return nil +} + func convertDollarToMicroDollars(dollar float64) int64 { return int64(dollar * 1000000) }