Merge pull request #21 from jessekrubin/main

pool `conn_for_each`

Author: ryanfowler

src/pool.rs

@@ -238,4 +238,37 @@ impl Pool {
         let n = self.state.counter.fetch_add(1, Relaxed);
         &self.state.clients[n as usize % self.state.clients.len()]
     }
+
+    /// Runs a function on all connections in the pool asynchronously.
+    ///
+    /// The function is executed on each connection concurrently.
+    pub async fn conn_for_each<F, T>(&self, func: F) -> Vec<Result<T, Error>>
+    where
+        F: Fn(&Connection) -> Result<T, rusqlite::Error> + Send + Sync + 'static,
+        T: Send + 'static,
+    {
+        let func = Arc::new(func);
+        let futures = self.state.clients.iter().map(|client| {
+            let func = func.clone();
+            async move { client.conn(move |conn| func(conn)).await }
+        });
+        join_all(futures).await
+    }
+
+    /// Runs a function on all connections in the pool, blocking the current thread.
+    pub fn conn_for_each_blocking<F, T>(&self, func: F) -> Vec<Result<T, Error>>
+    where
+        F: Fn(&Connection) -> Result<T, rusqlite::Error> + Send + Sync + 'static,
+        T: Send + 'static,
+    {
+        let func = Arc::new(func);
+        self.state
+            .clients
+            .iter()
+            .map(|client| {
+                let func = func.clone();
+                client.conn_blocking(move |conn| func(conn))
+            })
+            .collect()
+    }
 }

tests/tests.rs

@@ -83,6 +83,7 @@ macro_rules! async_test {
 async_test!(test_journal_mode);
 async_test!(test_concurrency);
 async_test!(test_pool);
+async_test!(test_pool_conn_for_each);
 
 async fn test_journal_mode() {
     let tmp_dir = tempfile::tempdir().unwrap();
@@ -166,3 +167,63 @@ async fn test_pool() {
         .collect::<Result<(), Error>>()
         .expect("collecting query results");
 }
+
+async fn test_pool_conn_for_each() {
+    // make dummy db
+    let tmp_dir = tempfile::tempdir().unwrap();
+    {
+        let client = ClientBuilder::new()
+            .journal_mode(JournalMode::Wal)
+            .path(tmp_dir.path().join("sqlite.db"))
+            .open_blocking()
+            .expect("client unable to be opened");
+
+        client
+            .conn_blocking(|conn| {
+                conn.execute(
+                    "CREATE TABLE testing (id INTEGER PRIMARY KEY, val TEXT NOT NULL)",
+                    (),
+                )?;
+                conn.execute("INSERT INTO testing VALUES (1, ?)", ["value1"])
+            })
+            .expect("writing schema and seed data");
+    }
+
+    let pool = PoolBuilder::new()
+        .path(tmp_dir.path().join("another-sqlite.db"))
+        .num_conns(2)
+        .open()
+        .await
+        .expect("pool unable to be opened");
+
+    let dummy_db_path = tmp_dir.path().join("sqlite.db");
+    let attach_fn = move |conn: &rusqlite::Connection| {
+        conn.execute(
+            "ATTACH DATABASE ? AS dummy",
+            [dummy_db_path.to_str().unwrap()],
+        )
+    };
+    // attach to the dummy db via conn_for_each
+    pool.conn_for_each(attach_fn).await;
+
+    // check that the dummy db is attached
+    fn check_fn(conn: &rusqlite::Connection) -> Result<Vec<String>, rusqlite::Error> {
+        let mut stmt = conn
+            .prepare_cached("SELECT name FROM dummy.sqlite_master WHERE type='table'")
+            .unwrap();
+        let names = stmt
+            .query_map([], |row| row.get(0))
+            .unwrap()
+            .map(|r| r.unwrap())
+            .collect::<Vec<String>>();
+
+        Ok(names)
+    }
+    let res = pool.conn_for_each(check_fn).await;
+    for r in res {
+        assert_eq!(r.unwrap(), vec!["testing"]);
+    }
+
+    // cleanup
+    pool.close().await.expect("closing client conn");
+}