garage/src/db/sqlite_adapter.rs

461 lines
10 KiB
Rust
Raw Normal View History

2022-06-02 17:58:47 +00:00
use core::ops::Bound;
2022-06-02 21:14:10 +00:00
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::ptr::NonNull;
use std::sync::{Arc, Mutex, MutexGuard, RwLock};
2022-06-02 17:58:47 +00:00
use log::trace;
2022-06-02 21:14:10 +00:00
use rusqlite::{params, Connection, Rows, Statement, Transaction};
2022-06-02 17:58:47 +00:00
2022-06-03 09:44:41 +00:00
use crate::{Db, Error, IDb, ITx, ITxFn, Result, TxError, TxFnResult, TxResult, Value, ValueIter};
2022-06-02 17:58:47 +00:00
pub use rusqlite;
2022-06-03 08:44:54 +00:00
// --- err
2022-06-02 17:58:47 +00:00
impl From<rusqlite::Error> for Error {
fn from(e: rusqlite::Error) -> Error {
Error(format!("{}", e).into())
}
}
impl<T> From<rusqlite::Error> for TxError<T> {
fn from(e: rusqlite::Error) -> TxError<T> {
TxError::Db(e.into())
}
}
2022-06-03 08:44:54 +00:00
// -- db
2022-06-02 17:58:47 +00:00
pub struct SqliteDb {
db: Mutex<Connection>,
trees: RwLock<Vec<String>>,
}
impl SqliteDb {
2022-06-03 09:14:24 +00:00
pub fn init(db: rusqlite::Connection) -> Db {
2022-06-02 17:58:47 +00:00
let s = Self {
db: Mutex::new(db),
trees: RwLock::new(Vec::new()),
};
Db(Arc::new(s))
}
fn get_tree(&self, i: usize) -> Result<String> {
self.trees
.read()
.unwrap()
.get(i)
.cloned()
2022-06-03 09:14:24 +00:00
.ok_or_else(|| Error("invalid tree id".into()))
2022-06-02 17:58:47 +00:00
}
}
impl IDb for SqliteDb {
fn open_tree(&self, name: &str) -> Result<usize> {
let name = format!("tree_{}", name.replace(':', "_COLON_"));
2022-06-03 09:44:41 +00:00
2022-06-02 17:58:47 +00:00
let mut trees = self.trees.write().unwrap();
2022-06-03 09:44:41 +00:00
if let Some(i) = trees.iter().position(|x| x == &name) {
2022-06-02 17:58:47 +00:00
Ok(i)
} else {
trace!("open tree {}: lock db", name);
let db = self.db.lock().unwrap();
trace!("create table {}", name);
db.execute(
2022-06-02 17:58:47 +00:00
&format!(
"CREATE TABLE IF NOT EXISTS {} (
k BLOB PRIMARY KEY,
v BLOB
)",
name
),
[],
)?;
trace!("table created: {}", name);
2022-06-02 17:58:47 +00:00
let i = trees.len();
trees.push(name.to_string());
Ok(i)
}
}
2022-06-03 09:44:41 +00:00
fn list_trees(&self) -> Result<Vec<String>> {
let mut trees = vec![];
trace!("list_trees: lock db");
2022-06-03 09:44:41 +00:00
let db = self.db.lock().unwrap();
trace!("list_trees: lock acquired");
2022-06-03 09:44:41 +00:00
let mut stmt = db.prepare(
"SELECT name FROM sqlite_schema WHERE type = 'table' AND name LIKE 'tree_%'",
)?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
let name = row.get::<_, String>(0)?;
let name = name.replace("_COLON_", ":");
trees.push(name);
}
Ok(trees)
}
2022-06-02 17:58:47 +00:00
// ----
2022-06-03 09:14:24 +00:00
fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value<'_>>> {
2022-06-02 17:58:47 +00:00
let tree = self.get_tree(tree)?;
2022-06-03 10:35:12 +00:00
trace!("get {}: lock db", tree);
2022-06-02 17:58:47 +00:00
let db = self.db.lock().unwrap();
2022-06-03 10:35:12 +00:00
trace!("get {}: lock acquired", tree);
2022-06-02 17:58:47 +00:00
let mut stmt = db.prepare(&format!("SELECT v FROM {} WHERE k = ?1", tree))?;
let mut res_iter = stmt.query([key])?;
match res_iter.next()? {
None => Ok(None),
Some(v) => Ok(Some(v.get::<_, Vec<u8>>(0)?.into())),
}
}
fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
let tree = self.get_tree(tree)?;
2022-06-03 10:35:12 +00:00
trace!("remove {}: lock db", tree);
2022-06-02 17:58:47 +00:00
let db = self.db.lock().unwrap();
2022-06-03 10:35:12 +00:00
trace!("remove {}: lock acquired", tree);
2022-06-02 17:58:47 +00:00
let res = db.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
Ok(res > 0)
}
fn len(&self, tree: usize) -> Result<usize> {
let tree = self.get_tree(tree)?;
2022-06-03 10:35:12 +00:00
trace!("len {}: lock db", tree);
2022-06-02 17:58:47 +00:00
let db = self.db.lock().unwrap();
2022-06-03 10:35:12 +00:00
trace!("len {}: lock acquired", tree);
2022-06-02 17:58:47 +00:00
let mut stmt = db.prepare(&format!("SELECT COUNT(*) FROM {}", tree))?;
let mut res_iter = stmt.query([])?;
match res_iter.next()? {
None => Ok(0),
2022-06-03 09:14:24 +00:00
Some(v) => Ok(v.get::<_, usize>(0)?),
2022-06-02 17:58:47 +00:00
}
}
fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
let tree = self.get_tree(tree)?;
2022-06-03 10:35:12 +00:00
trace!("insert {}: lock db", tree);
2022-06-02 17:58:47 +00:00
let db = self.db.lock().unwrap();
2022-06-03 10:35:12 +00:00
trace!("insert {}: lock acquired", tree);
2022-06-02 17:58:47 +00:00
db.execute(
&format!("INSERT OR REPLACE INTO {} (k, v) VALUES (?1, ?2)", tree),
params![key, value],
)?;
Ok(())
}
2022-06-03 09:14:24 +00:00
fn iter(&self, tree: usize) -> Result<ValueIter<'_>> {
2022-06-02 17:58:47 +00:00
let tree = self.get_tree(tree)?;
let sql = format!("SELECT k, v FROM {} ORDER BY k ASC", tree);
trace!("iter {}: lock db", tree);
let db = self.db.lock().unwrap();
trace!("iter {}: lock acquired", tree);
DbValueIterator::make(db, &sql, [])
2022-06-02 17:58:47 +00:00
}
2022-06-03 09:14:24 +00:00
fn iter_rev(&self, tree: usize) -> Result<ValueIter<'_>> {
2022-06-02 17:58:47 +00:00
let tree = self.get_tree(tree)?;
2022-06-02 21:14:10 +00:00
let sql = format!("SELECT k, v FROM {} ORDER BY k DESC", tree);
trace!("iter_rev {}: lock db", tree);
let db = self.db.lock().unwrap();
trace!("iter_rev {}: lock acquired", tree);
DbValueIterator::make(db, &sql, [])
2022-06-02 17:58:47 +00:00
}
2022-06-03 09:14:24 +00:00
fn range<'r>(
&self,
2022-06-02 17:58:47 +00:00
tree: usize,
low: Bound<&'r [u8]>,
high: Bound<&'r [u8]>,
2022-06-03 09:14:24 +00:00
) -> Result<ValueIter<'_>> {
2022-06-02 17:58:47 +00:00
let tree = self.get_tree(tree)?;
2022-06-02 21:46:28 +00:00
let (bounds_sql, params) = bounds_sql(low, high);
let sql = format!("SELECT k, v FROM {} {} ORDER BY k ASC", tree, bounds_sql);
let params = params
.iter()
.map(|x| x as &dyn rusqlite::ToSql)
.collect::<Vec<_>>();
trace!("range {}: lock db", tree);
let db = self.db.lock().unwrap();
trace!("range {}: lock acquired", tree);
DbValueIterator::make::<&[&dyn rusqlite::ToSql]>(db, &sql, params.as_ref())
2022-06-02 17:58:47 +00:00
}
2022-06-03 09:14:24 +00:00
fn range_rev<'r>(
&self,
2022-06-02 17:58:47 +00:00
tree: usize,
low: Bound<&'r [u8]>,
high: Bound<&'r [u8]>,
2022-06-03 09:14:24 +00:00
) -> Result<ValueIter<'_>> {
2022-06-02 17:58:47 +00:00
let tree = self.get_tree(tree)?;
2022-06-02 21:46:28 +00:00
let (bounds_sql, params) = bounds_sql(low, high);
let sql = format!("SELECT k, v FROM {} {} ORDER BY k DESC", tree, bounds_sql);
let params = params
.iter()
.map(|x| x as &dyn rusqlite::ToSql)
.collect::<Vec<_>>();
trace!("range_rev {}: lock db", tree);
let db = self.db.lock().unwrap();
trace!("range_rev {}: lock acquired", tree);
DbValueIterator::make::<&[&dyn rusqlite::ToSql]>(db, &sql, params.as_ref())
2022-06-02 17:58:47 +00:00
}
// ----
fn transaction(&self, f: &dyn ITxFn) -> TxResult<(), ()> {
let trees = self.trees.read().unwrap();
trace!("transaction: lock db");
2022-06-02 17:58:47 +00:00
let mut db = self.db.lock().unwrap();
trace!("transaction: lock acquired");
2022-06-02 17:58:47 +00:00
let tx = SqliteTx {
tx: db.transaction()?,
trees: trees.as_ref(),
};
let res = match f.try_on(&tx) {
2022-06-02 17:58:47 +00:00
TxFnResult::Ok => {
tx.tx.commit()?;
Ok(())
}
TxFnResult::Abort => {
tx.tx.rollback()?;
Err(TxError::Abort(()))
}
TxFnResult::DbErr => {
tx.tx.rollback()?;
Err(TxError::Db(Error(
"(this message will be discarded)".into(),
)))
}
};
trace!("transaction done");
res
2022-06-02 17:58:47 +00:00
}
}
// ----
struct SqliteTx<'a> {
tx: Transaction<'a>,
trees: &'a [String],
}
impl<'a> SqliteTx<'a> {
fn get_tree(&self, i: usize) -> Result<String> {
2022-06-03 09:14:24 +00:00
self.trees.get(i).cloned().ok_or_else(|| {
Error(
"invalid tree id (it might have been openned after the transaction started)".into(),
)
})
2022-06-02 17:58:47 +00:00
}
}
impl<'a> ITx<'a> for SqliteTx<'a> {
fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value<'a>>> {
let tree = self.get_tree(tree)?;
let mut stmt = self
.tx
.prepare(&format!("SELECT v FROM {} WHERE k = ?1", tree))?;
let mut res_iter = stmt.query([key])?;
match res_iter.next()? {
None => Ok(None),
Some(v) => Ok(Some(v.get::<_, Vec<u8>>(0)?.into())),
}
}
fn len(&self, tree: usize) -> Result<usize> {
let tree = self.get_tree(tree)?;
let mut stmt = self.tx.prepare(&format!("SELECT COUNT(*) FROM {}", tree))?;
let mut res_iter = stmt.query([])?;
match res_iter.next()? {
None => Ok(0),
2022-06-03 09:14:24 +00:00
Some(v) => Ok(v.get::<_, usize>(0)?),
2022-06-02 17:58:47 +00:00
}
}
fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> {
let tree = self.get_tree(tree)?;
self.tx.execute(
&format!("INSERT OR REPLACE INTO {} (k, v) VALUES (?1, ?2)", tree),
params![key, value],
)?;
Ok(())
}
fn remove(&self, tree: usize, key: &[u8]) -> Result<bool> {
let tree = self.get_tree(tree)?;
let res = self
.tx
.execute(&format!("DELETE FROM {} WHERE k = ?1", tree), params![key])?;
Ok(res > 0)
}
fn iter(&self, _tree: usize) -> Result<ValueIter<'a>> {
unimplemented!();
}
fn iter_rev(&self, _tree: usize) -> Result<ValueIter<'a>> {
unimplemented!();
}
fn range<'r>(
&self,
_tree: usize,
_low: Bound<&'r [u8]>,
_high: Bound<&'r [u8]>,
) -> Result<ValueIter<'a>> {
unimplemented!();
}
fn range_rev<'r>(
&self,
_tree: usize,
_low: Bound<&'r [u8]>,
_high: Bound<&'r [u8]>,
) -> Result<ValueIter<'a>> {
unimplemented!();
}
}
2022-06-02 21:14:10 +00:00
// ----
struct DbValueIterator<'a> {
2022-06-03 02:44:40 +00:00
db: MutexGuard<'a, Connection>,
2022-06-02 21:14:10 +00:00
stmt: Option<Statement<'a>>,
iter: Option<Rows<'a>>,
_pin: PhantomPinned,
}
impl<'a> DbValueIterator<'a> {
2022-06-03 09:14:24 +00:00
fn make<P: rusqlite::Params>(
2022-06-02 21:14:10 +00:00
db: MutexGuard<'a, Connection>,
sql: &str,
args: P,
) -> Result<ValueIter<'a>> {
let res = DbValueIterator {
2022-06-03 09:14:24 +00:00
db,
2022-06-02 21:14:10 +00:00
stmt: None,
iter: None,
_pin: PhantomPinned,
};
let mut boxed = Box::pin(res);
2022-06-03 11:35:02 +00:00
trace!("make iterator with sql: {}", sql);
2022-06-02 21:14:10 +00:00
unsafe {
let db = NonNull::from(&boxed.db);
2022-06-03 09:14:24 +00:00
let stmt = db.as_ref().prepare(sql)?;
2022-06-02 21:14:10 +00:00
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed);
Pin::get_unchecked_mut(mut_ref).stmt = Some(stmt);
let mut stmt = NonNull::from(&boxed.stmt);
let iter = stmt.as_mut().as_mut().unwrap().query(args)?;
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut boxed);
Pin::get_unchecked_mut(mut_ref).iter = Some(iter);
}
Ok(Box::new(DbValueIteratorPin(boxed)))
}
}
2022-06-03 02:44:40 +00:00
impl<'a> Drop for DbValueIterator<'a> {
fn drop(&mut self) {
trace!("drop iter");
2022-06-03 02:44:40 +00:00
drop(self.iter.take());
drop(self.stmt.take());
}
}
2022-06-02 21:14:10 +00:00
struct DbValueIteratorPin<'a>(Pin<Box<DbValueIterator<'a>>>);
impl<'a> Iterator for DbValueIteratorPin<'a> {
type Item = Result<(Value<'a>, Value<'a>)>;
fn next(&mut self) -> Option<Self::Item> {
let next = unsafe {
let mut_ref: Pin<&mut DbValueIterator<'a>> = Pin::as_mut(&mut self.0);
Pin::get_unchecked_mut(mut_ref).iter.as_mut()?.next()
};
let row = match next {
Err(e) => return Some(Err(e.into())),
2022-06-03 02:44:40 +00:00
Ok(None) => return None,
2022-06-02 21:14:10 +00:00
Ok(Some(r)) => r,
};
let k = match row.get::<_, Vec<u8>>(0) {
Err(e) => return Some(Err(e.into())),
Ok(x) => x,
};
let v = match row.get::<_, Vec<u8>>(1) {
Err(e) => return Some(Err(e.into())),
Ok(y) => y,
};
Some(Ok((k.into(), v.into())))
}
}
2022-06-02 21:46:28 +00:00
// ----
fn bounds_sql<'r>(low: Bound<&'r [u8]>, high: Bound<&'r [u8]>) -> (String, Vec<Vec<u8>>) {
let mut sql = String::new();
let mut params: Vec<Vec<u8>> = vec![];
match low {
Bound::Included(b) => {
sql.push_str(" WHERE k >= ?1");
params.push(b.to_vec());
}
Bound::Excluded(b) => {
sql.push_str(" WHERE k > ?1");
params.push(b.to_vec());
}
Bound::Unbounded => (),
};
match high {
Bound::Included(b) => {
if !params.is_empty() {
sql.push_str(" AND k <= ?2");
} else {
sql.push_str(" WHERE k <= ?1");
}
params.push(b.to_vec());
}
Bound::Excluded(b) => {
if !params.is_empty() {
sql.push_str(" AND k < ?2");
} else {
sql.push_str(" WHERE k < ?1");
}
params.push(b.to_vec());
}
Bound::Unbounded => (),
}
(sql, params)
}