use rusqlite::{Connection, LoadExtensionGuard};
use std::env;
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::path::PathBuf;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
use once_cell::sync::Lazy;
use crate::error::VecXError;
static EXT_LOAD_MUTEX: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
pub fn load_sqlite_vector_extension(conn: &mut Connection) -> Result<(), VecXError> {
load_sqlite_vector_extension_with_bytes(conn, embedded_lib_bytes())
}
fn create_unique_temp_filename(attempts: u8) -> PathBuf {
#[cfg(target_os = "linux")]
let ext = "so";
#[cfg(target_os = "macos")]
let ext = "dylib";
#[cfg(target_os = "windows")]
let ext = "dll";
let mut tmp_path = env::temp_dir();
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default();
let pid = std::process::id();
let file_name = format!("vectorlite_{}_{}_{}.{}", pid, now, attempts, ext);
tmp_path.push(file_name);
tmp_path
}
fn load_sqlite_vector_extension_with_bytes(
conn: &mut Connection,
lib_bytes: &[u8],
) -> Result<(), VecXError> {
let _guard = EXT_LOAD_MUTEX.lock().expect("mutex poisoned");
let mut attempts = 0u8;
let mut tmp_path: PathBuf;
loop {
tmp_path = create_unique_temp_filename(attempts);
let open_res = OpenOptions::new()
.write(true)
.create_new(true) .open(&tmp_path);
match open_res {
Ok(mut file) => {
if let Err(e) = file.write_all(lib_bytes) {
return Err(VecXError::from(e));
}
if let Err(e) = file.sync_all() {
return Err(VecXError::from(e));
}
break;
}
Err(e) => {
if e.kind() == std::io::ErrorKind::AlreadyExists && attempts < 5 {
attempts += 1;
continue;
} else {
return Err(VecXError::from(e));
}
}
}
}
unsafe {
let _guard = LoadExtensionGuard::new(conn)?;
conn.load_extension(
tmp_path
.to_str()
.ok_or_else(|| VecXError::ExtensionLoadError("invalid temporary path".into()))?,
None::<&str>,
)?;
}
#[cfg(unix)]
{
let _ = fs::remove_file(&tmp_path);
}
#[cfg(windows)]
{
}
Ok(())
}
fn embedded_lib_bytes() -> &'static [u8] {
#[cfg(target_os = "linux")]
{
include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/vectorlite.so"))
}
#[cfg(target_os = "macos")]
{
include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/assets/vectorlite.dylib"
))
}
#[cfg(target_os = "windows")]
{
include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/assets/vectorlite.dll"
))
}
}