Add support for different SQLite journal modes (#193)

awslabs · Sep 30, 2024 · 1a1fa84 · 1a1fa84
1 parent cb5b9e9
commit 1a1fa84
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 3 deletions.
diff --git a/mls-rs-ffi/Cargo.toml b/mls-rs-ffi/Cargo.toml
@@ -22,6 +22,6 @@ x509 = ["mls-rs-identity-x509"]
 mls-rs = { path = "../mls-rs", version = "0.41.2", features = ["ffi"] }
 mls-rs-crypto-openssl = { path = "../mls-rs-crypto-openssl", version = "0.10.0", optional = true }
 mls-rs-identity-x509 = { path = "../mls-rs-identity-x509", version = "0.12.0", optional = true }
-mls-rs-provider-sqlite = { path = "../mls-rs-provider-sqlite", version = "0.13.0", default-features = false, optional = true }
+mls-rs-provider-sqlite = { path = "../mls-rs-provider-sqlite", version = "0.13.1", default-features = false, optional = true }
 safer-ffi = { version = "0.1.7", default-features = false }
 safer-ffi-gen = { version = "0.9.2", default-features = false }
diff --git a/mls-rs-provider-sqlite/Cargo.toml b/mls-rs-provider-sqlite/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "mls-rs-provider-sqlite"
-version = "0.13.0"
+version = "0.13.1"
 edition = "2021"
 description = "SQLite based state storage for mls-rs"
 homepage = "https://github.com/awslabs/mls-rs"

diff --git a/mls-rs-provider-sqlite/src/lib.rs b/mls-rs-provider-sqlite/src/lib.rs
@@ -54,13 +54,39 @@ impl mls_rs_core::error::IntoAnyError for SqLiteDataStorageError {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum JournalMode {
+    Delete,
+    Truncate,
+    Persist,
+    Memory,
+    Wal,
+    Off,
+}
+
+/// Note: for in-memory dbs (such as what the tests use), the only available options are MEMORY or OFF
+/// Invalid modes do not error, only no-op
+impl JournalMode {
+    fn as_str(&self) -> &'static str {
+        match self {
+            JournalMode::Delete => "DELETE",
+            JournalMode::Truncate => "TRUNCATE",
+            JournalMode::Persist => "PERSIST",
+            JournalMode::Memory => "MEMORY",
+            JournalMode::Wal => "WAL",
+            JournalMode::Off => "OFF",
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 /// SQLite data storage engine.
 pub struct SqLiteDataStorageEngine<CS>
 where
     CS: ConnectionStrategy,
 {
     connection_strategy: CS,
+    journal_mode: Option<JournalMode>,
 }
 
 impl<CS> SqLiteDataStorageEngine<CS>
@@ -72,9 +98,18 @@ where
     ) -> Result<SqLiteDataStorageEngine<CS>, SqLiteDataStorageError> {
         Ok(SqLiteDataStorageEngine {
             connection_strategy,
+            journal_mode: None,
         })
     }
 
+    /// A `journal_mode` of `None` means the SQLite default is used.
+    pub fn with_journal_mode(self, journal_mode: Option<JournalMode>) -> Self {
+        Self {
+            journal_mode,
+            ..self
+        }
+    }
+
     fn create_connection(&self) -> Result<Connection, SqLiteDataStorageError> {
         let connection = self.connection_strategy.make_connection()?;
 
@@ -83,6 +118,12 @@ where
             .pragma_query_value(None, "user_version", |rows| rows.get::<_, u32>(0))
             .map_err(|e| SqLiteDataStorageError::SqlEngineError(e.into()))?;
 
+        if let Some(journal_mode) = &self.journal_mode {
+            connection
+                .pragma_update(None, "journal_mode", journal_mode.as_str())
+                .map_err(|e| SqLiteDataStorageError::SqlEngineError(e.into()))?;
+        }
+
         if current_schema != 1 {
             create_tables_v1(&connection)?;
         }
@@ -152,7 +193,12 @@ fn create_tables_v1(connection: &Connection) -> Result<(), SqLiteDataStorageErro
 
 #[cfg(test)]
 mod tests {
-    use crate::{connection_strategy::MemoryStrategy, SqLiteDataStorageEngine};
+    use tempfile::tempdir;
+
+    use crate::{
+        connection_strategy::{FileConnectionStrategy, MemoryStrategy},
+        SqLiteDataStorageEngine,
+    };
 
     #[test]
     pub fn user_version_test() {
@@ -170,4 +216,26 @@ mod tests {
 
         assert_eq!(current_schema, 1);
     }
+
+    #[test]
+    pub fn journal_mode_test() {
+        let temp = tempdir().unwrap();
+
+        // Connect with journal_mode other than the default of MEMORY
+        let database = SqLiteDataStorageEngine::new(FileConnectionStrategy::new(
+            &temp.path().join("test_db.sqlite"),
+        ))
+        .unwrap();
+
+        let connection = database
+            .with_journal_mode(Some(crate::JournalMode::Truncate))
+            .create_connection()
+            .unwrap();
+
+        let journal_mode = connection
+            .pragma_query_value(None, "journal_mode", |rows| rows.get::<_, String>(0))
+            .unwrap();
+
+        assert_eq!(journal_mode, "truncate");
+    }
 }