defter 0.4.3

Self-hosted Zotero sync and file storage server, with metadata in SQLite and attachments in S3.
//! The SQLite database: one write connection behind a mutex, so writes are
//! serialized (the library version check and bump stay atomic), and a small
//! pool of read connections that WAL lets run alongside it. All access runs on
//! Tokio's blocking pool.

use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use rusqlite::{Connection, OpenFlags, OptionalExtension};

use crate::error::AppError;

/// Idle read connections kept between requests.
const MAX_IDLE_READERS: usize = 16;

/// Schema migrations, applied in order; `PRAGMA user_version` records how many
/// have run. Append new ones, never edit an applied one.
const MIGRATIONS: &[&str] = &[
    include_str!("../migrations/0001_init.sql"),
    include_str!("../migrations/0002_file_compressed.sql"),
];

pub struct Db {
    path: PathBuf,
    writer: Mutex<Connection>,
    readers: Mutex<Vec<Connection>>,
}

impl Db {
    /// Open (creating if needed) the database at `path` and migrate it.
    pub fn open(path: &Path) -> 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)?;
        Ok(Arc::new(Db {
            path: path.to_path_buf(),
            writer: Mutex::new(conn),
            readers: Mutex::new(Vec::new()),
        }))
    }

    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() < MAX_IDLE_READERS {
            idle.push(conn);
        }
    }

    /// Run `f` on a read connection.
    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}")))?
    }

    /// Run `f` on the write connection, one writer at a time.
    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}")))?
    }

    /// The installation's stable `Zotero-Server-ID`, minted on first start.
    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) -> rusqlite::Result<()> {
    let applied: i64 = conn.query_row("PRAGMA user_version", [], |r| r.get(0))?;
    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);
    }
}