use std::sync::Arc;
use serde::{Serialize, de::DeserializeOwned};
use eventuary_core::io::reader::{CheckpointKey, CheckpointScope, CheckpointStore};
use eventuary_core::io::{Cursor, CursorId};
use eventuary_core::{Error, Result};
use crate::database::SqliteConn;
use crate::relation::SqliteRelationName;
use crate::schema::{Migration, RelationReplacement};
const CHECKPOINT_STORE_0001_INIT_SQL: &str = r#"
CREATE TABLE IF NOT EXISTS {offsets} (
consumer_group_id TEXT NOT NULL,
stream_id TEXT NOT NULL DEFAULT 'default',
cursor_id TEXT NOT NULL,
cursor TEXT NOT NULL,
cursor_order BLOB NOT NULL DEFAULT x'',
PRIMARY KEY (consumer_group_id, stream_id, cursor_id)
);
"#;
const CHECKPOINT_STORE_MIGRATIONS: &[Migration] = &[Migration {
name: "0001_init",
sql: CHECKPOINT_STORE_0001_INIT_SQL,
}];
#[derive(Debug, Clone)]
pub struct SqliteCheckpointStoreConfig {
pub offsets_relation: SqliteRelationName,
}
impl Default for SqliteCheckpointStoreConfig {
fn default() -> Self {
Self {
offsets_relation: SqliteRelationName::new("consumer_offsets")
.expect("default offsets relation"),
}
}
}
pub struct SqliteCheckpointStore<C> {
conn: SqliteConn,
relation: Arc<String>,
_cursor: std::marker::PhantomData<fn() -> C>,
}
impl<C> Clone for SqliteCheckpointStore<C> {
fn clone(&self) -> Self {
Self {
conn: Arc::clone(&self.conn),
relation: Arc::clone(&self.relation),
_cursor: std::marker::PhantomData,
}
}
}
impl<C> SqliteCheckpointStore<C> {
pub fn new(conn: SqliteConn, config: SqliteCheckpointStoreConfig) -> Self {
Self {
conn,
relation: Arc::new(config.offsets_relation.render()),
_cursor: std::marker::PhantomData,
}
}
pub fn connect(conn: SqliteConn, config: SqliteCheckpointStoreConfig) -> Result<Self> {
Self::prepare_schema(&conn, &config)?;
Ok(Self::new(conn, config))
}
pub fn prepare_schema(conn: &SqliteConn, config: &SqliteCheckpointStoreConfig) -> Result<()> {
let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
crate::schema::apply_schema(
&guard,
CHECKPOINT_STORE_MIGRATIONS,
&[RelationReplacement {
token: "{offsets}",
relation: &config.offsets_relation,
}],
)
}
pub fn schema_sql(config: &SqliteCheckpointStoreConfig) -> String {
crate::schema::render_schema_sql(
CHECKPOINT_STORE_MIGRATIONS,
&[RelationReplacement {
token: "{offsets}",
relation: &config.offsets_relation,
}],
)
}
}
fn encode_cursor_id(cursor_id: &CursorId) -> String {
cursor_id.as_str().to_owned()
}
fn decode_cursor_id(value: &str) -> CursorId {
CursorId::new(value).unwrap_or_else(|_| CursorId::global())
}
fn encode_cursor<C: Serialize>(cursor: &C) -> Result<String> {
serde_json::to_string(cursor)
.map_err(|e| Error::Serialization(format!("checkpoint encode: {e}")))
}
fn decode_cursor<C: DeserializeOwned>(value: String) -> Result<C> {
serde_json::from_str(&value)
.map_err(|e| Error::Serialization(format!("checkpoint decode: {e}")))
}
impl<C> CheckpointStore<C> for SqliteCheckpointStore<C>
where
C: Cursor + Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
{
async fn load(&self, key: &CheckpointKey) -> Result<Option<C>> {
let conn = Arc::clone(&self.conn);
let relation = Arc::clone(&self.relation);
let cursor_id = encode_cursor_id(&key.cursor_id);
let group = key.scope.consumer_group_id.as_str().to_owned();
let stream = key.scope.stream_id.as_str().to_owned();
tokio::task::spawn_blocking(move || {
let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
let sql = format!(
"SELECT cursor FROM {relation} \
WHERE consumer_group_id = ?1 \
AND stream_id = ?2 \
AND cursor_id = ?3"
);
let row = guard
.query_row(&sql, rusqlite::params![group, stream, cursor_id], |r| {
r.get::<_, String>(0)
})
.map(Some)
.or_else(|e| match e {
rusqlite::Error::QueryReturnedNoRows => Ok(None),
other => Err(other),
})
.map_err(|e| Error::Store(e.to_string()))?;
match row {
Some(json) => Ok(Some(decode_cursor::<C>(json)?)),
None => Ok(None),
}
})
.await
.map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
}
async fn load_scope(&self, scope: &CheckpointScope) -> Result<Vec<(CursorId, C)>> {
let conn = Arc::clone(&self.conn);
let relation = Arc::clone(&self.relation);
let group = scope.consumer_group_id.as_str().to_owned();
let stream = scope.stream_id.as_str().to_owned();
tokio::task::spawn_blocking(move || {
let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
let sql = format!(
"SELECT cursor_id, cursor FROM {relation} \
WHERE consumer_group_id = ?1 AND stream_id = ?2"
);
let mut stmt = guard
.prepare(&sql)
.map_err(|e| Error::Store(e.to_string()))?;
let rows = stmt
.query_map(rusqlite::params![group, stream], |r| {
Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
})
.map_err(|e| Error::Store(e.to_string()))?;
let mut out = Vec::new();
for row in rows {
let (cursor_id_str, json) = row.map_err(|e| Error::Store(e.to_string()))?;
out.push((decode_cursor_id(&cursor_id_str), decode_cursor::<C>(json)?));
}
Ok(out)
})
.await
.map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
}
async fn commit(&self, key: &CheckpointKey, cursor: C) -> Result<()> {
let conn = Arc::clone(&self.conn);
let relation = Arc::clone(&self.relation);
let cursor_id = encode_cursor_id(&key.cursor_id);
let group = key.scope.consumer_group_id.as_str().to_owned();
let stream = key.scope.stream_id.as_str().to_owned();
let cursor_json = encode_cursor(&cursor)?;
let cursor_order = cursor.order_key();
tokio::task::spawn_blocking(move || {
let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
let sql = format!(
"INSERT INTO {relation} \
(consumer_group_id, stream_id, cursor_id, cursor, cursor_order) \
VALUES (?1, ?2, ?3, ?4, ?5) \
ON CONFLICT (consumer_group_id, stream_id, cursor_id) \
DO UPDATE SET cursor = excluded.cursor, \
cursor_order = excluded.cursor_order \
WHERE {relation}.cursor_order < excluded.cursor_order"
);
guard
.execute(
&sql,
rusqlite::params![
group,
stream,
cursor_id,
cursor_json,
cursor_order.as_bytes()
],
)
.map_err(|e| Error::Store(e.to_string()))?;
Ok(())
})
.await
.map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
}
}
#[cfg(test)]
mod tests {
use super::*;
use eventuary_core::Partition;
#[test]
fn schema_sql_contains_expected_table() {
let sql = SqliteCheckpointStore::<crate::reader::SqliteCursor>::schema_sql(
&SqliteCheckpointStoreConfig::default(),
);
assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"consumer_offsets\""));
}
use eventuary_core::io::cursor::CursorOrder;
use eventuary_core::io::reader::CheckpointScope;
use eventuary_core::io::{ConsumerGroupId, StreamId};
use std::num::NonZeroU32;
use crate::database::SqliteDatabase;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct SeqCursor(i64);
impl Cursor for SeqCursor {
fn order_key(&self) -> CursorOrder {
CursorOrder::from_i64(self.0)
}
}
fn checkpoint_key() -> CheckpointKey {
CheckpointKey::new(
CheckpointScope::new(
ConsumerGroupId::new("test-group").unwrap(),
StreamId::new("test-stream").unwrap(),
),
CursorId::global(),
)
}
fn make_store() -> SqliteCheckpointStore<SeqCursor> {
let db = SqliteDatabase::open_in_memory().unwrap();
let conn = db.conn();
SqliteCheckpointStore::<SeqCursor>::prepare_schema(
&conn,
&SqliteCheckpointStoreConfig::default(),
)
.unwrap();
SqliteCheckpointStore::new(conn, SqliteCheckpointStoreConfig::default())
}
#[tokio::test]
async fn commit_rejects_older_cursor() {
let store = make_store();
let key = checkpoint_key();
store.commit(&key, SeqCursor(100)).await.unwrap();
store.commit(&key, SeqCursor(50)).await.unwrap();
let loaded = store.load(&key).await.unwrap().unwrap();
assert_eq!(loaded.0, 100);
}
#[tokio::test]
async fn commit_advances_forward() {
let store = make_store();
let key = checkpoint_key();
store.commit(&key, SeqCursor(100)).await.unwrap();
store.commit(&key, SeqCursor(200)).await.unwrap();
let loaded = store.load(&key).await.unwrap().unwrap();
assert_eq!(loaded.0, 200);
}
#[tokio::test]
async fn commit_is_idempotent_for_equal_cursor() {
let store = make_store();
let key = checkpoint_key();
store.commit(&key, SeqCursor(100)).await.unwrap();
store.commit(&key, SeqCursor(100)).await.unwrap();
let loaded = store.load(&key).await.unwrap().unwrap();
assert_eq!(loaded.0, 100);
store.commit(&key, SeqCursor(150)).await.unwrap();
let loaded = store.load(&key).await.unwrap().unwrap();
assert_eq!(loaded.0, 150);
}
#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
struct WrappedCursor {
sequence: i64,
partition: Partition,
}
#[test]
fn encode_cursor_preserves_nested_json() {
let partition = Partition::new(2, NonZeroU32::new(4).unwrap()).unwrap();
let cursor = WrappedCursor {
sequence: 42,
partition,
};
let value = encode_cursor(&cursor).unwrap();
let decoded: WrappedCursor = decode_cursor(value).unwrap();
assert_eq!(decoded, cursor);
}
#[test]
fn cursor_id_global_encodes_as_plain_string() {
assert_eq!(encode_cursor_id(&CursorId::global()), "global");
assert_eq!(decode_cursor_id("global"), CursorId::global());
}
#[test]
fn cursor_id_named_roundtrips_unquoted() {
let id = CursorId::partition(
eventuary_core::partition::Partition::new(17, std::num::NonZeroU32::new(100).unwrap())
.unwrap(),
);
let encoded = encode_cursor_id(&id);
assert_eq!(encoded, "partition:100:17");
assert_eq!(decode_cursor_id(&encoded), id);
}
}