[WIP] Refactored webext-storage database close function

components/webext-storage/src/api.rs

@@ -438,8 +438,9 @@ mod tests {
     #[test]
     fn test_simple() -> Result<()> {
         let ext_id = "x";
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
 
         // an empty store.
         for q in vec![JsonValue::Null, json!("foo"), json!(["foo"])].into_iter() {
@@ -529,8 +530,9 @@ mod tests {
     fn test_check_get_impl() -> Result<()> {
         // This is a port of checkGetImpl in test_ext_storage.js in Desktop.
         let ext_id = "x";
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
 
         let prop = "test-prop";
         let value = "test-value";
@@ -584,8 +586,9 @@ mod tests {
     fn test_bug_1621162() -> Result<()> {
         // apparently Firefox, unlike Chrome, will not optimize the changes.
         // See bug 1621162 for more!
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
 
         set(&tx, ext_id, json!({"foo": "bar" }))?;
@@ -599,8 +602,9 @@ mod tests {
 
     #[test]
     fn test_quota_maxitems() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         for i in 1..SYNC_MAX_ITEMS + 1 {
             set(
@@ -619,8 +623,9 @@ mod tests {
 
     #[test]
     fn test_quota_bytesperitem() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         // A string 5 bytes less than the max. This should be counted as being
         // 3 bytes less than the max as the quotes are counted. Plus the length
@@ -645,8 +650,9 @@ mod tests {
 
     #[test]
     fn test_quota_bytes() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         let val = "x".repeat(SYNC_QUOTA_BYTES + 1);
 
@@ -682,8 +688,9 @@ mod tests {
 
     #[test]
     fn test_get_bytes_in_use() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
 
         assert_eq!(get_bytes_in_use(&tx, ext_id, json!(null))?, 0);
@@ -714,8 +721,9 @@ mod tests {
 
     #[test]
     fn test_usage() {
-        let mut db = new_mem_db();
-        let tx = db.transaction().unwrap();
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction().unwrap();
         // '{"a":"a","b":"bb","c":"ccc","n":999999}': 39 bytes
         set(&tx, "xyz", json!({ "a": "a" })).unwrap();
         set(&tx, "xyz", json!({ "b": "bb" })).unwrap();
@@ -727,7 +735,7 @@ mod tests {
 
         tx.commit().unwrap();
 
-        let usage = usage(&db).unwrap();
+        let usage = usage(conn).unwrap();
         let expect = [
             UsageInfo {
                 ext_id: "abc".to_string(),

components/webext-storage/src/db.rs

@@ -11,7 +11,7 @@ use rusqlite::Connection;
 use rusqlite::OpenFlags;
 use sql_support::open_database::open_database_with_flags;
 use sql_support::ConnExt;
-use std::ops::{Deref, DerefMut};
+use std::ops::Deref;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use url::Url;
@@ -23,10 +23,16 @@ use url::Url;
 ///
 /// We only support a single writer connection - so that's the only thing we
 /// store. It's still a bit overkill, but there's only so many yaks in a day.
+pub enum WebExtStorageDb {
+    Open(Connection),
+    Closed,
+}
+
 pub struct StorageDb {
-    writer: Connection,
+    pub writer: WebExtStorageDb,
     interrupt_handle: Arc<SqlInterruptHandle>,
 }
+
 impl StorageDb {
     /// Create a new, or fetch an already open, StorageDb backed by a file on disk.
     pub fn new(db_path: impl AsRef<Path>) -> Result<Self> {
@@ -54,7 +60,7 @@ impl StorageDb {
         let conn = open_database_with_flags(db_path, flags, &schema::WebExtMigrationLogin)?;
         Ok(Self {
             interrupt_handle: Arc::new(SqlInterruptHandle::new(&conn)),
-            writer: conn,
+            writer: WebExtStorageDb::Open(conn),
         })
     }
 
@@ -73,29 +79,20 @@ impl StorageDb {
     /// underlying connection so the caller can retry but (a) that's very tricky
     /// in an Arc<Mutex<>> world and (b) we never actually took advantage of
     /// that retry capability.
-    pub fn close(self) -> Result<()> {
-        self.writer.close().map_err(|(writer, err)| {
-            // In rusqlite 0.28.0 and earlier, if we just let `writer` drop,
-            // the close would panic on failure.
-            // Later rusqlite versions will not panic, but this behavior doesn't
-            // hurt there.
-            std::mem::forget(writer);
-            err.into()
-        })
-    }
-}
-
-impl Deref for StorageDb {
-    type Target = Connection;
-
-    fn deref(&self) -> &Self::Target {
-        &self.writer
+    pub fn close(&mut self) -> Result<()> {
+        let conn = match std::mem::replace(&mut self.writer, WebExtStorageDb::Closed) {
+            WebExtStorageDb::Open(conn) => conn,
+            WebExtStorageDb::Closed => return Ok(()),
+        };
+        conn.close().map_err(|(_, y)| Error::SqlError(y))
     }
-}
 
-impl DerefMut for StorageDb {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.writer
+    pub(crate) fn get_connection(&self) -> Result<&Connection> {
+        let db = &self.writer;
+        match db {
+            WebExtStorageDb::Open(y) => Ok(y),
+            WebExtStorageDb::Closed => Err(Error::DatabaseConnectionClosed),
+        }
     }
 }
 
@@ -290,12 +287,13 @@ mod tests {
 
     #[test]
     fn test_meta() -> Result<()> {
-        let writer = new_mem_db();
-        assert_eq!(get_meta::<String>(&writer, "foo")?, None);
-        put_meta(&writer, "foo", &"bar".to_string())?;
-        assert_eq!(get_meta(&writer, "foo")?, Some("bar".to_string()));
-        delete_meta(&writer, "foo")?;
-        assert_eq!(get_meta::<String>(&writer, "foo")?, None);
+        let db = new_mem_db();
+        let conn = &db.get_connection()?;
+        assert_eq!(get_meta::<String>(conn, "foo")?, None);
+        put_meta(conn, "foo", &"bar".to_string())?;
+        assert_eq!(get_meta(conn, "foo")?, Some("bar".to_string()));
+        delete_meta(conn, "foo")?;
+        assert_eq!(get_meta::<String>(conn, "foo")?, None);
         Ok(())
     }
 }

components/webext-storage/src/migration.rs

@@ -346,8 +346,9 @@ mod tests {
         init_source_db(path, f);
 
         // now migrate
-        let mut db = new_mem_db();
-        let tx = db.transaction().expect("tx should work");
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction().expect("tx should work");
 
         let mi = migrate(&tx, &tmpdir.path().join("source.db")).expect("migrate should work");
         tx.commit().expect("should work");
@@ -384,17 +385,18 @@ mod tests {
     #[test]
     fn test_happy_paths() {
         // some real data.
-        let conn = do_migrate(HAPPY_PATH_MIGRATION_INFO, |c| {
+        let db = do_migrate(HAPPY_PATH_MIGRATION_INFO, |c| {
             c.execute_batch(HAPPY_PATH_SQL).expect("should populate")
         });
+        let conn = db.get_connection().expect("should retrieve connection");
 
         assert_has(
-            &conn,
+            conn,
             "{e7fefcf3-b39c-4f17-5215-ebfe120a7031}",
             json!({"userWelcomed": 1570659224457u64, "isWho": "4ec8109f"}),
         );
         assert_has(
-            &conn,
+            conn,
             "[email protected]",
             json!({"userRules": [], "ruleActiveStates": {}, "migration_version": 2}),
         );

components/webext-storage/src/schema.rs

@@ -94,22 +94,24 @@ mod tests {
     #[test]
     fn test_create_schema_twice() {
         let db = new_mem_db();
-        db.execute_batch(CREATE_SCHEMA_SQL)
+        let conn = db.get_connection().expect("should retrieve connection");
+        conn.execute_batch(CREATE_SCHEMA_SQL)
             .expect("should allow running twice");
     }
 
     #[test]
     fn test_create_empty_sync_temp_tables_twice() {
         let db = new_mem_db();
-        create_empty_sync_temp_tables(&db).expect("should work first time");
+        let conn = db.get_connection().expect("should retrieve connection");
+        create_empty_sync_temp_tables(conn).expect("should work first time");
         // insert something into our new temp table and check it's there.
-        db.execute_batch(
+        conn.execute_batch(
             "INSERT INTO temp.storage_sync_staging
                             (guid, ext_id) VALUES
                             ('guid', 'ext_id');",
         )
         .expect("should work once");
-        let count = db
+        let count = conn
             .query_row_and_then(
                 "SELECT COUNT(*) FROM temp.storage_sync_staging;",
                 [],
@@ -119,9 +121,9 @@ mod tests {
         assert_eq!(count, 1, "should be one row");
 
         // re-execute
-        create_empty_sync_temp_tables(&db).expect("should second first time");
+        create_empty_sync_temp_tables(conn).expect("should second first time");
         // and it should have deleted existing data.
-        let count = db
+        let count = conn
             .query_row_and_then(
                 "SELECT COUNT(*) FROM temp.storage_sync_staging;",
                 [],