launchbadge / sqlx

🧰 The Rust SQL Toolkit. An async, pure Rust SQL crate featuring compile-time checked queries without a DSL. Supports PostgreSQL, MySQL, and SQLite.
Apache License 2.0
12.48k stars 1.18k forks source link

Enabling loading of static compiled extensions #3147

Open prabirshrestha opened 3 months ago

prabirshrestha commented 3 months ago

I would like to load sqlite-vss extension for vector search. But rather than loading via an external library I would like to statically compile it as part of the binary so it can be embedded in the same rust program.

Is there an equivalent of this in sql?

use rusqlite::{ffi::sqlite3_auto_extension, Connection};
use sqlite_vss::{sqlite3_vector_init, sqlite3_vss_init};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    unsafe {
        sqlite3_auto_extension(Some(sqlite3_vector_init));
        sqlite3_auto_extension(Some(sqlite3_vss_init));
    }
}
abonander commented 3 months ago

I don't know if I like the sqlite3_auto_extension() mechanism. It seems like some spooky action-at-a-distance to me. To be fair, it's similar to how our Any driver works now, but at least that has some cross-checks in it: you have to explicitly request a driver via database URL whereas a SQLite3 extension can run arbitrary code on the database when it's initialized.

If you don't care, there's nothing stopping you from adding libsqlite3-sys to your dependencies and calling sqlite3_auto_extension yourself. Note that we consider the linkage to libsqlite3-sys to be semver-exempt, so you'll want to pin your sqlx dependency to avoid accidental breakage: https://docs.rs/sqlx/latest/sqlx/sqlite/index.html#note-linkage-is-semver-exempt

However, according to the SQLite documentation, you can also simply call the extension's init function with your database after it's initialized: https://www.sqlite.org/loadext.html#statically_linking_a_run_time_loadable_extension

You can get access to the raw sqlite3* pointer by calling SqliteConnection::lock_handle(). Things like this are what it's there for.

Strangely, the functions exported from sqlite-vss have the wrong signature for this. Given that libsqlite3_sys::sqlite3_auto_extension() expects the correct signature, though, this suggests that the types of extern fn pointers can be coerced to any other signature which I didn't realize was possible.

I would accept adding a method to SqliteConnectOptions to automatically call an extension's init function after creation, similar to sqlite3_auto_extension but constrained to just the database connections created from those options:

pub type SqliteExtensionInit = extern "C" fn(*mut sqlite3, *mut *const c_char, *mut sqlite3_api_routines) -> c_int;

impl SqliteConnectOptions {

    /// Add a statically linked extension to the database by directly calling its init function.
    ///
    /// This is similar to [`sqlite3_auto_extension()`](https://www.sqlite.org/c3ref/auto_extension.html)
    /// but is constrained just to databases created with this `SqliteConnectOptions`.
    ///
    /// ### SAFETY
    /// The function pointer must be safe to call from any thread for the lifetime of this `SqliteConnectOptions`.
    ///
    /// If you wish to dynamically load an extension, consider [`Self::extension()`] or [`Self::extension_with_entrypoint()`] instead.
    pub unsafe fn extension_with_init(
        &mut self,
        init: SqliteExtensionInit,
    ) -> &mut Self {
        // ..
    }
}
ospfranco commented 2 months ago

Thank you both for the pointers and the comments! After some finagling I was able to load the extension by directly calling libsqlite3_sys. Even though the signature does not match one can just unsafely cast it to whatever Rust neeeds:

use sqlite_vss::{sqlite3_vector_init, sqlite3_vss_init};

unsafe {
    let vss_vector_init = sqlite3_vector_init as *const ();
    let vss_vector_init_correct: extern "C" fn(
        db: *mut sqlite3,
        pzErrMsg: *mut *const c_char,
        pThunk: *const sqlite3_api_routines,
    ) -> i32 = std::mem::transmute(vss_vector_init);
    libsqlite3_sys::sqlite3_auto_extension(Some(vss_vector_init_correct));

    let vss_init = sqlite3_vss_init as *const ();
    let vss_init_correct: extern "C" fn(
        db: *mut sqlite3,
        pzErrMsg: *mut *const c_char,
        pThunk: *const sqlite3_api_routines,
    ) -> i32 = std::mem::transmute(vss_init);
    libsqlite3_sys::sqlite3_auto_extension(Some(vss_init_correct));
}

I also had to add the following to the build.rs after installing libomp from homebrew:

fn main() {
    if cfg!(target_os = "macos") {
        println!("cargo:rustc-link-arg=-Wl,-undefined,dynamic_lookup,-lomp,-L/opt/homebrew/opt/libomp/lib");
    } else if cfg!(target_os = "linux") {
        println!("cargo:rustc-link-arg=-Wl,-undefined,dynamic_lookup,-lstdc++");
    }
}

This should also work with rusqlite since it's directly calling the C binding to the sqlite3_auto_extension` function.