aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-02-08 09:44:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

Introduce VTable

Browse Source 
This allows storing table arguments parsed in the VTabModule::create
method.
This commit is contained in:
Piotr Rzysko
2025-05-19 06:29:20 +02:00
parent 6b454ea36f
commit 9c1dca72db
14 changed files with 424 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
core/lib.rs
View File

@@ -45,6 +45,7 @@ use limbo_ext::{ConstraintInfo, IndexInfo, OrderByInfo, ResultCode, VTabKind, VT
use limbo_sqlite3_parser::{ast, ast::Cmd, lexer::sql::Parser};
use parking_lot::RwLock;
use schema::{Column, Schema};
use std::ffi::c_void;
use std::{
borrow::Cow,
cell::{Cell, RefCell, UnsafeCell},
@@ -674,6 +675,7 @@ pub struct VirtualTable {
pub implementation: Rc<VTabModuleImpl>,
columns: Vec<Column>,
kind: VTabKind,
table_ptr: *const c_void,
}
impl VirtualTable {
@@ -719,7 +721,7 @@ impl VirtualTable {
)));
}
};
let schema = module.implementation.as_ref().init_schema(args)?;
let (schema, table_ptr) = module.implementation.as_ref().create(args)?;
let mut parser = Parser::new(schema.as_bytes());
if let ast::Cmd::Stmt(ast::Stmt::CreateTable { body, .. }) = parser.next()?.ok_or(
LimboError::ParseError("Failed to parse schema from virtual table module".to_string()),
@@ -731,6 +733,7 @@ impl VirtualTable {
columns,
args: exprs,
kind,
table_ptr,
});
return Ok(vtab);
}
@@ -740,7 +743,7 @@ impl VirtualTable {
}
pub fn open(&self) -> crate::Result<VTabOpaqueCursor> {
let cursor = unsafe { (self.implementation.open)(self.implementation.ctx) };
let cursor = unsafe { (self.implementation.open)(self.table_ptr) };
VTabOpaqueCursor::new(cursor, self.implementation.close)
}
@@ -797,10 +800,9 @@ impl VirtualTable {
let arg_count = args.len();
let ext_args = args.iter().map(|arg| arg.to_ffi()).collect::<Vec<_>>();
let newrowid = 0i64;
let implementation = self.implementation.as_ref();
let rc = unsafe {
(self.implementation.update)(
implementation as *const VTabModuleImpl as *const std::ffi::c_void,
self.table_ptr,
arg_count as i32,
ext_args.as_ptr(),
&newrowid as *const _ as *mut i64,
@@ -819,12 +821,7 @@ impl VirtualTable {
}
pub fn destroy(&self) -> Result<()> {
let implementation = self.implementation.as_ref();
let rc = unsafe {
(self.implementation.destroy)(
implementation as *const VTabModuleImpl as *const std::ffi::c_void,
)
};
let rc = unsafe { (self.implementation.destroy)(self.table_ptr) };
match rc {
ResultCode::OK => Ok(()),
_ => Err(LimboError::ExtensionError(rc.to_string())),

6
core/translate/emitter.rs
View File

@@ -461,7 +461,7 @@ fn emit_delete_insns(
dest: key_reg,
});
if let Some(vtab) = table_reference.virtual_table() {
if let Some(_) = table_reference.virtual_table() {
let conflict_action = 0u16;
let start_reg = key_reg;
@@ -474,7 +474,6 @@ fn emit_delete_insns(
cursor_id,
arg_count: 2,
start_reg,
vtab_ptr: vtab.implementation.as_ref().ctx as usize,
conflict_action,
});
} else {
@@ -1039,13 +1038,12 @@ fn emit_update_insns(
flag: 0,
table_name: table_ref.identifier.clone(),
});
} else if let Some(vtab) = table_ref.virtual_table() {
} else if let Some(_) = table_ref.virtual_table() {
let arg_count = table_ref.columns().len() + 2;
program.emit_insn(Insn::VUpdate {
cursor_id,
arg_count,
start_reg: beg,
vtab_ptr: vtab.implementation.as_ref().ctx as usize,
conflict_action: 0u16,
});
}

1
core/translate/insert.rs
View File

@@ -714,7 +714,6 @@ fn translate_virtual_table_insert(
cursor_id,
arg_count: column_mappings.len() + 2,
start_reg: registers_start,
vtab_ptr: virtual_table.implementation.as_ref().ctx as usize,
conflict_action,
});

2
core/translate/schema.rs
View File

@@ -505,7 +505,7 @@ fn create_vtable_body_to_str(vtab: &CreateVirtualTable, module: Rc<VTabImpl>) ->
.collect::<Vec<_>>();
let schema = module
.implementation
.init_schema(ext_args)
.create_schema(ext_args)
.unwrap_or_default();
let vtab_args = if let Some(first_paren) = schema.find('(') {
let closing_paren = schema.rfind(')').unwrap_or_default();

5
core/vdbe/explain.rs
View File

@@ -416,14 +416,13 @@ pub fn insn_to_str(
cursor_id,
arg_count, // P2: Number of arguments in argv[]
start_reg, // P3: Start register for argv[]
vtab_ptr, // P4: vtab pointer
conflict_action, // P5: Conflict resolution flags
conflict_action, // P4: Conflict resolution flags
} => (
"VUpdate",
*cursor_id as i32,
*arg_count as i32,
*start_reg as i32,
Value::build_text(&format!("vtab:{}", vtab_ptr)),
Value::build_text(""),
*conflict_action,
format!("args=r[{}..{}]", start_reg, start_reg + arg_count - 1),
),

3
core/vdbe/insn.rs
View File

@@ -323,8 +323,7 @@ pub enum Insn {
cursor_id: usize, // P1: Virtual table cursor number
arg_count: usize, // P2: Number of arguments in argv[]
start_reg: usize, // P3: Start register for argv[]
vtab_ptr: usize, // P4: vtab pointer
conflict_action: u16, // P5: Conflict resolution flags
conflict_action: u16, // P4: Conflict resolution flags
},
/// Advance the virtual table cursor to the next row.

97
extensions/completion/src/lib.rs
View File

@@ -4,10 +4,12 @@
mod keywords;
use keywords::KEYWORDS;
use limbo_ext::{register_extension, ResultCode, VTabCursor, VTabModule, VTabModuleDerive, Value};
use limbo_ext::{
register_extension, ResultCode, VTabCursor, VTabModule, VTabModuleDerive, VTable, Value,
};
register_extension! {
vtabs: { CompletionVTab }
vtabs: { CompletionVTabModule }
}
macro_rules! try_option {
@@ -57,61 +59,34 @@ impl Into<i64> for CompletionPhase {
/// A virtual table that generates candidate completions
#[derive(Debug, Default, VTabModuleDerive)]
struct CompletionVTab {}
struct CompletionVTabModule {}
impl VTabModule for CompletionVTab {
type VCursor = CompletionCursor;
impl VTabModule for CompletionVTabModule {
type Table = CompletionTable;
const NAME: &'static str = "completion";
const VTAB_KIND: limbo_ext::VTabKind = limbo_ext::VTabKind::TableValuedFunction;
type Error = ResultCode;
fn create_schema(_args: &[Value]) -> String {
"CREATE TABLE completion(
fn create(_args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
let schema = "CREATE TABLE completion(
candidate TEXT,
prefix TEXT HIDDEN,
wholeline TEXT HIDDEN,
phase INT HIDDEN
)"
.to_string()
.to_string();
Ok((schema, CompletionTable {}))
}
}
fn open(&self) -> Result<Self::VCursor, Self::Error> {
struct CompletionTable {}
impl VTable for CompletionTable {
type Cursor = CompletionCursor;
type Error = ResultCode;
fn open(&self) -> Result<Self::Cursor, Self::Error> {
Ok(CompletionCursor::default())
}
fn filter(cursor: &mut Self::VCursor, args: &[Value], _: Option<(&str, i32)>) -> ResultCode {
if args.is_empty() || args.len() > 2 {
return ResultCode::InvalidArgs;
}
cursor.reset();
let prefix = try_option!(args[0].to_text(), ResultCode::InvalidArgs);
let wholeline = args.get(1).map(|v| v.to_text().unwrap_or("")).unwrap_or("");
cursor.line = wholeline.to_string();
cursor.prefix = prefix.to_string();
// Currently best index is not implemented so the correct arg parsing is not done here
if !cursor.line.is_empty() && cursor.prefix.is_empty() {
let mut i = cursor.line.len();
while let Some(ch) = cursor.line.chars().next() {
if i > 0 && (ch.is_alphanumeric() || ch == '_') {
i -= 1;
} else {
break;
}
}
if cursor.line.len() - i > 0 {
// TODO see if need to inclusive range
cursor.prefix = cursor.line[..i].to_string();
}
}
cursor.rowid = 0;
cursor.phase = CompletionPhase::Keywords;
cursor.next()
}
}
/// The cursor for iterating over the completions
@@ -138,6 +113,40 @@ impl CompletionCursor {
impl VTabCursor for CompletionCursor {
type Error = ResultCode;
fn filter(&mut self, args: &[Value], _: Option<(&str, i32)>) -> ResultCode {
if args.is_empty() || args.len() > 2 {
return ResultCode::InvalidArgs;
}
self.reset();
let prefix = try_option!(args[0].to_text(), ResultCode::InvalidArgs);
let wholeline = args.get(1).map(|v| v.to_text().unwrap_or("")).unwrap_or("");
self.line = wholeline.to_string();
self.prefix = prefix.to_string();
// Currently best index is not implemented so the correct arg parsing is not done here
if !self.line.is_empty() && self.prefix.is_empty() {
let mut i = self.line.len();
while let Some(ch) = self.line.chars().next() {
if i > 0 && (ch.is_alphanumeric() || ch == '_') {
i -= 1;
} else {
break;
}
}
if self.line.len() - i > 0 {
// TODO see if need to inclusive range
self.prefix = self.line[..i].to_string();
}
}
self.rowid = 0;
self.phase = CompletionPhase::Keywords;
self.next()
}
fn next(&mut self) -> ResultCode {
self.rowid += 1;

38
extensions/core/README.md
View File

@@ -172,29 +172,35 @@ impl AggFunc for Percentile {
/// Example: A virtual table that operates on a CSV file as a database table.
/// This example assumes that the CSV file is located at "data.csv" in the current directory.
#[derive(Debug, VTabModuleDerive)]
struct CsvVTable;
struct CsvVTableModule;
impl VTabModule for CsvVTable {
type VCursor = CsvCursor;
/// Define your error type. Must impl Display and match VCursor::Error
type Error = &'static str;
impl VTabModule for CsvVTableModule {
type Table = CsvTable;
/// Declare the name for your virtual table
const NAME: &'static str = "csv_data";
/// Declare the type of vtable (TableValuedFunction or VirtualTable)
const VTAB_KIND: VTabKind = VTabKind::VirtualTable;
/// Function to initialize the schema of your vtable
fn create_schema(_args: &[Value]) -> &'static str {
"CREATE TABLE csv_data(
/// Declare your virtual table and its schema
fn create(args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
let schema = "CREATE TABLE csv_data(
name TEXT,
age TEXT,
city TEXT
)"
)".into();
Ok((schema, CsvTable {}))
}
}
struct CsvTable {}
impl VTable for CsvTable {
type Cursor = CsvCursor;
/// Define your error type. Must impl Display and match Cursor::Error
type Error = &'static str;
/// Open to return a new cursor: In this simple example, the CSV file is read completely into memory on connect.
fn open(&self) -> Result<Self::VCursor, Self::Error> {
fn open(&self) -> Result<Self::Cursor, Self::Error> {
// Read CSV file contents from "data.csv"
let csv_content = fs::read_to_string("data.csv").unwrap_or_default();
// For simplicity, we'll ignore the header row.
@@ -210,11 +216,6 @@ impl VTabModule for CsvVTable {
Ok(CsvCursor { rows, index: 0 })
}
/// Filter through result columns. (not used in this simple example)
fn filter(_cursor: &mut Self::VCursor, _args: &[Value]) -> ResultCode {
ResultCode::OK
}
/// *Optional* methods for non-readonly tables
/// Update the value at rowid
@@ -243,6 +244,11 @@ struct CsvCursor {
impl VTabCursor for CsvCursor {
type Error = &'static str;
/// Filter through result columns. (not used in this simple example)
fn filter(&mut self, _args: &[Value], _idx_info: Option<(&str, i32)>) -> ResultCode {
ResultCode::OK
}
/// Next advances the cursor to the next row.
fn next(&mut self) -> ResultCode {
if self.index < self.rows.len() - 1 {

2
extensions/core/src/lib.rs
View File

@@ -17,7 +17,7 @@ pub use vfs_modules::{RegisterVfsFn, VfsExtension, VfsFile, VfsFileImpl, VfsImpl
use vtabs::RegisterModuleFn;
pub use vtabs::{
ConstraintInfo, ConstraintOp, ConstraintUsage, ExtIndexInfo, IndexInfo, OrderByInfo,
VTabCursor, VTabKind, VTabModule, VTabModuleImpl,
VTabCreateResult, VTabCursor, VTabKind, VTabModule, VTabModuleImpl, VTable,
};
pub type ExtResult<T> = std::result::Result<T, ResultCode>;

71
extensions/core/src/vtabs.rs
View File

@@ -11,9 +11,8 @@ pub type RegisterModuleFn = unsafe extern "C" fn(
#[repr(C)]
#[derive(Clone, Debug)]
pub struct VTabModuleImpl {
pub ctx: *const c_void,
pub name: *const c_char,
pub create_schema: VtabFnCreateSchema,
pub create: VtabFnCreate,
pub open: VtabFnOpen,
pub close: VtabFnClose,
pub filter: VtabFnFilter,
@@ -26,24 +25,50 @@ pub struct VTabModuleImpl {
pub best_idx: BestIdxFn,
}
#[repr(C)]
pub struct VTabCreateResult {
pub code: ResultCode,
pub schema: *const c_char,
pub table: *const c_void,
}
#[cfg(feature = "core_only")]
impl VTabModuleImpl {
pub fn init_schema(&self, args: Vec<Value>) -> crate::ExtResult<String> {
let schema = unsafe { (self.create_schema)(args.as_ptr(), args.len() as i32) };
if schema.is_null() {
return Err(ResultCode::InvalidArgs);
}
pub fn create(&self, args: Vec<Value>) -> crate::ExtResult<(String, *const c_void)> {
let result = unsafe { (self.create)(args.as_ptr(), args.len() as i32) };
for arg in args {
unsafe { arg.__free_internal_type() };
}
let schema = unsafe { std::ffi::CString::from_raw(schema) };
Ok(schema.to_string_lossy().to_string())
if !result.code.is_ok() {
return Err(result.code);
}
let schema = unsafe { std::ffi::CString::from_raw(result.schema as *mut _) };
Ok((schema.to_string_lossy().to_string(), result.table))
}
// TODO: This function is temporary and should eventually be removed.
// The only difference from `create` is that it takes ownership of the table instance.
// Currently, it is used to generate virtual table column names that are stored in
// `sqlite_schema` alongside the table's schema.
// However, storing column names is not necessary to match SQLite's behavior.
// SQLite computes the list of columns dynamically each time the `.schema` command
// is executed, using the `shell_add_schema` UDF function.
pub fn create_schema(&self, args: Vec<Value>) -> crate::ExtResult<String> {
self.create(args).and_then(|(schema, table)| {
// Drop the allocated table instance to avoid a memory leak.
let result = unsafe { (self.destroy)(table) };
if result.is_ok() {
Ok(schema)
} else {
Err(result)
}
})
}
}
pub type VtabFnCreateSchema = unsafe extern "C" fn(args: *const Value, argc: i32) -> *mut c_char;
pub type VtabFnCreate = unsafe extern "C" fn(args: *const Value, argc: i32) -> VTabCreateResult;
pub type VtabFnOpen = unsafe extern "C" fn(*const c_void) -> *const c_void;
pub type VtabFnOpen = unsafe extern "C" fn(table: *const c_void) -> *const c_void;
pub type VtabFnClose = unsafe extern "C" fn(cursor: *const c_void) -> ResultCode;
@@ -64,13 +89,14 @@ pub type VtabFnEof = unsafe extern "C" fn(cursor: *const c_void) -> bool;
pub type VtabRowIDFn = unsafe extern "C" fn(cursor: *const c_void) -> i64;
pub type VtabFnUpdate = unsafe extern "C" fn(
vtab: *const c_void,
table: *const c_void,
argc: i32,
argv: *const Value,
p_out_rowid: *mut i64,
) -> ResultCode;
pub type VtabFnDestroy = unsafe extern "C" fn(vtab: *const c_void) -> ResultCode;
pub type VtabFnDestroy = unsafe extern "C" fn(table: *const c_void) -> ResultCode;
pub type BestIdxFn = unsafe extern "C" fn(
constraints: *const ConstraintInfo,
constraint_len: i32,
@@ -86,18 +112,20 @@ pub enum VTabKind {
}
pub trait VTabModule: 'static {
type VCursor: VTabCursor<Error = Self::Error>;
type Table: VTable;
const VTAB_KIND: VTabKind;
const NAME: &'static str;
/// Creates a new instance of a virtual table.
/// Returns a tuple where the first element is the table's schema.
fn create(args: &[Value]) -> Result<(String, Self::Table), ResultCode>;
}
pub trait VTable {
type Cursor: VTabCursor<Error = Self::Error>;
type Error: std::fmt::Display;
fn create_schema(args: &[Value]) -> String;
fn open(&self) -> Result<Self::VCursor, Self::Error>;
fn filter(
cursor: &mut Self::VCursor,
args: &[Value],
idx_info: Option<(&str, i32)>,
) -> ResultCode;
fn open(&self) -> Result<Self::Cursor, Self::Error>;
fn update(&mut self, _rowid: i64, _args: &[Value]) -> Result<(), Self::Error> {
Ok(())
}
@@ -130,6 +158,7 @@ pub trait VTabModule: 'static {
pub trait VTabCursor: Sized {
type Error: std::fmt::Display;
fn filter(&mut self, args: &[Value], idx_info: Option<(&str, i32)>) -> ResultCode;
fn rowid(&self) -> i64;
fn column(&self, idx: u32) -> Result<Value, Self::Error>;
fn eof(&self) -> bool;

107
extensions/series/src/lib.rs
View File

@@ -1,9 +1,10 @@
use limbo_ext::{
register_extension, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, Value,
register_extension, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, VTable,
Value,
};
register_extension! {
vtabs: { GenerateSeriesVTab }
vtabs: { GenerateSeriesVTabModule }
}
macro_rules! try_option {
@@ -17,26 +18,32 @@ macro_rules! try_option {
/// A virtual table that generates a sequence of integers
#[derive(Debug, VTabModuleDerive, Default)]
struct GenerateSeriesVTab;
struct GenerateSeriesVTabModule;
impl VTabModule for GenerateSeriesVTab {
type VCursor = GenerateSeriesCursor;
type Error = ResultCode;
impl VTabModule for GenerateSeriesVTabModule {
type Table = GenerateSeriesTable;
const NAME: &'static str = "generate_series";
const VTAB_KIND: VTabKind = VTabKind::TableValuedFunction;
fn create_schema(_args: &[Value]) -> String {
// Create table schema
"CREATE TABLE generate_series(
fn create(_args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
let schema = "CREATE TABLE generate_series(
value INTEGER,
start INTEGER HIDDEN,
stop INTEGER HIDDEN,
step INTEGER HIDDEN
)"
.into()
.into();
Ok((schema, GenerateSeriesTable {}))
}
}
fn open(&self) -> Result<Self::VCursor, Self::Error> {
struct GenerateSeriesTable {}
impl VTable for GenerateSeriesTable {
type Cursor = GenerateSeriesCursor;
type Error = ResultCode;
fn open(&self) -> Result<Self::Cursor, Self::Error> {
Ok(GenerateSeriesCursor {
start: 0,
stop: 0,
@@ -44,41 +51,6 @@ impl VTabModule for GenerateSeriesVTab {
current: 0,
})
}
fn filter(cursor: &mut Self::VCursor, args: &[Value], _: Option<(&str, i32)>) -> ResultCode {
// args are the start, stop, and step
if args.is_empty() || args.len() > 3 {
return ResultCode::InvalidArgs;
}
let start = try_option!(args[0].to_integer(), ResultCode::InvalidArgs);
let stop = try_option!(
args.get(1).map(|v| v.to_integer().unwrap_or(i64::MAX)),
ResultCode::EOF // Sqlite returns an empty series for wacky args
);
let mut step = args
.get(2)
.map(|v| v.to_integer().unwrap_or(1))
.unwrap_or(1);
// Convert zero step to 1, matching SQLite behavior
if step == 0 {
step = 1;
}
cursor.start = start;
cursor.step = step;
cursor.stop = stop;
// Set initial value based on range validity
// For invalid input SQLite returns an empty series
cursor.current = if cursor.is_invalid_range() {
return ResultCode::EOF;
} else {
start
};
ResultCode::OK
}
}
/// The cursor for iterating over the generated sequence
@@ -116,6 +88,41 @@ impl GenerateSeriesCursor {
impl VTabCursor for GenerateSeriesCursor {
type Error = ResultCode;
fn filter(&mut self, args: &[Value], _: Option<(&str, i32)>) -> ResultCode {
// args are the start, stop, and step
if args.is_empty() || args.len() > 3 {
return ResultCode::InvalidArgs;
}
let start = try_option!(args[0].to_integer(), ResultCode::InvalidArgs);
let stop = try_option!(
args.get(1).map(|v| v.to_integer().unwrap_or(i64::MAX)),
ResultCode::EOF // Sqlite returns an empty series for wacky args
);
let mut step = args
.get(2)
.map(|v| v.to_integer().unwrap_or(1))
.unwrap_or(1);
// Convert zero step to 1, matching SQLite behavior
if step == 0 {
step = 1;
}
self.start = start;
self.step = step;
self.stop = stop;
// Set initial value based on range validity
// For invalid input SQLite returns an empty series
self.current = if self.is_invalid_range() {
return ResultCode::EOF;
} else {
start
};
ResultCode::OK
}
fn next(&mut self) -> ResultCode {
if self.eof() {
return ResultCode::EOF;
@@ -217,7 +224,7 @@ mod tests {
}
// Helper function to collect all values from a cursor, returns Result with error code
fn collect_series(series: Series) -> Result<Vec<i64>, ResultCode> {
let tbl = GenerateSeriesVTab;
let tbl = GenerateSeriesTable {};
let mut cursor = tbl.open()?;
// Create args array for filter
@@ -228,7 +235,7 @@ mod tests {
];
// Initialize cursor through filter
match GenerateSeriesVTab::filter(&mut cursor, &args, None) {
match cursor.filter(&args, None) {
ResultCode::OK => (),
ResultCode::EOF => return Ok(vec![]),
err => return Err(err),
@@ -534,7 +541,7 @@ mod tests {
let start = series.start;
let stop = series.stop;
let step = series.step;
let tbl = GenerateSeriesVTab {};
let tbl = GenerateSeriesTable {};
let mut cursor = tbl.open().unwrap();
let args = vec![
@@ -544,7 +551,7 @@ mod tests {
];
// Initialize cursor through filter
GenerateSeriesVTab::filter(&mut cursor, &args, None);
cursor.filter(&args, None);
let mut rowids = vec![];
while !cursor.eof() {

208
extensions/tests/src/lib.rs
View File

@@ -1,7 +1,8 @@
use lazy_static::lazy_static;
use limbo_ext::{
register_extension, scalar, ConstraintInfo, ConstraintOp, ConstraintUsage, ExtResult,
IndexInfo, OrderByInfo, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, Value,
IndexInfo, OrderByInfo, ResultCode, VTabCursor, VTabKind, VTabModule, VTabModuleDerive, VTable,
Value,
};
#[cfg(not(target_family = "wasm"))]
use limbo_ext::{VfsDerive, VfsExtension, VfsFile};
@@ -11,7 +12,7 @@ use std::io::{Read, Seek, SeekFrom, Write};
use std::sync::Mutex;
register_extension! {
vtabs: { KVStoreVTab },
vtabs: { KVStoreVTabModule },
scalars: { test_scalar },
vfs: { TestFS },
}
@@ -21,7 +22,7 @@ lazy_static! {
}
#[derive(VTabModuleDerive, Default)]
pub struct KVStoreVTab;
pub struct KVStoreVTabModule;
/// the cursor holds a snapshot of (rowid, key, value) in memory.
pub struct KVStoreCursor {
@@ -29,17 +30,114 @@ pub struct KVStoreCursor {
index: Option<usize>,
}
impl VTabModule for KVStoreVTab {
type VCursor = KVStoreCursor;
impl VTabModule for KVStoreVTabModule {
type Table = KVStoreTable;
const VTAB_KIND: VTabKind = VTabKind::VirtualTable;
const NAME: &'static str = "kv_store";
fn create(_args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
let schema = "CREATE TABLE x (key TEXT PRIMARY KEY, value TEXT);".to_string();
Ok((schema, KVStoreTable {}))
}
}
fn hash_key(key: &str) -> i64 {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
key.hash(&mut hasher);
hasher.finish() as i64
}
impl VTabCursor for KVStoreCursor {
type Error = String;
fn create_schema(_args: &[Value]) -> String {
"CREATE TABLE x (key TEXT PRIMARY KEY, value TEXT);".to_string()
fn filter(&mut self, args: &[Value], idx_str: Option<(&str, i32)>) -> ResultCode {
match idx_str {
Some(("key_eq", 1)) => {
let key = args
.first()
.and_then(|v| v.to_text())
.map(|s| s.to_string());
log::debug!("idx_str found: key_eq\n value: {:?}", key);
if let Some(key) = key {
let rowid = hash_key(&key);
let store = GLOBAL_STORE.lock().unwrap();
if let Some((k, v)) = store.get(&rowid) {
self.rows.push((rowid, k.clone(), v.clone()));
self.index = Some(0);
} else {
self.rows.clear();
self.index = None;
return ResultCode::EOF;
}
return ResultCode::OK;
}
self.rows.clear();
self.index = None;
ResultCode::OK
}
_ => {
let store = GLOBAL_STORE.lock().unwrap();
self.rows = store
.iter()
.map(|(&rowid, (k, v))| (rowid, k.clone(), v.clone()))
.collect();
self.rows.sort_by_key(|(rowid, _, _)| *rowid);
if self.rows.is_empty() {
self.index = None;
ResultCode::EOF
} else {
self.index = Some(0);
ResultCode::OK
}
}
}
}
fn open(&self) -> Result<Self::VCursor, Self::Error> {
fn rowid(&self) -> i64 {
if self.index.is_some_and(|c| c < self.rows.len()) {
self.rows[self.index.unwrap_or(0)].0
} else {
log::error!("rowid: -1");
-1
}
}
fn column(&self, idx: u32) -> Result<Value, Self::Error> {
if self.index.is_some_and(|c| c >= self.rows.len()) {
return Err("cursor out of range".into());
}
if let Some((_, ref key, ref val)) = self.rows.get(self.index.unwrap_or(0)) {
match idx {
0 => Ok(Value::from_text(key.clone())), // key
1 => Ok(Value::from_text(val.clone())), // value
_ => Err("Invalid column".into()),
}
} else {
Err("Invalid Column".into())
}
}
fn eof(&self) -> bool {
self.index.is_some_and(|s| s >= self.rows.len()) || self.index.is_none()
}
fn next(&mut self) -> ResultCode {
self.index = Some(self.index.unwrap_or(0) + 1);
if self.index.is_some_and(|c| c >= self.rows.len()) {
return ResultCode::EOF;
}
ResultCode::OK
}
}
pub struct KVStoreTable {}
impl VTable for KVStoreTable {
type Cursor = KVStoreCursor;
type Error = String;
fn open(&self) -> Result<Self::Cursor, Self::Error> {
let _ = env_logger::try_init();
Ok(KVStoreCursor {
rows: Vec::new(),
@@ -88,53 +186,6 @@ impl VTabModule for KVStoreVTab {
}
}
fn filter(
cursor: &mut Self::VCursor,
args: &[Value],
idx_str: Option<(&str, i32)>,
) -> ResultCode {
match idx_str {
Some(("key_eq", 1)) => {
let key = args
.first()
.and_then(|v| v.to_text())
.map(|s| s.to_string());
log::debug!("idx_str found: key_eq\n value: {:?}", key);
if let Some(key) = key {
let rowid = hash_key(&key);
let store = GLOBAL_STORE.lock().unwrap();
if let Some((k, v)) = store.get(&rowid) {
cursor.rows.push((rowid, k.clone(), v.clone()));
cursor.index = Some(0);
} else {
cursor.rows.clear();
cursor.index = None;
return ResultCode::EOF;
}
return ResultCode::OK;
}
cursor.rows.clear();
cursor.index = None;
ResultCode::OK
}
_ => {
let store = GLOBAL_STORE.lock().unwrap();
cursor.rows = store
.iter()
.map(|(&rowid, (k, v))| (rowid, k.clone(), v.clone()))
.collect();
cursor.rows.sort_by_key(|(rowid, _, _)| *rowid);
if cursor.rows.is_empty() {
cursor.index = None;
ResultCode::EOF
} else {
cursor.index = Some(0);
ResultCode::OK
}
}
}
}
fn insert(&mut self, values: &[Value]) -> Result<i64, Self::Error> {
let key = values
.first()
@@ -175,53 +226,6 @@ impl VTabModule for KVStoreVTab {
}
}
fn hash_key(key: &str) -> i64 {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
key.hash(&mut hasher);
hasher.finish() as i64
}
impl VTabCursor for KVStoreCursor {
type Error = String;
fn rowid(&self) -> i64 {
if self.index.is_some_and(|c| c < self.rows.len()) {
self.rows[self.index.unwrap_or(0)].0
} else {
log::error!("rowid: -1");
-1
}
}
fn column(&self, idx: u32) -> Result<Value, Self::Error> {
if self.index.is_some_and(|c| c >= self.rows.len()) {
return Err("cursor out of range".into());
}
if let Some((_, ref key, ref val)) = self.rows.get(self.index.unwrap_or(0)) {
match idx {
0 => Ok(Value::from_text(key.clone())), // key
1 => Ok(Value::from_text(val.clone())), // value
_ => Err("Invalid column".into()),
}
} else {
Err("Invalid Column".into())
}
}
fn eof(&self) -> bool {
self.index.is_some_and(|s| s >= self.rows.len()) || self.index.is_none()
}
fn next(&mut self) -> ResultCode {
self.index = Some(self.index.unwrap_or(0) + 1);
if self.index.is_some_and(|c| c >= self.rows.len()) {
return ResultCode::EOF;
}
ResultCode::OK
}
}
pub struct TestFile {
file: File,
}

91
macros/src/ext/vtab_derive.rs
View File

@@ -7,7 +7,7 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
let struct_name = &ast.ident;
let register_fn_name = format_ident!("register_{}", struct_name);
let create_schema_fn_name = format_ident!("create_schema_{}", struct_name);
let create_fn_name = format_ident!("create_{}", struct_name);
let open_fn_name = format_ident!("open_{}", struct_name);
let close_fn_name = format_ident!("close_{}", struct_name);
let filter_fn_name = format_ident!("filter_{}", struct_name);
@@ -22,26 +22,40 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
let expanded = quote! {
impl #struct_name {
#[no_mangle]
unsafe extern "C" fn #create_schema_fn_name(
unsafe extern "C" fn #create_fn_name(
argv: *const ::limbo_ext::Value, argc: i32
) -> *mut ::std::ffi::c_char {
) -> ::limbo_ext::VTabCreateResult {
let args = if argv.is_null() {
&Vec::new()
} else {
::std::slice::from_raw_parts(argv, argc as usize)
};
let sql = <#struct_name as ::limbo_ext::VTabModule>::create_schema(&args);
::std::ffi::CString::new(sql).unwrap().into_raw()
match <#struct_name as ::limbo_ext::VTabModule>::create(&args) {
Ok((schema, table)) => {
::limbo_ext::VTabCreateResult {
code: ::limbo_ext::ResultCode::OK,
schema: ::std::ffi::CString::new(schema).unwrap().into_raw(),
table: ::std::boxed::Box::into_raw(::std::boxed::Box::new(table)) as *const ::std::ffi::c_void,
}
},
Err(e) => {
::limbo_ext::VTabCreateResult {
code: e,
schema: ::std::ptr::null(),
table: ::std::ptr::null(),
}
}
}
}
#[no_mangle]
unsafe extern "C" fn #open_fn_name(ctx: *const ::std::ffi::c_void) -> *const ::std::ffi::c_void {
if ctx.is_null() {
unsafe extern "C" fn #open_fn_name(table: *const ::std::ffi::c_void) -> *const ::std::ffi::c_void {
if table.is_null() {
return ::std::ptr::null();
}
let ctx = ctx as *const #struct_name;
let ctx: &#struct_name = &*ctx;
if let Ok(cursor) = <#struct_name as ::limbo_ext::VTabModule>::open(ctx) {
let table = table as *const <#struct_name as ::limbo_ext::VTabModule>::Table;
let table: &<#struct_name as ::limbo_ext::VTabModule>::Table = &*table;
if let Ok(cursor) = <#struct_name as ::limbo_ext::VTabModule>::Table::open(table) {
return ::std::boxed::Box::into_raw(::std::boxed::Box::new(cursor)) as *const ::std::ffi::c_void;
} else {
return ::std::ptr::null();
@@ -55,8 +69,9 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
if cursor.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let boxed_cursor = ::std::boxed::Box::from_raw(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor);
boxed_cursor.close()
let cursor = cursor as *mut <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor;
let cursor = ::std::boxed::Box::from_raw(cursor);
<<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::close(&*cursor)
}
#[no_mangle]
@@ -70,14 +85,14 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
if cursor.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let cursor = unsafe { &mut *(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor) };
let cursor = &mut *(cursor as *mut <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor);
let args = ::std::slice::from_raw_parts(argv, argc as usize);
let idx_str = if idx_str.is_null() {
None
} else {
Some((unsafe { ::std::ffi::CStr::from_ptr(idx_str).to_str().unwrap() }, idx_num))
};
<#struct_name as ::limbo_ext::VTabModule>::filter(cursor, args, idx_str)
<<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::filter(cursor, args, idx_str)
}
#[no_mangle]
@@ -88,8 +103,8 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
if cursor.is_null() {
return ::limbo_ext::Value::error(::limbo_ext::ResultCode::Error);
}
let cursor = unsafe { &mut *(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor) };
match <<#struct_name as ::limbo_ext::VTabModule>::VCursor as ::limbo_ext::VTabCursor>::column(cursor, idx) {
let cursor = &*(cursor as *const <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor);
match <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::column(cursor, idx) {
Ok(val) => val,
Err(e) => ::limbo_ext::Value::error_with_message(e.to_string())
}
@@ -102,8 +117,8 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
if cursor.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let cursor = &mut *(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor);
<<#struct_name as ::limbo_ext::VTabModule>::VCursor as ::limbo_ext::VTabCursor>::next(cursor)
let cursor = &mut *(cursor as *mut <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor);
<<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::next(cursor)
}
#[no_mangle]
@@ -113,22 +128,22 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
if cursor.is_null() {
return true;
}
let cursor = &mut *(cursor as *mut <#struct_name as ::limbo_ext::VTabModule>::VCursor);
<<#struct_name as ::limbo_ext::VTabModule>::VCursor as ::limbo_ext::VTabCursor>::eof(cursor)
let cursor = &*(cursor as *const <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor);
<<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::eof(cursor)
}
#[no_mangle]
unsafe extern "C" fn #update_fn_name(
vtab: *const ::std::ffi::c_void,
table: *const ::std::ffi::c_void,
argc: i32,
argv: *const ::limbo_ext::Value,
p_out_rowid: *mut i64,
) -> ::limbo_ext::ResultCode {
if vtab.is_null() {
if table.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let vtab = &mut *(vtab as *mut #struct_name);
let table = &mut *(table as *mut <#struct_name as ::limbo_ext::VTabModule>::Table);
let args = ::std::slice::from_raw_parts(argv, argc as usize);
let old_rowid = match args.get(0).map(|v| v.value_type()) {
@@ -143,21 +158,21 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
match (old_rowid, new_rowid) {
// DELETE: old_rowid provided, no new_rowid
(Some(old), None) => {
if <#struct_name as VTabModule>::delete(vtab, old).is_err() {
if <#struct_name as VTabModule>::Table::delete(table, old).is_err() {
return ::limbo_ext::ResultCode::Error;
}
return ::limbo_ext::ResultCode::OK;
}
// UPDATE: old_rowid provided and new_rowid may exist
(Some(old), Some(new)) => {
if <#struct_name as VTabModule>::update(vtab, old, &columns).is_err() {
if <#struct_name as VTabModule>::Table::update(table, old, &columns).is_err() {
return ::limbo_ext::ResultCode::Error;
}
return ::limbo_ext::ResultCode::OK;
}
// INSERT: no old_rowid (old_rowid = None)
(None, _) => {
if let Ok(rowid) = <#struct_name as VTabModule>::insert(vtab, &columns) {
if let Ok(rowid) = <#struct_name as VTabModule>::Table::insert(table, &columns) {
if !p_out_rowid.is_null() {
*p_out_rowid = rowid;
return ::limbo_ext::ResultCode::RowID;
@@ -170,24 +185,26 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
}
#[no_mangle]
pub unsafe extern "C" fn #rowid_fn_name(ctx: *const ::std::ffi::c_void) -> i64 {
if ctx.is_null() {
pub unsafe extern "C" fn #rowid_fn_name(cursor: *const ::std::ffi::c_void) -> i64 {
if cursor.is_null() {
return -1;
}
let cursor = &*(ctx as *const <#struct_name as ::limbo_ext::VTabModule>::VCursor);
<<#struct_name as ::limbo_ext::VTabModule>::VCursor as ::limbo_ext::VTabCursor>::rowid(cursor)
let cursor = &*(cursor as *const <<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor);
<<#struct_name as ::limbo_ext::VTabModule>::Table as ::limbo_ext::VTable>::Cursor::rowid(cursor)
}
#[no_mangle]
unsafe extern "C" fn #destroy_fn_name(
vtab: *const ::std::ffi::c_void,
table: *const ::std::ffi::c_void,
) -> ::limbo_ext::ResultCode {
if vtab.is_null() {
if table.is_null() {
return ::limbo_ext::ResultCode::Error;
}
let vtab = &mut *(vtab as *mut #struct_name);
if <#struct_name as VTabModule>::destroy(vtab).is_err() {
// Take ownership of the table so it can be properly dropped.
let mut table: ::std::boxed::Box<<#struct_name as ::limbo_ext::VTabModule>::Table> =
::std::boxed::Box::from_raw(table as *mut <#struct_name as ::limbo_ext::VTabModule>::Table);
if <#struct_name as VTabModule>::Table::destroy(&mut *table).is_err() {
return ::limbo_ext::ResultCode::Error;
}
@@ -203,7 +220,7 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
) -> ::limbo_ext::ExtIndexInfo {
let constraints = if n_constraints > 0 { std::slice::from_raw_parts(constraints, n_constraints as usize) } else { &[] };
let order_by = if n_order_by > 0 { std::slice::from_raw_parts(order_by, n_order_by as usize) } else { &[] };
<#struct_name as ::limbo_ext::VTabModule>::best_index(constraints, order_by).to_ffi()
<#struct_name as ::limbo_ext::VTabModule>::Table::best_index(constraints, order_by).to_ffi()
}
#[no_mangle]
@@ -216,11 +233,9 @@ pub fn derive_vtab_module(input: TokenStream) -> TokenStream {
let api = &*api;
let name = <#struct_name as ::limbo_ext::VTabModule>::NAME;
let name_c = ::std::ffi::CString::new(name).unwrap().into_raw() as *const ::std::ffi::c_char;
let table_instance = ::std::boxed::Box::into_raw(::std::boxed::Box::new(#struct_name::default()));
let module = ::limbo_ext::VTabModuleImpl {
ctx: table_instance as *const ::std::ffi::c_void,
name: name_c,
create_schema: Self::#create_schema_fn_name,
create: Self::#create_fn_name,
open: Self::#open_fn_name,
close: Self::#close_fn_name,
filter: Self::#filter_fn_name,

111
macros/src/lib.rs
View File

@@ -230,63 +230,77 @@ pub fn derive_agg_func(input: TokenStream) -> TokenStream {
/// Macro to derive a VTabModule for your extension. This macro will generate
/// the necessary functions to register your module with core. You must implement
/// the VTabModule trait for your struct, and the VTabCursor trait for your cursor.
/// the VTabModule, VTable, and VTabCursor traits.
/// ```ignore
///#[derive(Debug, VTabModuleDerive)]
///struct CsvVTab;
///impl VTabModule for CsvVTab {
/// type VCursor = CsvCursor;
/// const NAME: &'static str = "csv_data";
/// #[derive(Debug, VTabModuleDerive)]
/// struct CsvVTabModule;
///
/// impl VTabModule for CsvVTabModule {
/// type Table = CsvTable;
/// const NAME: &'static str = "csv_data";
/// const VTAB_KIND: VTabKind = VTabKind::VirtualTable;
///
/// /// Declare your virtual table and its schema
/// fn create(args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
/// let schema = "CREATE TABLE csv_data(
/// name TEXT,
/// age TEXT,
/// city TEXT
/// )".into();
/// Ok((schema, CsvTable {}))
/// }
/// }
///
/// struct CsvTable {}
///
/// // Implement the VTable trait for your virtual table
/// impl VTable for CsvTable {
/// type Cursor = CsvCursor;
/// type Error = &'static str;
///
/// /// Open the virtual table and return a cursor
/// fn open(&self) -> Result<Self::Cursor, Self::Error> {
/// let csv_content = fs::read_to_string("data.csv").unwrap_or_default();
/// let rows: Vec<Vec<String>> = csv_content
/// .lines()
/// .skip(1)
/// .map(|line| {
/// line.split(',')
/// .map(|s| s.trim().to_string())
/// .collect()
/// })
/// .collect();
/// Ok(CsvCursor { rows, index: 0 })
/// }
///
/// /// Declare the schema for your virtual table
/// fn create_schema(args: &[&str]) -> &'static str {
/// let sql = "CREATE TABLE csv_data(
/// name TEXT,
/// age TEXT,
/// city TEXT
/// )"
/// }
/// /// Open the virtual table and return a cursor
/// fn open() -> Self::VCursor {
/// let csv_content = fs::read_to_string("data.csv").unwrap_or_default();
/// let rows: Vec<Vec<String>> = csv_content
/// .lines()
/// .skip(1)
/// .map(|line| {
/// line.split(',')
/// .map(|s| s.trim().to_string())
/// .collect()
/// })
/// .collect();
/// CsvCursor { rows, index: 0 }
/// }
/// /// Filter the virtual table based on arguments (omitted here for simplicity)
/// fn filter(_cursor: &mut Self::VCursor, _arg_count: i32, _args: &[Value]) -> ResultCode {
/// ResultCode::OK
/// }
/// /// **Optional** methods for non-readonly tables:
///
/// /// Update the row with the provided values, return the new rowid
/// fn update(&mut self, rowid: i64, args: &[Value]) -> Result<Option<i64>, Self::Error> {
/// Ok(None)// return Ok(None) for read-only
/// }
///
/// /// Insert a new row with the provided values, return the new rowid
/// fn insert(&mut self, args: &[Value]) -> Result<(), Self::Error> {
/// Ok(()) //
/// }
///
/// /// Delete the row with the provided rowid
/// fn delete(&mut self, rowid: i64) -> Result<(), Self::Error> {
/// Ok(())
/// }
///
/// /// Destroy the virtual table. Any cleanup logic for when the table is deleted comes heres
/// fn destroy(&mut self) -> Result<(), Self::Error> {
/// Ok(())
/// }
/// }
///
/// #[derive(Debug)]
/// struct CsvCursor {
/// rows: Vec<Vec<String>>,
/// index: usize,
/// }
///
/// impl CsvCursor {
/// /// Returns the value for a given column index.
@@ -298,27 +312,40 @@ pub fn derive_agg_func(input: TokenStream) -> TokenStream {
/// Value::null()
/// }
/// }
/// }
///
/// // Implement the VTabCursor trait for your virtual cursor
/// impl VTabCursor for CsvCursor {
/// /// Move the cursor to the next row
/// fn next(&mut self) -> ResultCode {
/// if self.index < self.rows.len() - 1 {
/// self.index += 1;
/// ResultCode::OK
/// } else {
/// ResultCode::EOF
/// }
/// }
/// type Error = &'static str;
///
/// /// Filter the virtual table based on arguments (omitted here for simplicity)
/// fn filter(&mut self, _args: &[Value], _idx_info: Option<(&str, i32)>) -> ResultCode {
/// ResultCode::OK
/// }
///
/// /// Move the cursor to the next row
/// fn next(&mut self) -> ResultCode {
/// if self.index < self.rows.len() - 1 {
/// self.index += 1;
/// ResultCode::OK
/// } else {
/// ResultCode::EOF
/// }
/// }
///
/// fn eof(&self) -> bool {
/// self.index >= self.rows.len()
/// }
///
/// /// Return the value for a given column index
/// fn column(&self, idx: u32) -> Value {
/// fn column(&self, idx: u32) -> Result<Value, Self::Error> {
/// self.column(idx)
/// }
///
/// fn rowid(&self) -> i64 {
/// self.index as i64
/// }
/// }
///
#[proc_macro_derive(VTabModuleDerive)]
pub fn derive_vtab_module(input: TokenStream) -> TokenStream {