Merge pull request #16 from superfly/tweak-conn-init

Fix CrConn's Drop
superfly · Aug 9, 2023 · d69104a · d69104a
2 parents e3ac00b + e5cda76
commit d69104a
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 19 deletions.
diff --git a/crates/corro-agent/src/agent.rs b/crates/corro-agent/src/agent.rs
@@ -31,7 +31,7 @@ use corro_types::{
     config::{Config, DEFAULT_GOSSIP_PORT},
     members::{MemberEvent, Members},
     schema::init_schema,
-    sqlite::{init_cr_conn, CrConn, Migration, SqlitePoolError},
+    sqlite::{CrConn, Migration, SqlitePoolError},
     sync::{generate_sync, SyncMessageDecodeError, SyncMessageEncodeError},
 };
 
@@ -91,8 +91,7 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age
     }
 
     let actor_id = {
-        let mut conn = CrConn(Connection::open(&conf.db_path)?);
-        init_cr_conn(&mut conn)?;
+        let conn = CrConn::init(Connection::open(&conf.db_path)?)?;
         conn.query_row("SELECT crsql_site_id();", [], |row| {
             row.get::<_, ActorId>(0)
         })?
@@ -2320,9 +2319,7 @@ pub mod tests {
 
     #[test]
     fn test_in_memory_versions_compaction() -> eyre::Result<()> {
-        let mut conn = rusqlite::Connection::open_in_memory()?;
-
-        init_cr_conn(&mut conn)?;
+        let conn = CrConn::init(rusqlite::Connection::open_in_memory()?)?;
 
         conn.execute_batch(
             "CREATE TABLE foo (a INTEGER PRIMARY KEY, b INTEGER); SELECT crsql_as_crr('foo');",

diff --git a/crates/corro-types/src/pubsub.rs b/crates/corro-types/src/pubsub.rs
@@ -1167,7 +1167,7 @@ mod tests {
     use crate::{
         change::{row_to_change, SqliteValue},
         schema::{make_schema_inner, parse_sql},
-        sqlite::{init_cr_conn, setup_conn},
+        sqlite::{setup_conn, CrConn},
     };
 
     use super::*;
@@ -1185,9 +1185,7 @@ mod tests {
         let tmpdir = tempfile::tempdir()?;
         let db_path = tmpdir.path().join("test.db");
 
-        let mut conn = rusqlite::Connection::open(&db_path).expect("could not open conn");
-
-        init_cr_conn(&mut conn)?;
+        let mut conn = CrConn::init(rusqlite::Connection::open(&db_path)?)?;
 
         setup_conn(
             &mut conn,
@@ -1314,9 +1312,9 @@ mod tests {
         let tmpdir = tempfile::tempdir().unwrap();
         let db_path = tmpdir.path().join("test.db");
 
-        let mut conn = rusqlite::Connection::open(&db_path).expect("could not open conn");
-
-        init_cr_conn(&mut conn).unwrap();
+        let mut conn =
+            CrConn::init(rusqlite::Connection::open(&db_path).expect("could not open conn"))
+                .expect("could not init crsql");
 
         setup_conn(
             &mut conn,
@@ -1363,10 +1361,11 @@ mod tests {
         }
 
         {
-            let mut conn2 = rusqlite::Connection::open(tmpdir.path().join("test2.db"))
-                .expect("could not open conn");
-
-            init_cr_conn(&mut conn2).unwrap();
+            let mut conn2 = CrConn::init(
+                rusqlite::Connection::open(tmpdir.path().join("test2.db"))
+                    .expect("could not open conn"),
+            )
+            .expect("could not init crsql");
 
             setup_conn(
                 &mut conn2,

diff --git a/crates/corro-types/src/sqlite.rs b/crates/corro-types/src/sqlite.rs
@@ -137,9 +137,14 @@ impl ManageConnection for CrConnManager {
 }
 
 #[derive(Debug)]
-pub struct CrConn(pub Connection);
+pub struct CrConn(Connection);
 
 impl CrConn {
+    pub fn init(mut conn: Connection) -> Result<Self, rusqlite::Error> {
+        init_cr_conn(&mut conn)?;
+        Ok(Self(conn))
+    }
+
     pub fn transaction(&mut self) -> rusqlite::Result<Transaction> {
         self.0
             .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
@@ -168,7 +173,7 @@ impl Drop for CrConn {
     }
 }
 
-pub fn init_cr_conn(conn: &mut Connection) -> Result<(), rusqlite::Error> {
+fn init_cr_conn(conn: &mut Connection) -> Result<(), rusqlite::Error> {
     let ext_dir = &CRSQL_EXT_DIR;
     trace!(
         "loading crsqlite extension from path: {}",