jellytau/src-tauri/src/storage/db_service.rs

//! Database service: the single owner of the SQLite database.
//!
//! Every query in the app goes through [`RusqliteService`], which owns the
//! connections and hands work to them — callers never touch a `Connection`.
//!
//! - **Writes** (`execute`, `insert`, `transaction`, …) are sent as jobs to one
//!   dedicated writer thread that owns the read-write connection. SQLite allows
//!   one writer at a time anyway; owning it on one thread makes that explicit,
//!   keeps connection-wide state (pragmas) out of reach of concurrent callers,
//!   and parks no tokio blocking threads on a mutex while writes queue up.
//! - **Reads** (`query_*`) run on a small pool of read-only connections. The
//!   database is in WAL mode, so readers see the last committed state and never
//!   wait for the writer — a large catalog-cache transaction no longer stalls
//!   library pages, thumbnail lookups or settings reads.
//!
//! A service built with [`RusqliteService::new`] has no reader pool (in-memory
//! databases cannot be shared between connections) and routes reads through
//! the writer, which is the old single-connection behaviour tests rely on.
//!
//! See `docs/architecture/08-database-design.md` → "Connection ownership".

use crate::utils::lock::MutexSafe;
use async_trait::async_trait;
use log::{debug, error};
use rusqlite::{params_from_iter, Connection, Result as SqliteResult, Row};
use std::panic::AssertUnwindSafe;
use std::sync::mpsc;
use std::sync::{Arc, Condvar, Mutex};

/// Database query result type
pub type DbResult<T> = Result<T, String>;

/// Represents a database query that can be executed
#[derive(Clone)]
pub struct Query {
    pub sql: String,
    pub params: Vec<QueryParam>,
}

/// Query parameter types supported by the database
#[derive(Clone, Debug)]
pub enum QueryParam {
    String(String),
    Int(i32),
    Int64(i64),
    Float(f64),
    #[allow(dead_code)]
    Bool(bool),
    Null,
}

impl Query {
    pub fn new(sql: impl Into<String>) -> Self {
        Self {
            sql: sql.into(),
            params: Vec::new(),
        }
    }

    pub fn with_params(sql: impl Into<String>, params: Vec<QueryParam>) -> Self {
        Self {
            sql: sql.into(),
            params,
        }
    }
}

/// Database service trait - abstraction over database operations
#[async_trait]
pub trait DatabaseService: Send + Sync {
    /// Execute a query that doesn't return results (INSERT, UPDATE, DELETE)
    async fn execute(&self, query: Query) -> DbResult<usize>;

    /// Execute a batch of SQL statements (for migrations)
    #[allow(dead_code)]
    async fn execute_batch(&self, sql: &str) -> DbResult<()>;

    /// Query a single row
    async fn query_one<T, F>(&self, query: Query, mapper: F) -> DbResult<T>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;

    /// Query a single optional row
    async fn query_optional<T, F>(&self, query: Query, mapper: F) -> DbResult<Option<T>>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;

    /// Query multiple rows
    async fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;

    /// Run a transaction with multiple operations
    async fn transaction<F, T>(&self, f: F) -> DbResult<T>
    where
        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
        T: Send + 'static;

    /// Run a transaction with foreign-key enforcement switched off for its
    /// duration only.
    ///
    /// `PRAGMA foreign_keys` is per connection and is a no-op inside a
    /// transaction, so it has to be flipped around the `BEGIN`/`COMMIT` — and
    /// all of that must happen as one job on the writer, or any other write
    /// that got in between would run unchecked too.
    async fn transaction_without_foreign_keys<F, T>(&self, f: F) -> DbResult<T>
    where
        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
        T: Send + 'static;

    /// Execute an INSERT and return the rowid of the row it inserted.
    ///
    /// The rowid is read in the same job as the insert. Reading it with a
    /// second call would race every other write, returning someone else's id.
    async fn insert(&self, query: Query) -> DbResult<i64>;

    /// Queue a write without waiting for it — for best-effort bookkeeping (an
    /// LRU access time) that must not hold up the caller. It still runs in
    /// order with every other write; failures are only logged.
    fn execute_detached(&self, query: Query);
}

/// Transaction handle for batching multiple operations
pub struct Transaction<'a> {
    conn: &'a Connection,
}

impl<'a> Transaction<'a> {
    pub fn new(conn: &'a Connection) -> Self {
        Self { conn }
    }

    pub fn execute(&mut self, query: Query) -> DbResult<usize> {
        execute_query(self.conn, query)
    }

    pub fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    where
        F: Fn(&Row) -> SqliteResult<T>,
    {
        query_many(self.conn, query, mapper)
    }
}

type Job = Box<dyn FnOnce(&Connection) + Send>;

