toolhub_storage/
pool.rs

1//! r2d2 connection pool for the ToolHub SQLite DB.
2//!
3//! The pool is wired so every connection it hands out has the `sqlite-vec`
4//! extension auto-loaded (via [`crate::ensure_vec_extension`]) and finds the
5//! schema already migrated. Open the pool once at startup and clone it into
6//! axum/tower handlers.
7
8use std::path::Path;
9
10use anyhow::Context;
11use r2d2::Pool;
12use r2d2_sqlite::SqliteConnectionManager;
13
14/// Open and migrate the DB at `path`, returning a ready-to-use connection
15/// pool. Migrations run on a single bootstrap connection that is dropped
16/// before the pool is built, so pooled handles never see a half-applied DB.
17pub fn open_pool(path: &Path) -> anyhow::Result<Pool<SqliteConnectionManager>> {
18    crate::ensure_vec_extension();
19    // Bootstrap: run migrations on a single connection. `crate::open` is
20    // idempotent (refinery skips already-applied migrations), so this is
21    // cheap on subsequent calls too.
22    let _bootstrap =
23        crate::open(path).with_context(|| format!("bootstrap migrate {}", path.display()))?;
24    drop(_bootstrap);
25
26    let manager = SqliteConnectionManager::file(path);
27    let pool = Pool::builder()
28        .max_size(8)
29        .build(manager)
30        .with_context(|| format!("build r2d2 pool for {}", path.display()))?;
31    Ok(pool)
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    #[test]
39    fn open_pool_creates_db_and_runs_migrations() {
40        let dir = tempfile::tempdir().unwrap();
41        let path = dir.path().join("toolhub.sqlite");
42
43        let pool = open_pool(&path).unwrap();
44        let conn = pool.get().unwrap();
45
46        let mut stmt = conn
47            .prepare(
48                "SELECT name FROM sqlite_master \
49                 WHERE type IN ('table','view') ORDER BY name",
50            )
51            .unwrap();
52        let names: Vec<String> = stmt
53            .query_map([], |row| row.get::<_, String>(0))
54            .unwrap()
55            .map(|r| r.unwrap())
56            .collect();
57
58        for expected in [
59            "tools",
60            "usage_events",
61            "tool_scores",
62            "sources",
63            "tools_fts",
64            "agent_suggestions",
65        ] {
66            assert!(
67                names.contains(&expected.to_string()),
68                "missing {expected} in {names:?}"
69            );
70        }
71    }
72
73    #[test]
74    fn open_pool_is_idempotent() {
75        let dir = tempfile::tempdir().unwrap();
76        let path = dir.path().join("toolhub.sqlite");
77
78        let _p1 = open_pool(&path).unwrap();
79        let _p2 = open_pool(&path).unwrap();
80    }
81
82    #[test]
83    fn pooled_connection_can_run_vec_extension_query() {
84        let dir = tempfile::tempdir().unwrap();
85        let path = dir.path().join("toolhub.sqlite");
86
87        let pool = open_pool(&path).unwrap();
88        let conn = pool.get().unwrap();
89
90        // tools_vec is the sqlite-vec virtual table; if the extension didn't
91        // load on this pooled connection, this query would error.
92        let count: i64 = conn
93            .query_row("SELECT COUNT(*) FROM tools_vec", [], |row| row.get(0))
94            .unwrap();
95        assert_eq!(count, 0);
96    }
97
98    #[test]
99    fn pool_hands_out_independent_connections() {
100        let dir = tempfile::tempdir().unwrap();
101        let path = dir.path().join("toolhub.sqlite");
102
103        let pool = open_pool(&path).unwrap();
104        let c1 = pool.get().unwrap();
105        let c2 = pool.get().unwrap();
106
107        // Both must be usable concurrently for read.
108        c1.query_row("SELECT 1", [], |row| row.get::<_, i64>(0))
109            .unwrap();
110        c2.query_row("SELECT 1", [], |row| row.get::<_, i64>(0))
111            .unwrap();
112    }
113}
toolhub_storage/pool.rs

toolhub_storage/
pool.rs
