Jeffrey Smith
2f31a69756
Removes standalone "// v0.X.X:" comment lines and inline trailing version annotations. Keeps version references in config field docs and compatibility notes where they describe requirements. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
36 lines
1.0 KiB
Go
36 lines
1.0 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// RateLimitStore manages distributed rate limit counters in Postgres.
|
|
type RateLimitStore struct{}
|
|
|
|
func NewRateLimitStore() *RateLimitStore { return &RateLimitStore{} }
|
|
|
|
func (s *RateLimitStore) Allow(ctx context.Context, key string, rate float64, burst int) (bool, error) {
|
|
// Upsert the current-second window counter and return the new token count.
|
|
var tokens float64
|
|
err := DB.QueryRowContext(ctx, `
|
|
INSERT INTO rate_limit_counters (key, window_start, tokens)
|
|
VALUES ($1, date_trunc('second', NOW()), 1)
|
|
ON CONFLICT (key, window_start)
|
|
DO UPDATE SET tokens = rate_limit_counters.tokens + 1
|
|
RETURNING tokens`, key).Scan(&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 {
|
|
_, err := DB.ExecContext(ctx, `
|
|
DELETE FROM rate_limit_counters
|
|
WHERE window_start < NOW() - $1 * INTERVAL '1 second'`,
|
|
int(maxAge.Seconds()))
|
|
return err
|
|
}
|