/// The thread that owns the read-write connection. Jobs run one at a time, in
/// the order they were sent; the thread exits when the last service handle
/// (and so the last sender) is dropped.
struct Writer {
    jobs: mpsc::Sender<Job>,
}

impl Writer {
    fn spawn(conn: Arc<Mutex<Connection>>) -> Self {
        let (jobs, queue) = mpsc::channel::<Job>();
        std::thread::Builder::new()
            .name("db-writer".into())
            .spawn(move || {
                for job in queue {
                    // The connection stays behind a mutex only so migrations and
                    // tests can reach it; in the app this thread is its sole user.
                    // `lock_safe` so a poisoned lock is recovered, not fatal.
                    let conn = conn.lock_safe();
                    // A panicking row mapper must not take the owner down with
                    // it: the job's reply channel drops, its caller gets an
                    // error, and the next job runs normally. The guard lives
                    // outside the unwind, so the mutex is not poisoned either.
                    if std::panic::catch_unwind(AssertUnwindSafe(|| job(&conn))).is_err() {
                        error!("[db] a database job panicked; the writer carries on");
                        // Undo whatever connection state the job was midway
                        // through: an open transaction would make the next
                        // job's BEGIN fail, and a job that switched foreign
                        // keys off would leave them off for everyone.
                        if !conn.is_autocommit() {
                            let _ = conn.execute_batch("ROLLBACK");
                        }
                        let _ = conn.execute_batch("PRAGMA foreign_keys = ON");
                    }
                }
            })
            .expect("failed to spawn the database writer thread");
        Self { jobs }
    }

    async fn run<T, F>(&self, f: F) -> DbResult<T>
    where
        T: Send + 'static,
        F: FnOnce(&Connection) -> DbResult<T> + Send + 'static,
    {
        let (reply, result) = tokio::sync::oneshot::channel();
        self.jobs
            .send(Box::new(move |conn| {
                let _ = reply.send(f(conn));
            }))
            .map_err(|_| "database writer has stopped".to_string())?;
        result
            .await
            .map_err(|_| "database job panicked".to_string())?
    }

    fn run_detached(&self, f: impl FnOnce(&Connection) + Send + 'static) {
        if self.jobs.send(Box::new(f)).is_err() {
            debug!("[db] writer stopped; dropped a detached write");
        }
    }
}

/// Read-only connections, checked out one per query. WAL gives each a
/// snapshot of the last commit, so they never wait for the writer.
struct ReaderPool {
    idle: Mutex<Vec<Connection>>,
    returned: Condvar,
}

impl ReaderPool {
    /// Blocking: waits for a free connection. Call from `spawn_blocking`.
    fn run<T>(&self, f: impl FnOnce(&Connection) -> T) -> T {
        let conn = {
            let mut idle = self.idle.lock_safe();
            loop {
                if let Some(conn) = idle.pop() {
                    break conn;
                }
                idle = self
                    .returned
                    .wait(idle)
                    .unwrap_or_else(|poisoned| poisoned.into_inner());
            }
        };
        // Returned on drop, so a panicking mapper does not leak the connection.
        let checkout = Checkout {
            pool: self,
            conn: Some(conn),
        };
        f(checkout.conn.as_ref().expect("checked-out connection"))
    }
}

struct Checkout<'a> {
    pool: &'a ReaderPool,
    conn: Option<Connection>,
}

impl Drop for Checkout<'_> {
    fn drop(&mut self) {
        if let Some(conn) = self.conn.take() {
            self.pool.idle.lock_safe().push(conn);
            self.pool.returned.notify_one();
        }
    }
}

/// Rusqlite-based database service: a cheap, cloneable handle to the writer
/// thread and reader pool. See the module docs.
#[derive(Clone)]
pub struct RusqliteService {
    writer: Arc<Writer>,
    readers: Option<Arc<ReaderPool>>,
}

impl RusqliteService {
    /// A service over a single connection: writes *and* reads go through the
    /// writer thread. Used for in-memory databases, which cannot be shared
    /// between connections.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn new(conn: Arc<Mutex<Connection>>) -> Self {
        Self {
            writer: Arc::new(Writer::spawn(conn)),
            readers: None,
        }
    }

    /// A service whose reads run on `readers` — read-only connections to the
    /// same (file-backed, WAL-mode) database — alongside the writer.
    pub fn with_readers(conn: Arc<Mutex<Connection>>, readers: Vec<Connection>) -> Self {
        let readers = (!readers.is_empty()).then(|| {
            Arc::new(ReaderPool {
                idle: Mutex::new(readers),
                returned: Condvar::new(),
            })
        });
        Self {
            writer: Arc::new(Writer::spawn(conn)),
            readers,
        }
    }

    async fn read<T, F>(&self, f: F) -> DbResult<T>
    where
        T: Send + 'static,
        F: FnOnce(&Connection) -> DbResult<T> + Send + 'static,
    {
        match &self.readers {
            Some(pool) => {
                let pool = Arc::clone(pool);
                tokio::task::spawn_blocking(move || pool.run(f))
                    .await
                    .map_err(|e| format!("Task join error: {}", e))?
            }
            None => self.writer.run(f).await,
        }
    }
}

