1use std::sync::Arc;
2
3use serde::{Serialize, de::DeserializeOwned};
4
5use eventuary_core::io::reader::{CheckpointKey, CheckpointScope, CheckpointStore};
6use eventuary_core::io::{Cursor, CursorId};
7use eventuary_core::{Error, Result};
8
9use crate::database::SqliteConn;
10use crate::relation::SqliteRelationName;
11use crate::schema::{Migration, RelationReplacement};
12
13const CHECKPOINT_STORE_0001_INIT_SQL: &str = r#"
14CREATE TABLE IF NOT EXISTS {offsets} (
15 consumer_group_id TEXT NOT NULL,
16 stream_id TEXT NOT NULL DEFAULT 'default',
17 cursor_id TEXT NOT NULL,
18 cursor TEXT NOT NULL,
19 cursor_order BLOB NOT NULL DEFAULT x'',
20 PRIMARY KEY (consumer_group_id, stream_id, cursor_id)
21);
22"#;
23
24const CHECKPOINT_STORE_MIGRATIONS: &[Migration] = &[Migration {
25 name: "0001_init",
26 sql: CHECKPOINT_STORE_0001_INIT_SQL,
27}];
28
29#[derive(Debug, Clone)]
30pub struct SqliteCheckpointStoreConfig {
31 pub offsets_relation: SqliteRelationName,
32}
33
34impl Default for SqliteCheckpointStoreConfig {
35 fn default() -> Self {
36 Self {
37 offsets_relation: SqliteRelationName::new("consumer_offsets")
38 .expect("default offsets relation"),
39 }
40 }
41}
42
43pub struct SqliteCheckpointStore<C> {
44 conn: SqliteConn,
45 relation: Arc<String>,
46 _cursor: std::marker::PhantomData<fn() -> C>,
47}
48
49impl<C> Clone for SqliteCheckpointStore<C> {
50 fn clone(&self) -> Self {
51 Self {
52 conn: Arc::clone(&self.conn),
53 relation: Arc::clone(&self.relation),
54 _cursor: std::marker::PhantomData,
55 }
56 }
57}
58
59impl<C> SqliteCheckpointStore<C> {
60 pub fn new(conn: SqliteConn, config: SqliteCheckpointStoreConfig) -> Self {
61 Self {
62 conn,
63 relation: Arc::new(config.offsets_relation.render()),
64 _cursor: std::marker::PhantomData,
65 }
66 }
67
68 pub fn connect(conn: SqliteConn, config: SqliteCheckpointStoreConfig) -> Result<Self> {
69 Self::prepare_schema(&conn, &config)?;
70 Ok(Self::new(conn, config))
71 }
72
73 pub fn prepare_schema(conn: &SqliteConn, config: &SqliteCheckpointStoreConfig) -> Result<()> {
74 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
75 crate::schema::apply_schema(
76 &guard,
77 CHECKPOINT_STORE_MIGRATIONS,
78 &[RelationReplacement {
79 token: "{offsets}",
80 relation: &config.offsets_relation,
81 }],
82 )
83 }
84
85 pub fn schema_sql(config: &SqliteCheckpointStoreConfig) -> String {
86 crate::schema::render_schema_sql(
87 CHECKPOINT_STORE_MIGRATIONS,
88 &[RelationReplacement {
89 token: "{offsets}",
90 relation: &config.offsets_relation,
91 }],
92 )
93 }
94}
95
96fn encode_cursor_id(cursor_id: &CursorId) -> String {
97 cursor_id.as_str().to_owned()
98}
99
100fn decode_cursor_id(value: &str) -> CursorId {
101 CursorId::new(value).unwrap_or_else(|_| CursorId::global())
102}
103
104fn encode_cursor<C: Serialize>(cursor: &C) -> Result<String> {
105 serde_json::to_string(cursor)
106 .map_err(|e| Error::Serialization(format!("checkpoint encode: {e}")))
107}
108
109fn decode_cursor<C: DeserializeOwned>(value: String) -> Result<C> {
110 serde_json::from_str(&value)
111 .map_err(|e| Error::Serialization(format!("checkpoint decode: {e}")))
112}
113
114impl<C> CheckpointStore<C> for SqliteCheckpointStore<C>
115where
116 C: Cursor + Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
117{
118 async fn load(&self, key: &CheckpointKey) -> Result<Option<C>> {
119 let conn = Arc::clone(&self.conn);
120 let relation = Arc::clone(&self.relation);
121 let cursor_id = encode_cursor_id(&key.cursor_id);
122 let group = key.scope.consumer_group_id.as_str().to_owned();
123 let stream = key.scope.stream_id.as_str().to_owned();
124 tokio::task::spawn_blocking(move || {
125 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
126 let sql = format!(
127 "SELECT cursor FROM {relation} \
128 WHERE consumer_group_id = ?1 \
129 AND stream_id = ?2 \
130 AND cursor_id = ?3"
131 );
132 let row = guard
133 .query_row(&sql, rusqlite::params![group, stream, cursor_id], |r| {
134 r.get::<_, String>(0)
135 })
136 .map(Some)
137 .or_else(|e| match e {
138 rusqlite::Error::QueryReturnedNoRows => Ok(None),
139 other => Err(other),
140 })
141 .map_err(|e| Error::Store(e.to_string()))?;
142 match row {
143 Some(json) => Ok(Some(decode_cursor::<C>(json)?)),
144 None => Ok(None),
145 }
146 })
147 .await
148 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
149 }
150
151 async fn load_scope(&self, scope: &CheckpointScope) -> Result<Vec<(CursorId, C)>> {
152 let conn = Arc::clone(&self.conn);
153 let relation = Arc::clone(&self.relation);
154 let group = scope.consumer_group_id.as_str().to_owned();
155 let stream = scope.stream_id.as_str().to_owned();
156 tokio::task::spawn_blocking(move || {
157 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
158 let sql = format!(
159 "SELECT cursor_id, cursor FROM {relation} \
160 WHERE consumer_group_id = ?1 AND stream_id = ?2"
161 );
162 let mut stmt = guard
163 .prepare(&sql)
164 .map_err(|e| Error::Store(e.to_string()))?;
165 let rows = stmt
166 .query_map(rusqlite::params![group, stream], |r| {
167 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
168 })
169 .map_err(|e| Error::Store(e.to_string()))?;
170 let mut out = Vec::new();
171 for row in rows {
172 let (cursor_id_str, json) = row.map_err(|e| Error::Store(e.to_string()))?;
173 out.push((decode_cursor_id(&cursor_id_str), decode_cursor::<C>(json)?));
174 }
175 Ok(out)
176 })
177 .await
178 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
179 }
180
181 async fn commit(&self, key: &CheckpointKey, cursor: C) -> Result<()> {
182 let conn = Arc::clone(&self.conn);
183 let relation = Arc::clone(&self.relation);
184 let cursor_id = encode_cursor_id(&key.cursor_id);
185 let group = key.scope.consumer_group_id.as_str().to_owned();
186 let stream = key.scope.stream_id.as_str().to_owned();
187 let cursor_json = encode_cursor(&cursor)?;
188 let cursor_order = cursor.order_key();
189 tokio::task::spawn_blocking(move || {
190 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
191 let sql = format!(
192 "INSERT INTO {relation} \
193 (consumer_group_id, stream_id, cursor_id, cursor, cursor_order) \
194 VALUES (?1, ?2, ?3, ?4, ?5) \
195 ON CONFLICT (consumer_group_id, stream_id, cursor_id) \
196 DO UPDATE SET cursor = excluded.cursor, \
197 cursor_order = excluded.cursor_order \
198 WHERE {relation}.cursor_order < excluded.cursor_order"
199 );
200 guard
201 .execute(
202 &sql,
203 rusqlite::params![
204 group,
205 stream,
206 cursor_id,
207 cursor_json,
208 cursor_order.as_bytes()
209 ],
210 )
211 .map_err(|e| Error::Store(e.to_string()))?;
212 Ok(())
213 })
214 .await
215 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222 use eventuary_core::Partition;
223
224 #[test]
225 fn schema_sql_contains_expected_table() {
226 let sql = SqliteCheckpointStore::<crate::reader::SqliteCursor>::schema_sql(
227 &SqliteCheckpointStoreConfig::default(),
228 );
229 assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"consumer_offsets\""));
230 }
231 use eventuary_core::io::cursor::CursorOrder;
232 use eventuary_core::io::reader::CheckpointScope;
233 use eventuary_core::io::{ConsumerGroupId, StreamId};
234 use std::num::NonZeroU16;
235
236 use crate::database::SqliteDatabase;
237
238 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
239 struct SeqCursor(i64);
240
241 impl Cursor for SeqCursor {
242 fn order_key(&self) -> CursorOrder {
243 CursorOrder::from_i64(self.0)
244 }
245 }
246
247 fn checkpoint_key() -> CheckpointKey {
248 CheckpointKey::new(
249 CheckpointScope::new(
250 ConsumerGroupId::new("test-group").unwrap(),
251 StreamId::new("test-stream").unwrap(),
252 ),
253 CursorId::global(),
254 )
255 }
256
257 fn make_store() -> SqliteCheckpointStore<SeqCursor> {
258 let db = SqliteDatabase::open_in_memory().unwrap();
259 let conn = db.conn();
260 SqliteCheckpointStore::<SeqCursor>::prepare_schema(
261 &conn,
262 &SqliteCheckpointStoreConfig::default(),
263 )
264 .unwrap();
265 SqliteCheckpointStore::new(conn, SqliteCheckpointStoreConfig::default())
266 }
267
268 #[tokio::test]
269 async fn commit_rejects_older_cursor() {
270 let store = make_store();
271 let key = checkpoint_key();
272
273 store.commit(&key, SeqCursor(100)).await.unwrap();
274 store.commit(&key, SeqCursor(50)).await.unwrap();
275
276 let loaded = store.load(&key).await.unwrap().unwrap();
277 assert_eq!(loaded.0, 100);
278 }
279
280 #[tokio::test]
281 async fn commit_advances_forward() {
282 let store = make_store();
283 let key = checkpoint_key();
284
285 store.commit(&key, SeqCursor(100)).await.unwrap();
286 store.commit(&key, SeqCursor(200)).await.unwrap();
287
288 let loaded = store.load(&key).await.unwrap().unwrap();
289 assert_eq!(loaded.0, 200);
290 }
291
292 #[tokio::test]
293 async fn commit_is_idempotent_for_equal_cursor() {
294 let store = make_store();
295 let key = checkpoint_key();
296
297 store.commit(&key, SeqCursor(100)).await.unwrap();
298 store.commit(&key, SeqCursor(100)).await.unwrap();
299
300 let loaded = store.load(&key).await.unwrap().unwrap();
301 assert_eq!(loaded.0, 100);
302
303 store.commit(&key, SeqCursor(150)).await.unwrap();
304 let loaded = store.load(&key).await.unwrap().unwrap();
305 assert_eq!(loaded.0, 150);
306 }
307
308 #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
309 struct WrappedCursor {
310 sequence: i64,
311 partition: Partition,
312 }
313
314 #[test]
315 fn encode_cursor_preserves_nested_json() {
316 let partition = Partition::new(2, NonZeroU16::new(4).unwrap()).unwrap();
317 let cursor = WrappedCursor {
318 sequence: 42,
319 partition,
320 };
321
322 let value = encode_cursor(&cursor).unwrap();
323 let decoded: WrappedCursor = decode_cursor(value).unwrap();
324
325 assert_eq!(decoded, cursor);
326 }
327
328 #[test]
329 fn cursor_id_global_encodes_as_plain_string() {
330 assert_eq!(encode_cursor_id(&CursorId::global()), "global");
331 assert_eq!(decode_cursor_id("global"), CursorId::global());
332 }
333
334 #[test]
335 fn cursor_id_named_roundtrips_unquoted() {
336 let id = CursorId::partition(100, 17);
337 let encoded = encode_cursor_id(&id);
338 assert_eq!(encoded, "partition:100:17");
339 assert_eq!(decode_cursor_id(&encoded), id);
340 }
341}