shithub Public

Watch 1 Fork 0 Star 0
Go · 6751 bytes Raw Blame History
1 // Code generated by sqlc. DO NOT EDIT.
2 // versions:
3 // sqlc v1.31.1
4 // source: rate_limits.sql
5
6 package ratelimitdb
7
8 import (
9 "context"
10 "net/netip"
11
12 "github.com/jackc/pgx/v5/pgtype"
13 )
14
15 const acquireRateLimitLease = `-- name: AcquireRateLimitLease :one
16 INSERT INTO rate_limits (scope, key, hits, window_started_at)
17 VALUES ($1, $2, 1, now())
18 ON CONFLICT (scope, key)
19 DO UPDATE SET
20 hits = CASE
21 WHEN rate_limits.window_started_at < now() - $3::interval
22 THEN 1
23 ELSE rate_limits.hits + 1
24 END,
25 window_started_at = CASE
26 WHEN rate_limits.window_started_at < now() - $3::interval
27 THEN now()
28 ELSE rate_limits.window_started_at
29 END
30 WHERE rate_limits.window_started_at < now() - $3::interval
31 OR rate_limits.hits < $4::integer
32 RETURNING hits, window_started_at
33 `
34
35 type AcquireRateLimitLeaseParams struct {
36 Scope string
37 Key string
38 Ttl pgtype.Interval
39 MaxHits int32
40 }
41
42 type AcquireRateLimitLeaseRow struct {
43 Hits int32
44 WindowStartedAt pgtype.Timestamptz
45 }
46
47 // Concurrent-lease variant for long-lived streams. `hits` is the
48 // currently-held lease count. The ttl rolls stale rows forward so a process
49 // crash or severed TCP connection cannot consume capacity indefinitely.
50 func (q *Queries) AcquireRateLimitLease(ctx context.Context, db DBTX, arg AcquireRateLimitLeaseParams) (AcquireRateLimitLeaseRow, error) {
51 row := db.QueryRow(ctx, acquireRateLimitLease,
52 arg.Scope,
53 arg.Key,
54 arg.Ttl,
55 arg.MaxHits,
56 )
57 var i AcquireRateLimitLeaseRow
58 err := row.Scan(&i.Hits, &i.WindowStartedAt)
59 return i, err
60 }
61
62 const bumpRateLimit = `-- name: BumpRateLimit :one
63
64 INSERT INTO rate_limits (scope, key, hits, window_started_at)
65 VALUES ($1, $2, 1, now())
66 ON CONFLICT (scope, key)
67 DO UPDATE SET
68 hits = CASE
69 WHEN rate_limits.window_started_at < now() - $3::interval
70 THEN 1
71 ELSE rate_limits.hits + 1
72 END,
73 window_started_at = CASE
74 WHEN rate_limits.window_started_at < now() - $3::interval
75 THEN now()
76 ELSE rate_limits.window_started_at
77 END
78 RETURNING hits, window_started_at
79 `
80
81 type BumpRateLimitParams struct {
82 Scope string
83 Key string
84 Ttl pgtype.Interval
85 }
86
87 type BumpRateLimitRow struct {
88 Hits int32
89 WindowStartedAt pgtype.Timestamptz
90 }
91
92 // SPDX-License-Identifier: AGPL-3.0-or-later
93 //
94 // Generic rate-limit counter queries (S35). Two write paths:
95 // - BumpRateLimit — atomic UPSERT that rolls the window forward
96 // when stale, increments hits otherwise. Returns
97 // the post-update hits + window_started_at so
98 // the caller can compute Retry-After without a
99 // second round trip.
100 // - BumpSignupIPThrottle — same shape against signup_ip_throttle
101 // keyed by inet/CIDR.
102 //
103 // Reads (PeekRateLimit, PeekSignupIPThrottle) are kept around for
104 // the admin observability surface; the hot path uses Bump-and-decide.
105 // Roll-or-increment in one statement. The CASE in the UPDATE branch
106 // handles the window roll: when the existing window started before
107 // (now - $3 interval), we treat it as a new window and reset hits
108 // to 1; otherwise we increment in place.
109 func (q *Queries) BumpRateLimit(ctx context.Context, db DBTX, arg BumpRateLimitParams) (BumpRateLimitRow, error) {
110 row := db.QueryRow(ctx, bumpRateLimit, arg.Scope, arg.Key, arg.Ttl)
111 var i BumpRateLimitRow
112 err := row.Scan(&i.Hits, &i.WindowStartedAt)
113 return i, err
114 }
115
116 const bumpSignupIPThrottle = `-- name: BumpSignupIPThrottle :one
117 INSERT INTO signup_ip_throttle (cidr, hits, window_started_at)
118 VALUES ($1, 1, now())
119 ON CONFLICT (cidr)
120 DO UPDATE SET
121 hits = CASE
122 WHEN signup_ip_throttle.window_started_at < now() - $2::interval
123 THEN 1
124 ELSE signup_ip_throttle.hits + 1
125 END,
126 window_started_at = CASE
127 WHEN signup_ip_throttle.window_started_at < now() - $2::interval
128 THEN now()
129 ELSE signup_ip_throttle.window_started_at
130 END
131 RETURNING hits, window_started_at
132 `
133
134 type BumpSignupIPThrottleParams struct {
135 Cidr netip.Addr
136 Ttl pgtype.Interval
137 }
138
139 type BumpSignupIPThrottleRow struct {
140 Hits int32
141 WindowStartedAt pgtype.Timestamptz
142 }
143
144 // Same UPSERT shape against the inet-keyed signup throttle.
145 func (q *Queries) BumpSignupIPThrottle(ctx context.Context, db DBTX, arg BumpSignupIPThrottleParams) (BumpSignupIPThrottleRow, error) {
146 row := db.QueryRow(ctx, bumpSignupIPThrottle, arg.Cidr, arg.Ttl)
147 var i BumpSignupIPThrottleRow
148 err := row.Scan(&i.Hits, &i.WindowStartedAt)
149 return i, err
150 }
151
152 const peekRateLimit = `-- name: PeekRateLimit :one
153 SELECT scope, key, hits, window_started_at
154 FROM rate_limits
155 WHERE scope = $1 AND key = $2
156 `
157
158 type PeekRateLimitParams struct {
159 Scope string
160 Key string
161 }
162
163 func (q *Queries) PeekRateLimit(ctx context.Context, db DBTX, arg PeekRateLimitParams) (RateLimit, error) {
164 row := db.QueryRow(ctx, peekRateLimit, arg.Scope, arg.Key)
165 var i RateLimit
166 err := row.Scan(
167 &i.Scope,
168 &i.Key,
169 &i.Hits,
170 &i.WindowStartedAt,
171 )
172 return i, err
173 }
174
175 const pruneRateLimits = `-- name: PruneRateLimits :execrows
176 DELETE FROM rate_limits
177 WHERE window_started_at < now() - $1::interval
178 `
179
180 func (q *Queries) PruneRateLimits(ctx context.Context, db DBTX, retention pgtype.Interval) (int64, error) {
181 result, err := db.Exec(ctx, pruneRateLimits, retention)
182 if err != nil {
183 return 0, err
184 }
185 return result.RowsAffected(), nil
186 }
187
188 const pruneSignupIPThrottle = `-- name: PruneSignupIPThrottle :execrows
189 DELETE FROM signup_ip_throttle
190 WHERE window_started_at < now() - $1::interval
191 `
192
193 func (q *Queries) PruneSignupIPThrottle(ctx context.Context, db DBTX, retention pgtype.Interval) (int64, error) {
194 result, err := db.Exec(ctx, pruneSignupIPThrottle, retention)
195 if err != nil {
196 return 0, err
197 }
198 return result.RowsAffected(), nil
199 }
200
201 const releaseRateLimitLease = `-- name: ReleaseRateLimitLease :execrows
202 UPDATE rate_limits
203 SET hits = GREATEST(hits - 1, 0)
204 WHERE scope = $1 AND key = $2
205 `
206
207 type ReleaseRateLimitLeaseParams struct {
208 Scope string
209 Key string
210 }
211
212 func (q *Queries) ReleaseRateLimitLease(ctx context.Context, db DBTX, arg ReleaseRateLimitLeaseParams) (int64, error) {
213 result, err := db.Exec(ctx, releaseRateLimitLease, arg.Scope, arg.Key)
214 if err != nil {
215 return 0, err
216 }
217 return result.RowsAffected(), nil
218 }
219