#[async_trait]
impl DatabaseService for RusqliteService {
    async fn execute(&self, query: Query) -> DbResult<usize> {
        self.writer
            .run(move |conn| execute_query(conn, query))
            .await
    }

    async fn execute_batch(&self, sql: &str) -> DbResult<()> {
        let sql = sql.to_string();
        self.writer
            .run(move |conn| {
                conn.execute_batch(&sql)
                    .map_err(|e| format!("Execute batch failed: {}", e))
            })
            .await
    }

    async fn query_one<T, F>(&self, query: Query, mapper: F) -> DbResult<T>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    {
        self.read(move |conn| query_one(conn, query, mapper)).await
    }

    async fn query_optional<T, F>(&self, query: Query, mapper: F) -> DbResult<Option<T>>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    {
        self.read(move |conn| query_optional(conn, query, mapper))
            .await
    }

    async fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    {
        self.read(move |conn| query_many(conn, query, mapper)).await
    }

    async fn transaction<F, T>(&self, f: F) -> DbResult<T>
    where
        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
        T: Send + 'static,
    {
        self.writer.run(move |conn| run_transaction(conn, f)).await
    }

    async fn transaction_without_foreign_keys<F, T>(&self, f: F) -> DbResult<T>
    where
        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
        T: Send + 'static,
    {
        self.writer
            .run(move |conn| {
                conn.execute_batch("PRAGMA foreign_keys = OFF")
                    .map_err(|e| format!("Failed to disable foreign keys: {}", e))?;
                let result = run_transaction(conn, f);
                // Always restored, whatever the transaction did.
                if let Err(e) = conn.execute_batch("PRAGMA foreign_keys = ON") {
                    error!("[db] failed to re-enable foreign keys: {}", e);
                }
                result
            })
            .await
    }

    async fn insert(&self, query: Query) -> DbResult<i64> {
        self.writer
            .run(move |conn| {
                execute_query(conn, query)?;
                Ok(conn.last_insert_rowid())
            })
            .await
    }

    fn execute_detached(&self, query: Query) {
        self.writer.run_detached(move |conn| {
            if let Err(e) = execute_query(conn, query) {
                debug!("[db] detached write failed: {}", e);
            }
        });
    }
}

fn run_transaction<F, T>(conn: &Connection, f: F) -> DbResult<T>
where
    F: FnOnce(&mut Transaction) -> DbResult<T>,
{
    conn.execute("BEGIN TRANSACTION", [])
        .map_err(|e| format!("Failed to begin transaction: {}", e))?;

    let mut transaction = Transaction::new(conn);
    match f(&mut transaction) {
        Ok(value) => {
            conn.execute("COMMIT", [])
                .map_err(|e| format!("Failed to commit transaction: {}", e))?;
            Ok(value)
        }
        Err(e) => {
            conn.execute("ROLLBACK", [])
                .map_err(|e| format!("Failed to rollback transaction: {}", e))?;
            Err(e)
        }
    }
}

// Helper functions for executing queries synchronously

fn execute_query(conn: &Connection, query: Query) -> DbResult<usize> {
    let params = convert_params(&query.params);
    conn.execute(&query.sql, params_from_iter(params.iter()))
        .map_err(|e| format!("Execute failed: {}", e))
}

fn query_one<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<T>
where
    F: Fn(&Row) -> SqliteResult<T>,
{
    let params = convert_params(&query.params);
    conn.query_row(&query.sql, params_from_iter(params.iter()), mapper)
        .map_err(|e| format!("Query one failed: {}", e))
}

fn query_optional<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<Option<T>>
where
    F: Fn(&Row) -> SqliteResult<T>,
{
    match query_one(conn, query, mapper) {
        Ok(value) => Ok(Some(value)),
        Err(e) if e.contains("Query returned no rows") || e.contains("QueryReturnedNoRows") => {
            Ok(None)
        }
        Err(e) => Err(e),
    }
}

fn query_many<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<Vec<T>>
where
    F: Fn(&Row) -> SqliteResult<T>,
{
    let params = convert_params(&query.params);
    let mut stmt = conn
        .prepare(&query.sql)
        .map_err(|e| format!("Prepare failed: {}", e))?;

    let rows = stmt
        .query_map(params_from_iter(params.iter()), mapper)
        .map_err(|e| format!("Query map failed: {}", e))?;

    rows.collect::<SqliteResult<Vec<T>>>()
        .map_err(|e| format!("Collect failed: {}", e))
}

