eventuary_sqlite/
checkpoint_store.rs1use std::sync::Arc;
2
3use serde::{Serialize, de::DeserializeOwned};
4
5use eventuary_core::io::CursorId;
6use eventuary_core::io::reader::{CheckpointKey, CheckpointScope, CheckpointStore};
7use eventuary_core::{Error, Result};
8
9use crate::database::SqliteConn;
10use crate::relation::SqliteRelationName;
11
12#[derive(Debug, Clone)]
13pub struct SqliteCheckpointStoreConfig {
14 pub offsets_relation: SqliteRelationName,
15}
16
17impl Default for SqliteCheckpointStoreConfig {
18 fn default() -> Self {
19 Self {
20 offsets_relation: SqliteRelationName::new("consumer_offsets")
21 .expect("default offsets relation"),
22 }
23 }
24}
25
26pub struct SqliteCheckpointStore<C> {
27 conn: SqliteConn,
28 relation: Arc<String>,
29 _cursor: std::marker::PhantomData<fn() -> C>,
30}
31
32impl<C> Clone for SqliteCheckpointStore<C> {
33 fn clone(&self) -> Self {
34 Self {
35 conn: Arc::clone(&self.conn),
36 relation: Arc::clone(&self.relation),
37 _cursor: std::marker::PhantomData,
38 }
39 }
40}
41
42impl<C> SqliteCheckpointStore<C> {
43 pub fn new(conn: SqliteConn, config: SqliteCheckpointStoreConfig) -> Self {
44 Self {
45 conn,
46 relation: Arc::new(config.offsets_relation.render()),
47 _cursor: std::marker::PhantomData,
48 }
49 }
50}
51
52fn encode_cursor_id(cursor_id: &CursorId) -> String {
53 cursor_id.as_str().to_owned()
54}
55
56fn decode_cursor_id(value: &str) -> CursorId {
57 CursorId::new(value).unwrap_or_else(|_| CursorId::global())
58}
59
60fn encode_cursor<C: Serialize>(cursor: &C) -> Result<String> {
61 serde_json::to_string(cursor)
62 .map_err(|e| Error::Serialization(format!("checkpoint encode: {e}")))
63}
64
65fn decode_cursor<C: DeserializeOwned>(value: String) -> Result<C> {
66 serde_json::from_str(&value)
67 .map_err(|e| Error::Serialization(format!("checkpoint decode: {e}")))
68}
69
70impl<C> CheckpointStore<C> for SqliteCheckpointStore<C>
71where
72 C: Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
73{
74 async fn load(&self, key: &CheckpointKey) -> Result<Option<C>> {
75 let conn = Arc::clone(&self.conn);
76 let relation = Arc::clone(&self.relation);
77 let cursor_id = encode_cursor_id(&key.cursor_id);
78 let group = key.scope.consumer_group_id.as_str().to_owned();
79 let stream = key.scope.stream_id.as_str().to_owned();
80 tokio::task::spawn_blocking(move || {
81 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
82 let sql = format!(
83 "SELECT cursor FROM {relation} \
84 WHERE consumer_group_id = ?1 \
85 AND stream_id = ?2 \
86 AND cursor_id = ?3"
87 );
88 let row = guard
89 .query_row(&sql, rusqlite::params![group, stream, cursor_id], |r| {
90 r.get::<_, String>(0)
91 })
92 .map(Some)
93 .or_else(|e| match e {
94 rusqlite::Error::QueryReturnedNoRows => Ok(None),
95 other => Err(other),
96 })
97 .map_err(|e| Error::Store(e.to_string()))?;
98 match row {
99 Some(json) => Ok(Some(decode_cursor::<C>(json)?)),
100 None => Ok(None),
101 }
102 })
103 .await
104 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
105 }
106
107 async fn load_scope(&self, scope: &CheckpointScope) -> Result<Vec<(CursorId, C)>> {
108 let conn = Arc::clone(&self.conn);
109 let relation = Arc::clone(&self.relation);
110 let group = scope.consumer_group_id.as_str().to_owned();
111 let stream = scope.stream_id.as_str().to_owned();
112 tokio::task::spawn_blocking(move || {
113 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
114 let sql = format!(
115 "SELECT cursor_id, cursor FROM {relation} \
116 WHERE consumer_group_id = ?1 AND stream_id = ?2"
117 );
118 let mut stmt = guard
119 .prepare(&sql)
120 .map_err(|e| Error::Store(e.to_string()))?;
121 let rows = stmt
122 .query_map(rusqlite::params![group, stream], |r| {
123 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
124 })
125 .map_err(|e| Error::Store(e.to_string()))?;
126 let mut out = Vec::new();
127 for row in rows {
128 let (cursor_id_str, json) = row.map_err(|e| Error::Store(e.to_string()))?;
129 out.push((decode_cursor_id(&cursor_id_str), decode_cursor::<C>(json)?));
130 }
131 Ok(out)
132 })
133 .await
134 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
135 }
136
137 async fn commit(&self, key: &CheckpointKey, cursor: C) -> Result<()> {
138 let conn = Arc::clone(&self.conn);
139 let relation = Arc::clone(&self.relation);
140 let cursor_id = encode_cursor_id(&key.cursor_id);
141 let group = key.scope.consumer_group_id.as_str().to_owned();
142 let stream = key.scope.stream_id.as_str().to_owned();
143 let cursor_json = encode_cursor(&cursor)?;
144 tokio::task::spawn_blocking(move || {
145 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
146 let sql = format!(
147 "INSERT INTO {relation} (consumer_group_id, stream_id, cursor_id, cursor) \
148 VALUES (?1, ?2, ?3, ?4) \
149 ON CONFLICT (consumer_group_id, stream_id, cursor_id) \
150 DO UPDATE SET cursor = excluded.cursor"
151 );
152 guard
153 .execute(
154 &sql,
155 rusqlite::params![group, stream, cursor_id, cursor_json],
156 )
157 .map_err(|e| Error::Store(e.to_string()))?;
158 Ok(())
159 })
160 .await
161 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
162 }
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168 use eventuary_core::Partition;
169 use std::num::NonZeroU16;
170
171 #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
172 struct WrappedCursor {
173 sequence: i64,
174 partition: Partition,
175 }
176
177 #[test]
178 fn encode_cursor_preserves_nested_json() {
179 let partition = Partition::new(2, NonZeroU16::new(4).unwrap()).unwrap();
180 let cursor = WrappedCursor {
181 sequence: 42,
182 partition,
183 };
184
185 let value = encode_cursor(&cursor).unwrap();
186 let decoded: WrappedCursor = decode_cursor(value).unwrap();
187
188 assert_eq!(decoded, cursor);
189 }
190
191 #[test]
192 fn cursor_id_global_encodes_as_plain_string() {
193 assert_eq!(encode_cursor_id(&CursorId::global()), "global");
194 assert_eq!(decode_cursor_id("global"), CursorId::global());
195 }
196
197 #[test]
198 fn cursor_id_named_roundtrips_unquoted() {
199 let id = CursorId::partition(100, 17);
200 let encoded = encode_cursor_id(&id);
201 assert_eq!(encoded, "partition:100:17");
202 assert_eq!(decode_cursor_id(&encoded), id);
203 }
204}