jellytau_lib/storage/

db_service.rs

//! Database service abstraction layer
//!
//! This module provides an async database interface that abstracts away
//! the underlying database implementation. This makes it easy to:
//! - Switch between sync (rusqlite) and async (tokio-rusqlite) implementations
//! - Prevent blocking the async runtime with synchronous database calls
//! - Test with different database backends
//! - Migrate to other database systems in the future

use async_trait::async_trait;
use rusqlite::{params_from_iter, Connection, Result as SqliteResult, Row};
use std::sync::{Arc, 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;

    /// Get the row ID of the most recent successful INSERT
    async fn last_insert_rowid(&self) -> DbResult<i64>;
}

/// 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)
    }
}

/// Rusqlite-based database service implementation
///
/// This implementation wraps synchronous rusqlite operations in tokio::task::spawn_blocking
/// to prevent blocking the async runtime.
pub struct RusqliteService {
    conn: Arc<Mutex<Connection>>,
}

impl RusqliteService {
    pub fn new(conn: Arc<Mutex<Connection>>) -> Self {
        Self { conn }
    }
}

#[async_trait]
impl DatabaseService for RusqliteService {
    async fn execute(&self, query: Query) -> DbResult<usize> {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            execute_query(&conn, query)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn execute_batch(&self, sql: &str) -> DbResult<()> {
        let conn = Arc::clone(&self.conn);
        let sql = sql.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            conn.execute_batch(&sql)
                .map_err(|e| format!("Execute batch failed: {}", e))
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn query_one<T, F>(&self, query: Query, mapper: F) -> DbResult<T>
    where
        T: Send + 'static,
        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            query_one(&conn, query, mapper)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    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,
    {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            query_optional(&conn, query, mapper)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    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,
    {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            query_many(&conn, query, mapper)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn transaction<F, T>(&self, f: F) -> DbResult<T>
    where
        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
        T: Send + 'static,
    {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;

            conn.execute("BEGIN TRANSACTION", [])
                .map_err(|e| format!("Failed to begin transaction: {}", e))?;

            let mut transaction = Transaction::new(&conn);
            let result = f(&mut transaction);

            match result {
                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)
                }
            }
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn last_insert_rowid(&self) -> DbResult<i64> {
        let conn = Arc::clone(&self.conn);
        tokio::task::spawn_blocking(move || {
            let conn = conn
                .lock()
                .map_err(|e| format!("Failed to lock connection: {}", e))?;
            Ok(conn.last_insert_rowid())
        })
        .await
        .map_err(|e| format!("Task join error: {}", 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);
    }
}