Jeffrey Smith
3d4228f868
All checks were successful
Co-authored-by: Jeffrey Smith <jasafpro@gmail.com> Co-committed-by: Jeffrey Smith <jasafpro@gmail.com>
46 lines
1.4 KiB
Go
46 lines
1.4 KiB
Go
package sqlite
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// RateLimitStore manages rate limit counters in SQLite.
|
|
type RateLimitStore struct{}
|
|
|
|
func NewRateLimitStore() *RateLimitStore { return &RateLimitStore{} }
|
|
|
|
func (s *RateLimitStore) Allow(ctx context.Context, key string, rate float64, burst int) (bool, error) {
|
|
// SQLite: strftime truncates to the second for window bucketing.
|
|
window := time.Now().UTC().Truncate(time.Second).Format("2006-01-02 15:04:05")
|
|
|
|
// Upsert via INSERT OR REPLACE pattern. SQLite has no ON CONFLICT ... DO UPDATE
|
|
// with RETURNING, so read-then-write in a single connection (single-writer safe).
|
|
var tokens float64
|
|
err := DB.QueryRowContext(ctx, `
|
|
SELECT COALESCE(tokens, 0) FROM rate_limit_counters
|
|
WHERE key = ? AND window_start = ?`, key, window).Scan(&tokens)
|
|
if err != nil {
|
|
// No row yet — insert fresh
|
|
tokens = 0
|
|
}
|
|
tokens++
|
|
_, err = DB.ExecContext(ctx, `
|
|
INSERT INTO rate_limit_counters (key, window_start, tokens)
|
|
VALUES (?, ?, ?)
|
|
ON CONFLICT (key, window_start) DO UPDATE SET tokens = ?`,
|
|
key, window, tokens, tokens)
|
|
if err != nil {
|
|
return false, fmt.Errorf("rate limit upsert: %w", err)
|
|
}
|
|
return tokens <= float64(burst), nil
|
|
}
|
|
|
|
func (s *RateLimitStore) Cleanup(ctx context.Context, maxAge time.Duration) error {
|
|
cutoff := time.Now().UTC().Add(-maxAge).Format("2006-01-02 15:04:05")
|
|
_, err := DB.ExecContext(ctx, `
|
|
DELETE FROM rate_limit_counters WHERE window_start < ?`, cutoff)
|
|
return err
|
|
}
|