1use std::sync::OnceLock;
2
3const AUTO_ENABLE_ENV: &str = "BONES_SQLITE_VEC_AUTO";
4
5static REGISTRATION: OnceLock<Result<(), String>> = OnceLock::new();
6
7pub 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}