/// Convert QueryParam to rusqlite::types::Value
fn convert_params(params: &[QueryParam]) -> Vec<rusqlite::types::Value> {
    params
        .iter()
        .map(|p| match p {
            QueryParam::String(s) => rusqlite::types::Value::Text(s.clone()),
            QueryParam::Int(i) => rusqlite::types::Value::Integer(*i as i64),
            QueryParam::Int64(i) => rusqlite::types::Value::Integer(*i),
            QueryParam::Float(f) => rusqlite::types::Value::Real(*f),
            QueryParam::Bool(b) => rusqlite::types::Value::Integer(if *b { 1 } else { 0 }),
            QueryParam::Null => rusqlite::types::Value::Null,
        })
        .collect()
}

// TRACES: UR-002, UR-012 | DR-012 | UT-014, UT-015, UT-016, UT-017, UT-018, UT-019, UT-020, UT-021, UT-022, UT-023, UT-025
#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_execute_query() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();

        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));

        let query = Query::with_params(
            "INSERT INTO test (name) VALUES (?)",
            vec![QueryParam::String("Alice".to_string())],
        );

        let rows = service.execute(query).await.unwrap();
        assert_eq!(rows, 1);
    }

    #[tokio::test]
    async fn test_query_one() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();
        conn.execute("INSERT INTO test (name) VALUES ('Bob')", [])
            .unwrap();

        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));

        let query = Query::new("SELECT name FROM test WHERE id = 1");
        let name: String = service.query_one(query, |row| row.get(0)).await.unwrap();

        assert_eq!(name, "Bob");
    }

    #[tokio::test]
    async fn test_query_many() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();
        conn.execute("INSERT INTO test (name) VALUES ('Alice')", [])
            .unwrap();
        conn.execute("INSERT INTO test (name) VALUES ('Bob')", [])
            .unwrap();

        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));

        let query = Query::new("SELECT name FROM test ORDER BY id");
        let names: Vec<String> = service.query_many(query, |row| row.get(0)).await.unwrap();

        assert_eq!(names, vec!["Alice", "Bob"]);
    }

    #[tokio::test]
    async fn test_query_optional() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();

        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));

        let query = Query::new("SELECT name FROM test WHERE id = 999");
        let result: Option<String> = service
            .query_optional(query, |row| row.get(0))
            .await
            .unwrap();

        assert_eq!(result, None);
    }

    #[tokio::test]
    async fn test_transaction() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();

        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));

        let result = service
            .transaction(|tx| {
                tx.execute(Query::with_params(
                    "INSERT INTO test (name) VALUES (?)",
                    vec![QueryParam::String("Alice".to_string())],
                ))?;
                tx.execute(Query::with_params(
                    "INSERT INTO test (name) VALUES (?)",
                    vec![QueryParam::String("Bob".to_string())],
                ))?;
                Ok(())
            })
            .await;

        assert!(result.is_ok());

        // Verify both rows were inserted
        let query = Query::new("SELECT COUNT(*) FROM test");
        let count: i32 = service.query_one(query, |row| row.get(0)).await.unwrap();
        assert_eq!(count, 2);
    }
    /// A panic while the connection guard is held must not brick every later
    /// query.
    ///
    /// This is the single busiest lock in the app — every async DB operation
    /// goes through it. With a raw `.lock()`, one panic under the guard poisons
    /// the mutex and every subsequent call returns "poisoned lock" until the
    /// process restarts, which for a database-backed app means the whole UI
    /// stops working. `utils::lock` exists precisely to stop that cascade, and
    /// `storage::Database` already used it; this path did not.
    ///
    /// TRACES: UR-002 | DR-012 | UT-014
    #[tokio::test]
    async fn a_poisoned_connection_still_serves_queries() {
        let conn = Arc::new(Mutex::new(Connection::open_in_memory().unwrap()));
        {
            let c = conn.lock_safe();
            c.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY);")
                .unwrap();
        }

        // Poison the mutex the way a panicking row mapper would.
        let poisoner = Arc::clone(&conn);
        let hook = std::panic::take_hook();
        std::panic::set_hook(Box::new(|_| {}));
        let _ = std::thread::spawn(move || {
            let _guard = poisoner.lock().unwrap();
            panic!("a row mapper blew up while holding the connection");
        })
        .join();
        std::panic::set_hook(hook);
        assert!(conn.lock().is_err(), "the mutex should now be poisoned");

        // Every operation must still work.
        let service = RusqliteService::new(Arc::clone(&conn));
        service
            .execute(Query::new("INSERT INTO test (id) VALUES (1)"))
            .await
            .expect("execute must survive a poisoned connection");
        let count: i32 = service
            .query_one(Query::new("SELECT COUNT(*) FROM test"), |row| row.get(0))
            .await
            .expect("query_one must survive a poisoned connection");
        assert_eq!(count, 1);
    }
}