use std::cell::RefCell;
use std::collections::HashMap;
use areev_core::error::{AreevError, Result};
use turso::Value;
pub(crate) fn db_err<E: std::fmt::Display>(e: E) -> AreevError {
AreevError::Storage(e.to_string())
}
pub(crate) struct Row(pub(crate) Vec<Value>);
impl Row {
pub(crate) fn i64(&self, i: usize) -> Option<i64> {
match self.0.get(i) {
Some(Value::Integer(x)) => Some(*x),
_ => None,
}
}
pub(crate) fn f64(&self, i: usize) -> Option<f64> {
match self.0.get(i) {
Some(Value::Real(r)) => Some(*r),
Some(Value::Integer(x)) => Some(*x as f64),
_ => None,
}
}
pub(crate) fn blob(&self, i: usize) -> Option<Vec<u8>> {
match self.0.get(i) {
Some(Value::Blob(b)) => Some(b.clone()),
_ => None,
}
}
pub(crate) fn text(&self, i: usize) -> Option<&str> {
match self.0.get(i) {
Some(Value::Text(t)) => Some(t.as_str()),
_ => None,
}
}
}
pub(crate) trait Db: Send {
fn execute(&self, sql: &str, params: Vec<Value>) -> Result<u64>;
fn query(&self, sql: &str, params: Vec<Value>) -> Result<Vec<Row>>;
fn execute_hot(&self, sql: &'static str, params: Vec<Value>) -> Result<u64>;
fn query_hot(&self, sql: &'static str, params: Vec<Value>) -> Result<Vec<Row>>;
fn begin(&self) -> Result<()>;
fn commit(&self) -> Result<()>;
fn rollback(&self) -> Result<()>;
fn prefers_batched_reads(&self) -> bool {
false
}
fn ensure_embeddings(&self, _dim: usize) -> Result<()> {
Ok(())
}
fn ensure_ann_index(&self, _m: usize, _ef_construction: usize, _ef_search: usize) -> Result<()> {
Err(AreevError::AnnIndexUnsupported(
"the embedded engine stores vectors as BLOBs and scans them exactly; \
scope the query by namespace or subject instead, or use the postgres backend"
.into(),
))
}
fn drop_ann_index(&self) -> Result<()> {
Ok(())
}
fn set_ann_ef_search(&self, _ef_search: usize) -> Result<()> {
Err(AreevError::AnnIndexUnsupported(
"no ANN index on this backend to tune".into(),
))
}
fn ann_ef_search(&self) -> Result<Option<usize>> {
Ok(None)
}
fn set_exact_vector_scan(&self, _on: bool) -> Result<()> {
Ok(())
}
fn ann_index_name(&self) -> Result<Option<String>> {
Ok(None)
}
fn reserve_write(
&self,
_n_seq: i64,
_n_op: i64,
_n_hlc: i64,
_wall_hlc: i64,
_hlc_floor: i64,
) -> Result<Option<WriteIds>> {
Ok(None)
}
fn intern_term(&self, _term: &str) -> Result<Option<i64>> {
Ok(None)
}
fn lookup_term(&self, _term: &str) -> Result<Option<i64>> {
Ok(None)
}
fn lookup_term_str(&self, _id: i64) -> Result<Option<String>> {
Ok(None)
}
fn scrub_term(&self, id: i64, placeholder: &str) -> Result<u64> {
self.execute(
"UPDATE terms SET term = ?2 WHERE id = ?1 \
AND NOT EXISTS (SELECT 1 FROM triples WHERE s=?1 OR p=?1 OR o=?1) \
AND NOT EXISTS (SELECT 1 FROM thread_idx WHERE session=?1) \
AND NOT EXISTS (SELECT 1 FROM run_idx WHERE run=?1) \
AND NOT EXISTS (SELECT 1 FROM grains WHERE ns=?1) \
AND term NOT LIKE ?3",
vec![crate::pi(id), crate::pt(placeholder), crate::pt("\u{1}erased:%")],
)
}
fn collection_stats(&self) -> Result<Option<(i64, i64)>> {
Ok(None)
}
fn for_update(&self) -> &'static str {
""
}
fn terms_with_prefix(&self, _prefix: &str) -> Result<Option<Vec<(i64, String)>>> {
Ok(None)
}
fn seeds_state_at_open(&self) -> bool {
true
}
fn may_have_legacy_fts_index(&self) -> bool {
true
}
}
pub(crate) struct WriteIds {
pub(crate) seq0: i64,
pub(crate) op0: i64,
pub(crate) hlc0: i64,
}
pub(crate) fn with_txn<T>(db: &dyn Db, f: impl FnOnce() -> Result<T>) -> Result<T> {
db.begin()?;
match f() {
Ok(v) => {
db.commit()?;
Ok(v)
}
Err(e) => {
let _ = db.rollback();
Err(e)
}
}
}
pub(crate) struct TursoDb {
rt: std::mem::ManuallyDrop<tokio::runtime::Runtime>,
_db: turso::Database,
conn: turso::Connection,
cache: RefCell<HashMap<&'static str, turso::Statement>>,
}
impl TursoDb {
pub(crate) fn open(path: &str, enc_key: Option<&[u8; 32]>) -> Result<Self> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(db_err)?;
let built = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
rt.block_on(async {
let mut b = turso::Builder::new_local(path).experimental_index_method(true);
if let Some(k) = enc_key {
let hexkey = zeroize::Zeroizing::new(hex32(k));
b = b.experimental_encryption(true).with_encryption(
turso::EncryptionOpts {
cipher: "aes256gcm".to_string(),
hexkey: (*hexkey).clone(),
},
);
}
let db = b.build().await.map_err(db_err)?;
let conn = db.connect().map_err(db_err)?;
Ok::<_, AreevError>((db, conn))
})
}));
let (db, conn) = match built {
Ok(r) => r?,
Err(_) => {
std::thread::spawn(move || drop(rt));
return Err(AreevError::AsyncContext(
"Areev is blocking and drives its own runtime, so it cannot be opened from \
an async task. Use AsyncAreev::open(path).await for the store, or \
areev_cal::AsyncFacade::open(path, opts).await for the governed facade — \
both open, call and tear down off the executor. To keep the blocking API, \
open it on a plain thread (std::thread::spawn) or inside \
tokio::task::spawn_blocking."
.to_string(),
));
}
};
Ok(Self {
rt: std::mem::ManuallyDrop::new(rt),
_db: db,
conn,
cache: RefCell::new(HashMap::new()),
})
}
fn hot(&self, sql: &'static str) -> Result<turso::Statement> {
if let Some(st) = self.cache.borrow().get(sql) {
return Ok(st.clone());
}
let st = self.rt.block_on(self.conn.prepare(sql)).map_err(db_err)?;
self.cache.borrow_mut().insert(sql, st.clone());
Ok(st)
}
}
async fn drain(rows: &mut turso::Rows) -> Result<Vec<Row>> {
let mut out = Vec::new();
while let Some(r) = rows.next().await.map_err(db_err)? {
let n = r.column_count();
let mut vals = Vec::with_capacity(n);
for i in 0..n {
vals.push(r.get_value(i).map_err(db_err)?);
}
out.push(Row(vals));
}
Ok(out)
}
impl Drop for TursoDb {
fn drop(&mut self) {
let rt = unsafe { std::mem::ManuallyDrop::take(&mut self.rt) };
if tokio::runtime::Handle::try_current().is_ok() {
std::thread::spawn(move || drop(rt));
} else {
drop(rt);
}
}
}
impl Db for TursoDb {
fn execute(&self, sql: &str, params: Vec<Value>) -> Result<u64> {
self.rt.block_on(self.conn.execute(sql, params)).map_err(db_err)
}
fn query(&self, sql: &str, params: Vec<Value>) -> Result<Vec<Row>> {
self.rt.block_on(async {
let mut rows = self.conn.query(sql, params).await.map_err(db_err)?;
drain(&mut rows).await
})
}
fn execute_hot(&self, sql: &'static str, params: Vec<Value>) -> Result<u64> {
let mut st = self.hot(sql)?;
self.rt.block_on(st.execute(params)).map_err(db_err)
}
fn query_hot(&self, sql: &'static str, params: Vec<Value>) -> Result<Vec<Row>> {
let mut st = self.hot(sql)?;
self.rt.block_on(async {
let mut rows = st.query(params).await.map_err(db_err)?;
drain(&mut rows).await
})
}
fn begin(&self) -> Result<()> {
self.execute("BEGIN", vec![]).map(|_| ())
}
fn commit(&self) -> Result<()> {
self.execute("COMMIT", vec![]).map(|_| ())
}
fn rollback(&self) -> Result<()> {
self.execute("ROLLBACK", vec![]).map(|_| ())
}
}
fn hex32(k: &[u8; 32]) -> String {
use std::fmt::Write;
let mut s = String::with_capacity(64);
for b in k {
let _ = write!(s, "{b:02x}");
}
s
}
#[cfg(test)]
mod tests {
use super::*;
fn mem() -> (TursoDb, tempfile::TempDir) {
let d = tempfile::TempDir::new().unwrap();
let db = TursoDb::open(d.path().join("t.db").to_str().unwrap(), None).unwrap();
(db, d)
}
#[test]
fn execute_query_roundtrip() {
let (db, _d) = mem();
db.execute("CREATE TABLE t(a INTEGER, b TEXT, c BLOB, d REAL)", vec![]).unwrap();
db.execute(
"INSERT INTO t(a,b,c,d) VALUES (?1,?2,?3,?4)",
vec![
Value::Integer(7),
Value::Text("x".into()),
Value::Blob(vec![1, 2]),
Value::Real(1.5),
],
)
.unwrap();
let rows = db.query("SELECT a, b, c, d FROM t", vec![]).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].i64(0), Some(7));
assert_eq!(rows[0].text(1), Some("x"));
assert_eq!(rows[0].blob(2), Some(vec![1, 2]));
assert_eq!(rows[0].f64(3), Some(1.5));
assert_eq!(rows[0].i64(1), None);
assert_eq!(rows[0].blob(0), None);
assert_eq!(rows[0].f64(0), Some(7.0)); assert_eq!(rows[0].i64(9), None); }
#[test]
fn hot_statements_are_cached_and_reentrant() {
let (db, _d) = mem();
db.execute("CREATE TABLE t(a INTEGER)", vec![]).unwrap();
for i in 0..5 {
db.execute_hot("INSERT INTO t(a) VALUES (?1)", vec![Value::Integer(i)]).unwrap();
}
for i in 0..5 {
let rows = db
.query_hot("SELECT a FROM t WHERE a = ?1", vec![Value::Integer(i)])
.unwrap();
assert_eq!(rows[0].i64(0), Some(i));
}
assert_eq!(db.cache.borrow().len(), 2);
}
#[test]
fn with_txn_commits_on_ok_and_rolls_back_on_err() {
let (db, _d) = mem();
db.execute("CREATE TABLE t(a INTEGER)", vec![]).unwrap();
with_txn(&db, || {
db.execute("INSERT INTO t(a) VALUES (1)", vec![])?;
Ok(())
})
.unwrap();
let r: Result<()> = with_txn(&db, || {
db.execute("INSERT INTO t(a) VALUES (2)", vec![])?;
Err(AreevError::Storage("boom".into()))
});
assert!(r.is_err());
let rows = db.query("SELECT COUNT(*) FROM t", vec![]).unwrap();
assert_eq!(rows[0].i64(0), Some(1), "rolled-back insert must not persist");
db.execute("INSERT INTO t(a) VALUES (3)", vec![]).unwrap();
}
}