Writing SQLite Extensions in Rust

Posted on by David Vassallo

Explainer

  • SQLite allows you to create custom virtual tables which allow you to query an external data source using SQL, lowering the bar to entry and homogenizing your data sources
  • In the past we’ve looked at creating virtual tables in golang, and here we do something similar in rust.
  • One notable improvement here is the inclusion of a loadable extension.

Setup: Dependencies

  • In Cargo.toml:
[dependencies]
rusqlite = { version = "0.31.0", features = ["vtab", "loadable_extension", "trace"] }

[lib]
name = "demo"
crate-type   = ["cdylib"] #this is important as the default `rlib` wont work
  • We are creating a shared library, so no main.rs is needed. instead, create a lib.rs file in the src folder instead

Step 1 – Creating the VTab

  • The rusqlite documentation has the bare-bones procedure already:
    • Write implementation of VTab and VTabCursor traits.
    • Create an instance of the Module structure specialized for VTab impl. from step 1.
    • Register your Module structure using Connection::create_module.
    • Run a CREATE VIRTUAL TABLE command that specifies the new module in the USING clause.

use std::{ffi::c_int, marker::PhantomData};
use rusqlite::ffi;
use rusqlite::{ffi::{sqlite3_vtab, sqlite3_vtab_cursor}, vtab::{Context, IndexInfo, VTab, VTabConnection, VTabCursor, Values}, Connection};
use std::os::raw::c_char;
use rusqlite::{to_sqlite_error, Result};
#[repr(C)]
struct Test {
/// Base class. Must be first
base: sqlite3_vtab,
}
#[derive(Default)]
#[repr(C)]
struct TestCursor<'vtab> {
/// Base class. Must be first
base: sqlite3_vtab_cursor,
/// The rowid
row_id: i64,
phantom: PhantomData<&'vtab Test>,
data: Vec<String>, // this is where we load and store our "external" data – see `open`
}
// Write implementation of VTab trait.
// Step 1(a) from https://docs.rs/rusqlite/latest/rusqlite/vtab/index.html
unsafe impl<'vtab> VTab<'vtab> for Test {
type Aux = ();
type Cursor = TestCursor<'vtab>;
fn connect(
_: &mut VTabConnection,
_aux: Option<&()>,
_args: &[&[u8]],
) -> Result<(String, Test), rusqlite::Error> {
let vtab = Test {
base: sqlite3_vtab::default()
};
// our vtab schema is defined here
Ok(("CREATE TABLE test(id INT, name TEXT)".to_owned(), vtab))
}
fn best_index(&self, info: &mut IndexInfo) -> Result<(), rusqlite::Error> {
info.set_estimated_cost(1.);
Ok(())
}
// this is where we do external calls (e.g. APIs, files)
// to populate our external data
fn open(&'vtab mut self) -> Result<TestCursor<'vtab>, rusqlite::Error> {
let mut test_cursor = TestCursor::default();
test_cursor.data = vec!["a".to_owned(), "b".to_owned(), "c".to_owned()];
Ok(test_cursor)
}
}
// Write implementation of VTabCursor trait.
// Step 1(b) from https://docs.rs/rusqlite/latest/rusqlite/vtab/index.html
unsafe impl VTabCursor for TestCursor<'_> {
fn filter(
&mut self,
_idx_num: c_int,
_idx_str: Option<&str>,
_args: &Values<'_>,
) -> Result<(), rusqlite::Error> {
Ok(())
}
// next – how do we get the next record?
fn next(&mut self) -> Result<(), rusqlite::Error> {
self.row_id += 1;
Ok(())
}
// EOF – when should we stop calling `next`?
fn eof(&self) -> bool {
self.row_id >= self.data.len().try_into().unwrap()
}
// descibe the mappings between columns (expressed as numbers) and our data
// stored in the cursor
fn column(&self, ctx: &mut Context, col_number: c_int) -> Result<(), rusqlite::Error> {
match col_number {
0 => ctx.set_result(&self.row_id),
1 => ctx.set_result(&self.data[self.row_id as usize]),
_ => Err(rusqlite::Error::InvalidColumnName("n/a".to_owned())),
}
}
fn rowid(&self) -> Result<i64, rusqlite::Error> {
Ok(self.row_id)
}
}
view raw vtab_1.rs hosted with ❤ by GitHub

Step 2 – loadable extension

  • Once we have our vtab defined above, we need to make it available to the sqlite3 command
    • sqlite3 has the .load command which lets you load runtime extensions like our vtab
  • When .load is executed, sqlite looks for a specific function within our shared object. Unless otherwise specified this will default to sqlite3_extension_init, so this is the first function we need to define:
#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[no_mangle]
pub extern "C" fn sqlite3_extension_init(
db: *mut ffi::sqlite3,
pz_err_msg: *mut *mut c_char,
p_api: *mut ffi::sqlite3_api_routines,
) -> c_int {
if p_api.is_null() {
return ffi::SQLITE_ERROR;
} else if let Err(err) = extension_init(db, p_api) { // we handoff to `extenstion_init`
return unsafe { to_sqlite_error(&err, pz_err_msg) };
}
ffi::SQLITE_OK
}
fn extension_init(db: *mut ffi::sqlite3, p_api: *mut ffi::sqlite3_api_routines) -> Result<()> {
let db = unsafe { Connection::extension_init2(db, p_api)? };
let module = rusqlite::vtab::eponymous_only_module::<Test>();
db.create_module::<Test>("test", module, None)?;
rusqlite::trace::log(ffi::SQLITE_WARNING, "Rusqlite extension initialized");
Ok(())
}
view raw vtab_2.rs hosted with ❤ by GitHub
  • In line 10 above, not that if we get a successful connection to the DB, we hand off to the extension_init function which we define right after
  • A lot of this is made super easy due to the bridging provided by rusqlite

Step 3 – build and test

  • cargo build --lib
  • sqlite3
    • .load /full/path/to/demo.so

Using the vtab in your rust program

What if we need to use our vtab not as a loadable extenstion, but as part of our rust program? The process remains the same as above, using just step 1. Step 2 is not required. Instead:

  • Create main.rs
fn main() -> Result<(), rusqlite::Error> {
let db = Connection::open_in_memory()?;
let module = rusqlite::vtab::eponymous_only_module::<Test>();
db.create_module::<Test>("test", module, None)?;
// normal rusqlite code from here on out…
println!("Hello, world!");
let mut s = db.prepare("SELECT * FROM test WHERE id=1")?;
let mut rows = s.query([])?;
while let Some(row) = rows.next()? {
let id: i64 = row.get(0)?;
let name: String = row.get(1)?;
println!("id: {}, name: {}", id, name);
}
let mut s = db.prepare("SELECT name FROM test WHERE id > 2")?;
let mut rows = s.query([])?;
while let Some(row) = rows.next()? {
let name: String = row.get(0)?;
println!("name: {}", name);
}
Ok(())
}
view raw vtab_in_main.rs hosted with ❤ by GitHub

References

Alternatives