Skip to main content

#
jellytau_lib/storage/
db_service.rs

    1//! Database service abstraction layer
    2//!
    3//! This module provides an async database interface that abstracts away
    4//! the underlying database implementation. This makes it easy to:
    5//! - Switch between sync (rusqlite) and async (tokio-rusqlite) implementations
    6//! - Prevent blocking the async runtime with synchronous database calls
    7//! - Test with different database backends
    8//! - Migrate to other database systems in the future
    9
    10use async_trait::async_trait;
    11use rusqlite::{params_from_iter, Connection, Result as SqliteResult, Row};
    12use std::sync::{Arc, Mutex};
    13
    14/// Database query result type
    15pub type DbResult<T> = Result<T, String>;
    16
    17/// Represents a database query that can be executed
    18#[derive(Clone)]
    19pub struct Query {
    20    pub sql: String,
    21    pub params: Vec<QueryParam>,
    22}
    23
    24/// Query parameter types supported by the database
    25#[derive(Clone, Debug)]
    26pub enum QueryParam {
    27    String(String),
    28    Int(i32),
    29    Int64(i64),
    30    Float(f64),
    31    #[allow(dead_code)]
    32    Bool(bool),
    33    Null,
    34}
    35
    36impl Query {
    37    pub fn new(sql: impl Into<String>) -> Self {
    38        Self {
    39            sql: sql.into(),
    40            params: Vec::new(),
    41        }
    42    }
    43
    44    pub fn with_params(sql: impl Into<String>, params: Vec<QueryParam>) -> Self {
    45        Self {
    46            sql: sql.into(),
    47            params,
    48        }
    49    }
    50}
    51
    52/// Database service trait - abstraction over database operations
    53#[async_trait]
    54pub trait DatabaseService: Send + Sync {
    55    /// Execute a query that doesn't return results (INSERT, UPDATE, DELETE)
    56    async fn execute(&self, query: Query) -> DbResult<usize>;
    57
    58    /// Execute a batch of SQL statements (for migrations)
    59    #[allow(dead_code)]
    60    async fn execute_batch(&self, sql: &str) -> DbResult<()>;
    61
    62    /// Query a single row
    63    async fn query_one<T, F>(&self, query: Query, mapper: F) -> DbResult<T>
    64    where
    65        T: Send + 'static,
    66        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;
    67
    68    /// Query a single optional row
    69    async fn query_optional<T, F>(&self, query: Query, mapper: F) -> DbResult<Option<T>>
    70    where
    71        T: Send + 'static,
    72        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;
    73
    74    /// Query multiple rows
    75    async fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    76    where
    77        T: Send + 'static,
    78        F: Fn(&Row) -> SqliteResult<T> + Send + 'static;
    79
    80    /// Run a transaction with multiple operations
    81    async fn transaction<F, T>(&self, f: F) -> DbResult<T>
    82    where
    83        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
    84        T: Send + 'static;
    85
    86    /// Get the row ID of the most recent successful INSERT
    87    async fn last_insert_rowid(&self) -> DbResult<i64>;
    88}
    89
    90/// Transaction handle for batching multiple operations
    91pub struct Transaction<'a> {
    92    conn: &'a Connection,
    93}
    94
    95impl<'a> Transaction<'a> {
    96    pub fn new(conn: &'a Connection) -> Self {
    97        Self { conn }
    98    }
    99
    100    pub fn execute(&mut self, query: Query) -> DbResult<usize> {
    101        execute_query(self.conn, query)
    102    }
    103
    104    pub fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    105    where
    106        F: Fn(&Row) -> SqliteResult<T>,
    107    {
    108        query_many(self.conn, query, mapper)
    109    }
    110}
    111
    112/// Rusqlite-based database service implementation
    113///
    114/// This implementation wraps synchronous rusqlite operations in tokio::task::spawn_blocking
    115/// to prevent blocking the async runtime.
    116pub struct RusqliteService {
    117    conn: Arc<Mutex<Connection>>,
    118}
    119
    120impl RusqliteService {
    121    pub fn new(conn: Arc<Mutex<Connection>>) -> Self {
    122        Self { conn }
    123    }
    124}
    125
    126#[async_trait]
    127impl DatabaseService for RusqliteService {
    128    async fn execute(&self, query: Query) -> DbResult<usize> {
    129        let conn = Arc::clone(&self.conn);
    130        tokio::task::spawn_blocking(move || {
    131            let conn = conn
    132                .lock()
    133                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    134            execute_query(&conn, query)
    135        })
    136        .await
    137        .map_err(|e| format!("Task join error: {}", e))?
    138    }
    139
    140    async fn execute_batch(&self, sql: &str) -> DbResult<()> {
    141        let conn = Arc::clone(&self.conn);
    142        let sql = sql.to_string();
    143        tokio::task::spawn_blocking(move || {
    144            let conn = conn
    145                .lock()
    146                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    147            conn.execute_batch(&sql)
    148                .map_err(|e| format!("Execute batch failed: {}", e))
    149        })
    150        .await
    151        .map_err(|e| format!("Task join error: {}", e))?
    152    }
    153
    154    async fn query_one<T, F>(&self, query: Query, mapper: F) -> DbResult<T>
    155    where
    156        T: Send + 'static,
    157        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    158    {
    159        let conn = Arc::clone(&self.conn);
    160        tokio::task::spawn_blocking(move || {
    161            let conn = conn
    162                .lock()
    163                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    164            query_one(&conn, query, mapper)
    165        })
    166        .await
    167        .map_err(|e| format!("Task join error: {}", e))?
    168    }
    169
    170    async fn query_optional<T, F>(&self, query: Query, mapper: F) -> DbResult<Option<T>>
    171    where
    172        T: Send + 'static,
    173        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    174    {
    175        let conn = Arc::clone(&self.conn);
    176        tokio::task::spawn_blocking(move || {
    177            let conn = conn
    178                .lock()
    179                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    180            query_optional(&conn, query, mapper)
    181        })
    182        .await
    183        .map_err(|e| format!("Task join error: {}", e))?
    184    }
    185
    186    async fn query_many<T, F>(&self, query: Query, mapper: F) -> DbResult<Vec<T>>
    187    where
    188        T: Send + 'static,
    189        F: Fn(&Row) -> SqliteResult<T> + Send + 'static,
    190    {
    191        let conn = Arc::clone(&self.conn);
    192        tokio::task::spawn_blocking(move || {
    193            let conn = conn
    194                .lock()
    195                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    196            query_many(&conn, query, mapper)
    197        })
    198        .await
    199        .map_err(|e| format!("Task join error: {}", e))?
    200    }
    201
    202    async fn transaction<F, T>(&self, f: F) -> DbResult<T>
    203    where
    204        F: FnOnce(&mut Transaction) -> DbResult<T> + Send + 'static,
    205        T: Send + 'static,
    206    {
    207        let conn = Arc::clone(&self.conn);
    208        tokio::task::spawn_blocking(move || {
    209            let conn = conn
    210                .lock()
    211                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    212
    213            conn.execute("BEGIN TRANSACTION", [])
    214                .map_err(|e| format!("Failed to begin transaction: {}", e))?;
    215
    216            let mut transaction = Transaction::new(&conn);
    217            let result = f(&mut transaction);
    218
    219            match result {
    220                Ok(value) => {
    221                    conn.execute("COMMIT", [])
    222                        .map_err(|e| format!("Failed to commit transaction: {}", e))?;
    223                    Ok(value)
    224                }
    225                Err(e) => {
    226                    conn.execute("ROLLBACK", [])
    227                        .map_err(|e| format!("Failed to rollback transaction: {}", e))?;
    228                    Err(e)
    229                }
    230            }
    231        })
    232        .await
    233        .map_err(|e| format!("Task join error: {}", e))?
    234    }
    235
    236    async fn last_insert_rowid(&self) -> DbResult<i64> {
    237        let conn = Arc::clone(&self.conn);
    238        tokio::task::spawn_blocking(move || {
    239            let conn = conn
    240                .lock()
    241                .map_err(|e| format!("Failed to lock connection: {}", e))?;
    242            Ok(conn.last_insert_rowid())
    243        })
    244        .await
    245        .map_err(|e| format!("Task join error: {}", e))?
    246    }
    247}
    248
    249// Helper functions for executing queries synchronously
    250
    251fn execute_query(conn: &Connection, query: Query) -> DbResult<usize> {
    252    let params = convert_params(&query.params);
    253    conn.execute(&query.sql, params_from_iter(params.iter()))
    254        .map_err(|e| format!("Execute failed: {}", e))
    255}
    256
    257fn query_one<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<T>
    258where
    259    F: Fn(&Row) -> SqliteResult<T>,
    260{
    261    let params = convert_params(&query.params);
    262    conn.query_row(&query.sql, params_from_iter(params.iter()), mapper)
    263        .map_err(|e| format!("Query one failed: {}", e))
    264}
    265
    266fn query_optional<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<Option<T>>
    267where
    268    F: Fn(&Row) -> SqliteResult<T>,
    269{
    270    match query_one(conn, query, mapper) {
    271        Ok(value) => Ok(Some(value)),
    272        Err(e) if e.contains("Query returned no rows") || e.contains("QueryReturnedNoRows") => {
    273            Ok(None)
    274        }
    275        Err(e) => Err(e),
    276    }
    277}
    278
    279fn query_many<T, F>(conn: &Connection, query: Query, mapper: F) -> DbResult<Vec<T>>
    280where
    281    F: Fn(&Row) -> SqliteResult<T>,
    282{
    283    let params = convert_params(&query.params);
    284    let mut stmt = conn
    285        .prepare(&query.sql)
    286        .map_err(|e| format!("Prepare failed: {}", e))?;
    287
    288    let rows = stmt
    289        .query_map(params_from_iter(params.iter()), mapper)
    290        .map_err(|e| format!("Query map failed: {}", e))?;
    291
    292    rows.collect::<SqliteResult<Vec<T>>>()
    293        .map_err(|e| format!("Collect failed: {}", e))
    294}
    295
    296/// Convert QueryParam to rusqlite::types::Value
    297fn convert_params(params: &[QueryParam]) -> Vec<rusqlite::types::Value> {
    298    params
    299        .iter()
    300        .map(|p| match p {
    301            QueryParam::String(s) => rusqlite::types::Value::Text(s.clone()),
    302            QueryParam::Int(i) => rusqlite::types::Value::Integer(*i as i64),
    303            QueryParam::Int64(i) => rusqlite::types::Value::Integer(*i),
    304            QueryParam::Float(f) => rusqlite::types::Value::Real(*f),
    305            QueryParam::Bool(b) => rusqlite::types::Value::Integer(if *b { 1 } else { 0 }),
    306            QueryParam::Null => rusqlite::types::Value::Null,
    307        })
    308        .collect()
    309}
    310
    311// 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
    312#[cfg(test)]
    313mod tests {
    314    use super::*;
    315
    316    #[tokio::test]
    317    async fn test_execute_query() {
    318        let conn = Connection::open_in_memory().unwrap();
    319        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    320            .unwrap();
    321
    322        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));
    323
    324        let query = Query::with_params(
    325            "INSERT INTO test (name) VALUES (?)",
    326            vec![QueryParam::String("Alice".to_string())],
    327        );
    328
    329        let rows = service.execute(query).await.unwrap();
    330        assert_eq!(rows, 1);
    331    }
    332
    333    #[tokio::test]
    334    async fn test_query_one() {
    335        let conn = Connection::open_in_memory().unwrap();
    336        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    337            .unwrap();
    338        conn.execute("INSERT INTO test (name) VALUES ('Bob')", [])
    339            .unwrap();
    340
    341        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));
    342
    343        let query = Query::new("SELECT name FROM test WHERE id = 1");
    344        let name: String = service.query_one(query, |row| row.get(0)).await.unwrap();
    345
    346        assert_eq!(name, "Bob");
    347    }
    348
    349    #[tokio::test]
    350    async fn test_query_many() {
    351        let conn = Connection::open_in_memory().unwrap();
    352        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    353            .unwrap();
    354        conn.execute("INSERT INTO test (name) VALUES ('Alice')", [])
    355            .unwrap();
    356        conn.execute("INSERT INTO test (name) VALUES ('Bob')", [])
    357            .unwrap();
    358
    359        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));
    360
    361        let query = Query::new("SELECT name FROM test ORDER BY id");
    362        let names: Vec<String> = service.query_many(query, |row| row.get(0)).await.unwrap();
    363
    364        assert_eq!(names, vec!["Alice", "Bob"]);
    365    }
    366
    367    #[tokio::test]
    368    async fn test_query_optional() {
    369        let conn = Connection::open_in_memory().unwrap();
    370        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    371            .unwrap();
    372
    373        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));
    374
    375        let query = Query::new("SELECT name FROM test WHERE id = 999");
    376        let result: Option<String> = service
    377            .query_optional(query, |row| row.get(0))
    378            .await
    379            .unwrap();
    380
    381        assert_eq!(result, None);
    382    }
    383
    384    #[tokio::test]
    385    async fn test_transaction() {
    386        let conn = Connection::open_in_memory().unwrap();
    387        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
    388            .unwrap();
    389
    390        let service = RusqliteService::new(Arc::new(Mutex::new(conn)));
    391
    392        let result = service
    393            .transaction(|tx| {
    394                tx.execute(Query::with_params(
    395                    "INSERT INTO test (name) VALUES (?)",
    396                    vec![QueryParam::String("Alice".to_string())],
    397                ))?;
    398                tx.execute(Query::with_params(
    399                    "INSERT INTO test (name) VALUES (?)",
    400                    vec![QueryParam::String("Bob".to_string())],
    401                ))?;
    402                Ok(())
    403            })
    404            .await;
    405
    406        assert!(result.is_ok());
    407
    408        // Verify both rows were inserted
    409        let query = Query::new("SELECT COUNT(*) FROM test");
    410        let count: i32 = service.query_one(query, |row| row.get(0)).await.unwrap();
    411        assert_eq!(count, 2);
    412    }
    413}