use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use rusqlite::{Connection, OpenFlags, OptionalExtension};
use crate::error::AppError;
pub const LIBRARY: &[&str] = &[
include_str!("../migrations/0001_init.sql"),
include_str!("../migrations/0002_file_compressed.sql"),
];
pub const CONTROL: &[&str] = &[include_str!("../migrations/control/0001_init.sql")];
pub struct Db {
path: PathBuf,
writer: Mutex<Connection>,
readers: Mutex<Vec<Connection>>,
max_idle_readers: usize,
}
impl Db {
#[cfg(test)]
pub fn open(path: &Path) -> Result<Arc<Self>, AppError> {
Self::open_with(path, LIBRARY, 16)
}
pub fn open_with(
path: &Path,
migrations: &[&str],
max_idle_readers: usize,
) -> Result<Arc<Self>, AppError> {
if let Some(dir) = path.parent().filter(|d| !d.as_os_str().is_empty()) {
std::fs::create_dir_all(dir)
.map_err(|e| AppError::Internal(format!("create {}: {e}", dir.display())))?;
}
let mut conn = Connection::open(path)?;
conn.busy_timeout(std::time::Duration::from_secs(30))?;
conn.execute_batch(
"PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
PRAGMA foreign_keys = ON;",
)?;
migrate(&mut conn, migrations)?;
Ok(Arc::new(Db {
path: path.to_path_buf(),
writer: Mutex::new(conn),
readers: Mutex::new(Vec::new()),
max_idle_readers,
}))
}
fn reader(&self) -> Result<Connection, AppError> {
if let Some(conn) = self.readers.lock().unwrap_or_else(|e| e.into_inner()).pop() {
return Ok(conn);
}
let conn = Connection::open_with_flags(
&self.path,
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
)?;
conn.busy_timeout(std::time::Duration::from_secs(30))?;
Ok(conn)
}
fn give_back(&self, conn: Connection) {
let mut idle = self.readers.lock().unwrap_or_else(|e| e.into_inner());
if idle.len() < self.max_idle_readers {
idle.push(conn);
}
}
pub async fn read<T, F>(self: &Arc<Self>, f: F) -> Result<T, AppError>
where
T: Send + 'static,
F: FnOnce(&Connection) -> rusqlite::Result<T> + Send + 'static,
{
let db = self.clone();
tokio::task::spawn_blocking(move || {
let conn = db.reader()?;
let result = f(&conn);
db.give_back(conn);
Ok(result?)
})
.await
.map_err(|e| AppError::Internal(format!("blocking task: {e}")))?
}
pub async fn write<T, F>(self: &Arc<Self>, f: F) -> Result<T, AppError>
where
T: Send + 'static,
F: FnOnce(&mut Connection) -> rusqlite::Result<T> + Send + 'static,
{
let db = self.clone();
tokio::task::spawn_blocking(move || {
let mut conn = db.writer.lock().unwrap_or_else(|e| e.into_inner());
Ok(f(&mut conn)?)
})
.await
.map_err(|e| AppError::Internal(format!("blocking task: {e}")))?
}
pub async fn server_id(self: &Arc<Self>) -> Result<String, AppError> {
self.write(|c| {
if let Some(id) = get_meta(c, "server_id")? {
return Ok(id);
}
let id = crate::auth::random_hex(16);
set_meta(c, "server_id", &id)?;
Ok(id)
})
.await
}
}
pub fn get_meta(conn: &Connection, key: &str) -> rusqlite::Result<Option<String>> {
conn.query_row("SELECT value FROM meta WHERE key = ?1", [key], |r| r.get(0))
.optional()
}
pub fn set_meta(conn: &Connection, key: &str, value: &str) -> rusqlite::Result<()> {
conn.execute(
"INSERT INTO meta (key, value) VALUES (?1, ?2)
ON CONFLICT (key) DO UPDATE SET value = excluded.value",
[key, value],
)?;
Ok(())
}
pub(crate) fn migrate(conn: &mut Connection, migrations: &[&str]) -> rusqlite::Result<()> {
let applied: i64 = conn.query_row("PRAGMA user_version", [], |r| r.get(0))?;
if applied as usize > migrations.len() {
return Err(rusqlite::Error::SqliteFailure(
rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_MISMATCH),
Some(format!(
"the database is at schema {applied}, newer than this defter knows ({})",
migrations.len()
)),
));
}
for (index, sql) in migrations.iter().enumerate().skip(applied.max(0) as usize) {
let tx = conn.transaction()?;
tx.execute_batch(sql)?;
tx.pragma_update(None, "user_version", index as i64 + 1)?;
tx.commit()?;
tracing::info!(migration = index + 1, "applied database migration");
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn opens_migrates_and_keeps_the_server_id() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("defter.sqlite3");
let db = Db::open(&path).unwrap();
let id = db.server_id().await.unwrap();
assert_eq!(id.len(), 32);
drop(db);
let db = Db::open(&path).unwrap();
assert_eq!(db.server_id().await.unwrap(), id);
let version: i64 = db
.read(|c| c.query_row("SELECT version FROM library", [], |r| r.get(0)))
.await
.unwrap();
assert_eq!(version, 0);
}
}