Skip to main content

eventuary_sqlite/
checkpoint_store.rs

1use 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}