use std::path::{Path, PathBuf};
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::Mutex;
use turso::{Builder, Connection};
use tinyflow_api::error::StreamResult;
use tinyflow_api::state::{StateStore, state_db_err};
pub struct SqliteJobStateStore {
path: PathBuf,
conn: Mutex<Connection>,
}
impl SqliteJobStateStore {
pub async fn open(path: &Path) -> StreamResult<Arc<dyn StateStore>> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| state_db_err(format!("create state db parent dir: {}", e)))?;
}
let path_str = path
.to_str()
.ok_or_else(|| state_db_err("non-UTF-8 path for turso db".to_string()))?;
let db = Builder::new_local(path_str)
.build()
.await
.map_err(|e| state_db_err(format!("open turso db: {}", e)))?;
let conn = db
.connect()
.map_err(|e| state_db_err(format!("connect turso: {}", e)))?;
let mut stmt = conn
.prepare("PRAGMA journal_mode=WAL")
.await
.map_err(|e| state_db_err(format!("turso pragma prepare: {}", e)))?;
stmt.query(())
.await
.map_err(|e| state_db_err(format!("turso pragma query: {}", e)))?;
let mut stmt = conn
.prepare("PRAGMA busy_timeout=5000")
.await
.map_err(|e| state_db_err(format!("turso busy_timeout prepare: {}", e)))?;
stmt.query(())
.await
.map_err(|e| state_db_err(format!("turso busy_timeout query: {}", e)))?;
conn.execute(
"CREATE TABLE IF NOT EXISTS kv (
namespace TEXT NOT NULL,
key BLOB NOT NULL,
value BLOB NOT NULL,
PRIMARY KEY (namespace, key)
)",
(),
)
.await
.map_err(|e| state_db_err(format!("turso init schema: {}", e)))?;
let store: Arc<dyn StateStore> = Arc::new(Self {
path: path.to_path_buf(),
conn: Mutex::new(conn),
});
Ok(store)
}
}
fn prefix_upper_bound(prefix: &[u8]) -> Option<Vec<u8>> {
if prefix.is_empty() {
return None;
}
let mut upper = prefix.to_vec();
while let Some(last) = upper.last_mut() {
if *last < 255 {
*last += 1;
return Some(upper);
}
upper.pop();
}
None
}
fn hex_prefix_like(prefix: &[u8]) -> String {
prefix.iter().map(|b| format!("{:02x}", b)).collect()
}
#[async_trait]
impl StateStore for SqliteJobStateStore {
fn persistent_db_path(&self) -> Option<&Path> {
Some(self.path.as_path())
}
async fn put(&self, namespace: &str, key: &[u8], value: &[u8]) -> StreamResult<()> {
let conn = self.conn.lock().await;
conn.execute(
"INSERT OR REPLACE INTO kv (namespace, key, value) VALUES (?1, ?2, ?3)",
(namespace, key, value),
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
Ok(())
}
async fn get(&self, namespace: &str, key: &[u8]) -> StreamResult<Option<Vec<u8>>> {
let conn = self.conn.lock().await;
let mut stmt = conn
.prepare("SELECT value FROM kv WHERE namespace = ?1 AND key = ?2")
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
let mut rows = stmt
.query((namespace, key))
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
if let Some(row) = rows
.next()
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?
{
let blob = row
.get_value(0)
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
let v = blob
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
return Ok(Some(v));
}
Ok(None)
}
async fn delete(&self, namespace: &str, key: &[u8]) -> StreamResult<()> {
let conn = self.conn.lock().await;
conn.execute(
"DELETE FROM kv WHERE namespace = ?1 AND key = ?2",
(namespace, key),
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
Ok(())
}
async fn delete_prefix(&self, namespace: &str, key_prefix: &[u8]) -> StreamResult<()> {
let conn = self.conn.lock().await;
if let Some(upper) = prefix_upper_bound(key_prefix) {
conn.execute(
"DELETE FROM kv WHERE namespace = ?1 AND key >= ?2 AND key < ?3",
(namespace, key_prefix, upper.as_slice()),
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
} else if key_prefix.is_empty() {
conn.execute("DELETE FROM kv WHERE namespace = ?1", [namespace])
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
} else {
let pattern = format!("{}%", hex_prefix_like(key_prefix));
conn.execute(
"DELETE FROM kv WHERE namespace = ?1 AND hex(key) LIKE ?2",
(namespace, pattern.as_str()),
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
}
Ok(())
}
async fn list_prefix(
&self,
namespace: &str,
key_prefix: &[u8],
) -> StreamResult<Vec<(Vec<u8>, Vec<u8>)>> {
let conn = self.conn.lock().await;
let mut out = Vec::new();
if let Some(upper) = prefix_upper_bound(key_prefix) {
let mut stmt = conn
.prepare(
"SELECT key, value FROM kv WHERE namespace = ?1 AND key >= ?2 AND key < ?3 ORDER BY key",
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
let mut rows = stmt
.query((namespace, key_prefix, upper.as_slice()))
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
while let Some(row) = rows
.next()
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?
{
let k = row
.get_value(0)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
let v = row
.get_value(1)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
out.push((k, v));
}
} else if key_prefix.is_empty() {
let mut stmt = conn
.prepare("SELECT key, value FROM kv WHERE namespace = ?1 ORDER BY key")
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
let mut rows = stmt
.query([namespace])
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
while let Some(row) = rows
.next()
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?
{
let k = row
.get_value(0)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
let v = row
.get_value(1)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
out.push((k, v));
}
} else {
let pattern = format!("{}%", hex_prefix_like(key_prefix));
let mut stmt = conn
.prepare(
"SELECT key, value FROM kv WHERE namespace = ?1 AND hex(key) LIKE ?2 ORDER BY key",
)
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
let mut rows = stmt
.query((namespace, pattern.as_str()))
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?;
while let Some(row) = rows
.next()
.await
.map_err(|e| state_db_err(format!("turso: {}", e)))?
{
let k = row
.get_value(0)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
let v = row
.get_value(1)
.map_err(|e| state_db_err(format!("turso: {}", e)))?
.as_blob()
.ok_or_else(|| state_db_err("null value in NOT NULL column".to_string()))?
.to_vec();
out.push((k, v));
}
}
Ok(out)
}
}