Skip to main content

max / makenotwork

8.3 KB · 216 lines History Blame Raw
1 //! Reusable keyset (cursor) pagination.
2 //!
3 //! Keyset, not offset: a paged query fetches only `limit + 1` rows via an
4 //! indexed range scan on `(sort_col, id)`, and derives the next page's cursor
5 //! from the last row, it never fetches the full set and slices in memory. The
6 //! cursor is **opaque** (base64) so a consumer's client contract does not change
7 //! if the encoded sort key later does.
8 //!
9 //! When to reach for this vs. offset:
10 //! - **Keyset (this module)**, forward "load more" / infinite-scroll lists,
11 //! large or hot tables, anywhere a full-scan-per-page would be wrong.
12 //! - **Offset**, a numbered-page admin UI that needs "page 7 of 42" and a
13 //! total count (e.g. the admin user list). Offset is correct there; keyset
14 //! cannot express page numbers or totals. Don't convert those.
15 //!
16 //! Consumer contract: a keyset query MUST have an index on `(<sort_col>, id)` in
17 //! its sort direction, or the `LIMIT n+1` degrades to sort-the-world. The
18 //! generic machinery here (cursor encode/decode, page assembly, next-page
19 //! detection) is written once; each consumer writes ~3 lines of SQL:
20 //!
21 //! ```ignore
22 //! // WHERE <base predicates>
23 //! // AND ($2::timestamptz IS NULL OR (created_at, id) < ($2, $3))
24 //! // ORDER BY created_at DESC, id DESC
25 //! // LIMIT $1 -- bind params.fetch_limit()
26 //! let rows = sqlx::query_as::<_, Row>(SQL)
27 //! .bind(params.fetch_limit())
28 //! .bind(params.after.as_ref().map(|c| c.ts))
29 //! .bind(params.after.as_ref().map(|c| c.id))
30 //! .fetch_all(pool).await?;
31 //! Ok(Page::from_rows(rows, &params, |r| Cursor::new(r.created_at, r.id)))
32 //! ```
33
34 use base64::Engine;
35 use chrono::{DateTime, Utc};
36 use serde::Serialize;
37 use uuid::Uuid;
38
39 /// Default page size when a consumer does not specify one.
40 pub const DEFAULT_PAGE_SIZE: u32 = 50;
41 /// Upper clamp on page size so a client can't request an unbounded page.
42 pub const MAX_PAGE_SIZE: u32 = 200;
43
44 /// An opaque forward cursor: the `(sort-timestamp, id)` of the last row on a
45 /// page. Ordering is `(ts, id)` DESC, with `id` as the unique tiebreaker so the
46 /// keyset is a total order (no row is skipped or repeated across pages when two
47 /// rows share a timestamp).
48 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
49 pub struct Cursor {
50 pub ts: DateTime<Utc>,
51 pub id: Uuid,
52 }
53
54 impl Cursor {
55 pub fn new(ts: DateTime<Utc>, id: Uuid) -> Self {
56 Self { ts, id }
57 }
58
59 /// Encode as an opaque base64 token: `<unix_millis>:<uuid>`. The internal
60 /// format is deliberately not part of any contract, decode tolerates only
61 /// what encode produces, and both can change without touching callers.
62 pub fn encode(&self) -> String {
63 let raw = format!("{}:{}", self.ts.timestamp_millis(), self.id);
64 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw)
65 }
66
67 // (see `Serialize` impl below: a Cursor serializes AS its opaque token, so a
68 // `Page.next` field emits the JSON string a client passes straight back.)
69
70 /// Decode a token produced by [`encode`](Self::encode). Returns `None` for
71 /// any malformed / tampered token, a bad cursor is treated as "start from
72 /// the beginning", never an error (a client should never be able to 500 the
73 /// server by editing a cursor).
74 pub fn decode(token: &str) -> Option<Self> {
75 let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
76 .decode(token)
77 .ok()?;
78 let raw = String::from_utf8(bytes).ok()?;
79 let (millis, id) = raw.split_once(':')?;
80 let ts = DateTime::<Utc>::from_timestamp_millis(millis.parse().ok()?)?;
81 Some(Self::new(ts, id.parse().ok()?))
82 }
83 }
84
85 impl Serialize for Cursor {
86 /// A cursor serializes as its opaque token string, so `Page.next` renders as
87 /// the exact value a client hands back as `?after=`.
88 fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
89 serializer.serialize_str(&self.encode())
90 }
91 }
92
93 /// Parameters for a keyset page request: where to start, and how many rows.
94 #[derive(Debug, Clone)]
95 pub struct PageParams {
96 pub after: Option<Cursor>,
97 pub limit: u32,
98 }
99
100 impl PageParams {
101 /// Build from a raw client cursor token + optional page size. A malformed
102 /// cursor silently starts from the beginning; the size is clamped to
103 /// `[1, MAX_PAGE_SIZE]`.
104 pub fn new(after_token: Option<&str>, limit: Option<u32>) -> Self {
105 Self {
106 after: after_token.and_then(Cursor::decode),
107 limit: limit.unwrap_or(DEFAULT_PAGE_SIZE).clamp(1, MAX_PAGE_SIZE),
108 }
109 }
110
111 /// The SQL `LIMIT` to bind: one more than the page size, so a returned extra
112 /// row signals "there is a next page" without a second COUNT query.
113 pub fn fetch_limit(&self) -> i64 {
114 i64::from(self.limit) + 1
115 }
116 }
117
118 /// One page of results plus the cursor to fetch the next page (`None` = last).
119 #[derive(Debug, Serialize)]
120 pub struct Page<T> {
121 pub items: Vec<T>,
122 pub next: Option<Cursor>,
123 }
124
125 impl<T> Page<T> {
126 /// Assemble a page from up to `limit + 1` fetched rows (ordered by the same
127 /// key the query used). If the extra row is present there is a next page:
128 /// drop it and take the next cursor from the last *kept* row. `cursor_of`
129 /// extracts the `(ts, id)` sort key from a row.
130 pub fn from_rows(
131 mut rows: Vec<T>,
132 params: &PageParams,
133 cursor_of: impl Fn(&T) -> Cursor,
134 ) -> Self {
135 let has_more = rows.len() as u32 > params.limit;
136 rows.truncate(params.limit as usize);
137 let next = if has_more {
138 rows.last().map(&cursor_of)
139 } else {
140 None
141 };
142 Page { items: rows, next }
143 }
144
145 /// The opaque next-page token for a client, if any.
146 pub fn next_token(&self) -> Option<String> {
147 self.next.map(|c| c.encode())
148 }
149 }
150
151 #[cfg(test)]
152 mod tests {
153 use super::*;
154
155 fn ts(millis: i64) -> DateTime<Utc> {
156 DateTime::<Utc>::from_timestamp_millis(millis).unwrap()
157 }
158
159 #[test]
160 fn cursor_round_trips_through_opaque_token() {
161 let c = Cursor::new(ts(1_700_000_000_123), Uuid::from_u128(42));
162 let token = c.encode();
163 // Opaque: not the raw value.
164 assert_ne!(token, "1700000000123:00000000-0000-0000-0000-00000000002a");
165 assert_eq!(Cursor::decode(&token), Some(c));
166 }
167
168 #[test]
169 fn malformed_cursor_decodes_to_none_not_error() {
170 assert_eq!(Cursor::decode("not-base64!!!"), None);
171 assert_eq!(Cursor::decode(""), None);
172 // Valid base64 but wrong shape.
173 let garbage = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode("no-colon");
174 assert_eq!(Cursor::decode(&garbage), None);
175 }
176
177 #[test]
178 fn params_clamp_and_default_page_size() {
179 assert_eq!(PageParams::new(None, None).limit, DEFAULT_PAGE_SIZE);
180 assert_eq!(PageParams::new(None, Some(0)).limit, 1);
181 assert_eq!(PageParams::new(None, Some(10_000)).limit, MAX_PAGE_SIZE);
182 assert_eq!(PageParams::new(None, Some(25)).fetch_limit(), 26);
183 }
184
185 #[test]
186 fn bad_cursor_token_starts_from_beginning() {
187 let p = PageParams::new(Some("tampered"), Some(20));
188 assert!(
189 p.after.is_none(),
190 "a bad cursor must not error, just restart"
191 );
192 }
193
194 #[test]
195 fn from_rows_detects_next_page_and_trims_extra() {
196 let params = PageParams::new(None, Some(3)); // fetch_limit = 4
197 // 4 rows returned for a page size of 3 => there IS a next page.
198 let rows: Vec<(i64, Uuid)> = (0..4).map(|i| (i, Uuid::from_u128(i as u128))).collect();
199 let page = Page::from_rows(rows, &params, |(m, id)| Cursor::new(ts(*m), *id));
200 assert_eq!(page.items.len(), 3, "the extra probe row is trimmed");
201 // Next cursor is the last KEPT row (index 2), not the dropped probe row.
202 assert_eq!(page.next, Some(Cursor::new(ts(2), Uuid::from_u128(2))));
203 }
204
205 #[test]
206 fn from_rows_last_page_has_no_next() {
207 let params = PageParams::new(None, Some(5)); // fetch_limit = 6
208 // Only 3 rows for a page size of 5 => last page.
209 let rows: Vec<(i64, Uuid)> = (0..3).map(|i| (i, Uuid::from_u128(i as u128))).collect();
210 let page = Page::from_rows(rows, &params, |(m, id)| Cursor::new(ts(*m), *id));
211 assert_eq!(page.items.len(), 3);
212 assert_eq!(page.next, None);
213 assert_eq!(page.next_token(), None);
214 }
215 }
216