Enable WAL mode for extension connections (#11)

Author: bcho

absurd-sqlite-extension/src/lib.rs

@@ -68,6 +68,7 @@ fn absurd_init(db: *mut sqlite3) -> Result<()> {
             "jsonb() requires SQLite 3.45.0+; please upgrade SQLite",
         ));
     }
+    sql::ensure_wal_journal_mode(db)?;
     let flags = FunctionFlags::UTF8 | FunctionFlags::DETERMINISTIC;
     define_scalar_function(db, "absurd_version", 0, absurd_version, flags)?;
     define_scalar_function(db, "absurd_create_queue", 1, absurd_create_queue, flags)?;
@@ -229,6 +230,7 @@ mod tests {
     use super::*;
     use rusqlite::{ffi::sqlite3_auto_extension, params, Connection};
     use std::collections::HashMap;
+    use uuid::Uuid;
 
     #[test]
     fn test_absurd_version() {
@@ -1231,4 +1233,21 @@ mod tests {
             .unwrap();
         assert_eq!(run_state, "cancelled");
     }
+
+    #[test]
+    fn test_enables_wal_journal_mode_for_file_db() {
+        unsafe {
+            sqlite3_auto_extension(Some(std::mem::transmute(sqlite3_absurd_init as *const ())));
+        }
+
+        let db_path = std::env::temp_dir().join(format!("absurd-wal-{}.db", Uuid::new_v4()));
+        let conn = Connection::open(&db_path).unwrap();
+
+        let journal_mode: String = conn
+            .query_row("pragma journal_mode", [], |row| row.get(0))
+            .unwrap();
+        assert_eq!(journal_mode.to_lowercase(), "wal");
+
+        let _ = std::fs::remove_file(db_path);
+    }
 }

absurd-sqlite-extension/src/sql.rs

@@ -3,7 +3,7 @@ use sqlite_loadable::ext::{
     sqlite3ext_bind_text, sqlite3ext_finalize, sqlite3ext_prepare_v2, sqlite3ext_step,
 };
 use sqlite_loadable::{Error, Result, SQLITE_DONE, SQLITE_OKAY, SQLITE_ROW};
-use std::ffi::CString;
+use std::ffi::{CStr, CString};
 use std::os::raw::{c_char, c_void};
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -174,3 +174,40 @@ pub fn now_ms_from_db(db: *mut sqlite3) -> i64 {
 
     now_ms()
 }
+
+pub fn ensure_wal_journal_mode(db: *mut sqlite3) -> Result<()> {
+    let mut stmt = sqlite_loadable::exec::Statement::prepare(db, "pragma journal_mode=WAL")
+        .map_err(|err| Error::new_message(format!("failed to enable WAL mode: {:?}", err)))?;
+    let mut rows = stmt.execute();
+    let row = match rows.next() {
+        Some(Ok(row)) => row,
+        _ => {
+            return Err(Error::new_message(
+                "failed to read journal_mode after enabling WAL",
+            ))
+        }
+    };
+    let mode = row
+        .get::<String>(0)
+        .map_err(|err| Error::new_message(format!("failed to read journal_mode: {:?}", err)))?;
+    if mode.eq_ignore_ascii_case("wal") || db_is_in_memory(db) {
+        return Ok(());
+    }
+    Err(Error::new_message(format!(
+        "expected journal_mode to be WAL, found {}",
+        mode
+    )))
+}
+
+fn db_is_in_memory(db: *mut sqlite3) -> bool {
+    let main = CString::new("main").unwrap();
+    let filename_ptr = unsafe { sqlite3ext_sys::sqlite3_db_filename(db, main.as_ptr()) };
+    if filename_ptr.is_null() {
+        return true;
+    }
+    let filename = unsafe { CStr::from_ptr(filename_ptr) }
+        .to_string_lossy()
+        .into_owned();
+    let trimmed = filename.trim();
+    trimmed.is_empty() || trimmed == ":memory:"
+}