eventuary_postgres/
checkpoint.rs

1use std::sync::Arc;
2
3use serde::{Serialize, de::DeserializeOwned};
4use sqlx::{PgPool, Row};
5
6use eventuary_core::io::reader::{CheckpointKey, CheckpointScope, CheckpointStore};
7use eventuary_core::io::{Cursor, CursorId};
8use eventuary_core::{Error, Result};
9
10use crate::relation::PgRelationName;
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            JSONB NOT NULL,
19    cursor_order      BYTEA NOT NULL DEFAULT ''::bytea,
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 PgCheckpointStoreConfig {
31    pub offsets_relation: PgRelationName,
32}
33
34impl Default for PgCheckpointStoreConfig {
35    fn default() -> Self {
36        Self {
37            offsets_relation: PgRelationName::new("consumer_offsets")
38                .expect("default offsets relation"),
39        }
40    }
41}
42
43pub struct PgCheckpointStore<C> {
44    pool: PgPool,
45    relation: Arc<String>,
46    _cursor: std::marker::PhantomData<fn() -> C>,
47}
48
49impl<C> Clone for PgCheckpointStore<C> {
50    fn clone(&self) -> Self {
51        Self {
52            pool: self.pool.clone(),
53            relation: Arc::clone(&self.relation),
54            _cursor: std::marker::PhantomData,
55        }
56    }
57}
58
59impl<C> PgCheckpointStore<C> {
60    pub fn new(pool: PgPool, config: PgCheckpointStoreConfig) -> Self {
61        Self {
62            pool,
63            relation: Arc::new(config.offsets_relation.render()),
64            _cursor: std::marker::PhantomData,
65        }
66    }
67
68    pub async fn connect(pool: PgPool, config: PgCheckpointStoreConfig) -> Result<Self> {
69        Self::prepare_schema(&pool, &config).await?;
70        Ok(Self::new(pool, config))
71    }
72
73    pub async fn prepare_schema(pool: &PgPool, config: &PgCheckpointStoreConfig) -> Result<()> {
74        crate::schema::apply_schema(
75            pool,
76            CHECKPOINT_STORE_MIGRATIONS,
77            &[RelationReplacement {
78                token: "{offsets}",
79                relation: &config.offsets_relation,
80            }],
81        )
82        .await
83    }
84
85    pub fn schema_sql(config: &PgCheckpointStoreConfig) -> 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) -> &str {
97    cursor_id.as_str()
98}
99
100fn decode_cursor_id(value: String) -> CursorId {
101    CursorId::new(value).unwrap_or_else(|_| CursorId::global())
102}
103
104fn encode_cursor<C: Serialize>(cursor: &C) -> Result<serde_json::Value> {
105    serde_json::to_value(cursor)
106        .map_err(|e| Error::Serialization(format!("checkpoint encode: {e}")))
107}
108
109fn decode_cursor<C: DeserializeOwned>(value: serde_json::Value) -> Result<C> {
110    serde_json::from_value(value)
111        .map_err(|e| Error::Serialization(format!("checkpoint decode: {e}")))
112}
113
114impl<C> CheckpointStore<C> for PgCheckpointStore<C>
115where
116    C: Cursor + Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
117{
118    async fn load(&self, key: &CheckpointKey) -> Result<Option<C>> {
119        let cursor_id = encode_cursor_id(&key.cursor_id);
120        let sql = format!(
121            "SELECT cursor FROM {relation} \
122             WHERE consumer_group_id = $1 AND stream_id = $2 AND cursor_id = $3",
123            relation = self.relation
124        );
125        let row = sqlx::query(&sql)
126            .bind(key.scope.consumer_group_id.as_str())
127            .bind(key.scope.stream_id.as_str())
128            .bind(cursor_id)
129            .fetch_optional(&self.pool)
130            .await
131            .map_err(|e| Error::Store(e.to_string()))?;
132        match row {
133            Some(r) => Ok(Some(decode_cursor::<C>(
134                r.get::<serde_json::Value, _>("cursor"),
135            )?)),
136            None => Ok(None),
137        }
138    }
139
140    async fn load_scope(&self, scope: &CheckpointScope) -> Result<Vec<(CursorId, C)>> {
141        let sql = format!(
142            "SELECT cursor_id, cursor FROM {relation} \
143             WHERE consumer_group_id = $1 AND stream_id = $2",
144            relation = self.relation
145        );
146        let rows = sqlx::query(&sql)
147            .bind(scope.consumer_group_id.as_str())
148            .bind(scope.stream_id.as_str())
149            .fetch_all(&self.pool)
150            .await
151            .map_err(|e| Error::Store(e.to_string()))?;
152        let mut out = Vec::with_capacity(rows.len());
153        for row in rows {
154            let cursor_id: String = row.get("cursor_id");
155            let cursor: serde_json::Value = row.get("cursor");
156            out.push((decode_cursor_id(cursor_id), decode_cursor::<C>(cursor)?));
157        }
158        Ok(out)
159    }
160
161    async fn commit(&self, key: &CheckpointKey, cursor: C) -> Result<()> {
162        let cursor_id = encode_cursor_id(&key.cursor_id);
163        let cursor_json = encode_cursor(&cursor)?;
164        let cursor_order = cursor.order_key();
165        let sql = format!(
166            "INSERT INTO {relation} \
167               (consumer_group_id, stream_id, cursor_id, cursor, cursor_order) \
168             VALUES ($1, $2, $3, $4, $5) \
169             ON CONFLICT (consumer_group_id, stream_id, cursor_id) DO UPDATE \
170             SET cursor = EXCLUDED.cursor, \
171                 cursor_order = EXCLUDED.cursor_order \
172             WHERE {relation}.cursor_order < EXCLUDED.cursor_order",
173            relation = self.relation
174        );
175        sqlx::query(&sql)
176            .bind(key.scope.consumer_group_id.as_str())
177            .bind(key.scope.stream_id.as_str())
178            .bind(cursor_id)
179            .bind(cursor_json)
180            .bind(cursor_order.as_bytes())
181            .execute(&self.pool)
182            .await
183            .map_err(|e| Error::Store(e.to_string()))?;
184        Ok(())
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use eventuary_core::Partition;
192    use std::num::NonZeroU32;
193
194    #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
195    struct WrappedCursor {
196        sequence: i64,
197        partition: Partition,
198    }
199
200    #[test]
201    fn encode_cursor_preserves_nested_json() {
202        let partition = Partition::new(2, NonZeroU32::new(4).unwrap()).unwrap();
203        let cursor = WrappedCursor {
204            sequence: 42,
205            partition,
206        };
207
208        let value = encode_cursor(&cursor).unwrap();
209        let decoded: WrappedCursor = decode_cursor(value).unwrap();
210
211        assert_eq!(decoded, cursor);
212    }
213
214    #[test]
215    fn cursor_id_global_encodes_as_plain_string() {
216        assert_eq!(encode_cursor_id(&CursorId::global()), "global");
217        assert_eq!(decode_cursor_id("global".to_owned()), CursorId::global());
218    }
219
220    #[test]
221    fn cursor_id_named_roundtrips_unquoted() {
222        let id = CursorId::partition(
223            eventuary_core::partition::Partition::new(17, std::num::NonZeroU32::new(100).unwrap())
224                .unwrap(),
225        );
226        let encoded = encode_cursor_id(&id).to_owned();
227        assert_eq!(encoded, "partition:100:17");
228        assert_eq!(decode_cursor_id(encoded), id);
229    }
230
231    #[test]
232    fn schema_sql_contains_expected_table() {
233        let sql = PgCheckpointStore::<crate::reader::PgCursor>::schema_sql(
234            &PgCheckpointStoreConfig::default(),
235        );
236        assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"consumer_offsets\""));
237    }
238}
eventuary_postgres/checkpoint.rs

eventuary_postgres/
checkpoint.rs
