Skip to main content

bones_sqlite_vec/
lib.rs

1use std::sync::OnceLock;
2
3const AUTO_ENABLE_ENV: &str = "BONES_SQLITE_VEC_AUTO";
4
5static REGISTRATION: OnceLock<Result<(), String>> = OnceLock::new();
6
7/// Register sqlite-vec as a process-wide `SQLite` auto extension.
8///
9/// # Errors
10///
11/// Returns an error when auto-registration is disabled via
12/// `BONES_SQLITE_VEC_AUTO` or when `sqlite3_auto_extension` fails.
13pub fn register_auto_extension() -> Result<(), String> {
14    if matches!(
15        std::env::var(AUTO_ENABLE_ENV).ok().as_deref(),
16        Some("0" | "false" | "off")
17    ) {
18        return Err(format!(
19            "sqlite-vec auto-extension disabled by {AUTO_ENABLE_ENV}"
20        ));
21    }
22
23    REGISTRATION.get_or_init(register_once).clone()
24}
25
26fn register_once() -> Result<(), String> {
27    #[allow(clippy::transmute_ptr_to_ptr)]
28    let entrypoint: unsafe extern "C" fn(
29        *mut rusqlite::ffi::sqlite3,
30        *mut *const std::os::raw::c_char,
31        *const rusqlite::ffi::sqlite3_api_routines,
32    ) -> std::os::raw::c_int =
33        unsafe { std::mem::transmute(sqlite_vec::sqlite3_vec_init as *const ()) };
34
35    let rc = unsafe { rusqlite::ffi::sqlite3_auto_extension(Some(entrypoint)) };
36    if rc == rusqlite::ffi::SQLITE_OK {
37        Ok(())
38    } else {
39        Err(format!("sqlite3_auto_extension failed with rc={rc}"))
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use rusqlite::Connection;
47
48    #[test]
49    fn registration_makes_vec_version_available() {
50        let result = register_auto_extension();
51        assert!(result.is_ok(), "registration failed: {result:?}");
52
53        let conn = Connection::open_in_memory().expect("open in-memory sqlite");
54        let version = conn.query_row("SELECT vec_version()", [], |row| row.get::<_, String>(0));
55        assert!(
56            version.is_ok(),
57            "vec_version() should be available after registration"
58        );
59    }
60}