Skip to main content

ling_http/
pool.rs

1use r2d2::ManageConnection;
2use rusqlite::Connection;
3use std::path::PathBuf;
4
5/// A minimal `r2d2::ManageConnection` for rusqlite. Hand-written instead of
6/// depending on the `r2d2_sqlite` crate: that crate pins its own `rusqlite`
7/// version, which never lines up with the `rusqlite` version `ling-ai`
8/// already pulls in transitively (via the burn ML framework) — and two
9/// different `rusqlite`/`libsqlite3-sys` versions can't coexist in one
10/// Cargo.lock (both declare `links = "sqlite3"`). `r2d2` itself has no
11/// sqlite dependency, so implementing this ourselves keeps us on exactly
12/// one rusqlite version, chosen by us, everywhere in the workspace.
13pub enum SqliteConnectionManager {
14    File(PathBuf),
15    SharedMemory { uri: String },
16}
17
18impl SqliteConnectionManager {
19    pub fn file(path: impl Into<PathBuf>) -> Self {
20        Self::File(path.into())
21    }
22
23    /// A named, shared-cache in-memory database (`file:<name>?mode=memory&cache=shared`).
24    pub fn shared_memory(uri: impl Into<String>) -> Self {
25        Self::SharedMemory { uri: uri.into() }
26    }
27}
28
29impl ManageConnection for SqliteConnectionManager {
30    type Connection = Connection;
31    type Error = rusqlite::Error;
32
33    fn connect(&self) -> Result<Connection, rusqlite::Error> {
34        let conn = match self {
35            Self::File(path) => Connection::open(path)?,
36            Self::SharedMemory { uri } => Connection::open_with_flags(
37                uri,
38                rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE
39                    | rusqlite::OpenFlags::SQLITE_OPEN_CREATE
40                    | rusqlite::OpenFlags::SQLITE_OPEN_URI,
41            )?,
42        };
43        conn.execute_batch("PRAGMA foreign_keys = ON;")?;
44        Ok(conn)
45    }
46
47    fn is_valid(&self, conn: &mut Connection) -> Result<(), rusqlite::Error> {
48        conn.execute_batch("SELECT 1;")
49    }
50
51    fn has_broken(&self, _conn: &mut Connection) -> bool {
52        false
53    }
54}