diff --git a/core/translate/planner.rs b/core/translate/planner.rs index 0863b1c68..aa21e8a34 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -14,7 +14,7 @@ use crate::{ function::Func, schema::{Schema, Table}, translate::expr::walk_expr_mut, - util::{exprs_are_equivalent, normalize_ident, vtable_args}, + util::{exprs_are_equivalent, normalize_ident}, vdbe::{builder::TableRefIdCounter, BranchOffset}, Result, }; @@ -344,18 +344,7 @@ fn parse_from_clause_table<'a>( } ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => { let normalized_name = &normalize_ident(qualified_name.name.0.as_str()); - let args = match maybe_args { - Some(ref args) => vtable_args(args), - None => vec![], - }; - let vtab = crate::VirtualTable::from_args( - None, - normalized_name, - args, - syms, - limbo_ext::VTabKind::TableValuedFunction, - maybe_args, - )?; + let vtab = crate::VirtualTable::function(normalized_name, maybe_args, syms)?; let alias = maybe_alias .as_ref() .map(|a| match a { diff --git a/core/util.rs b/core/util.rs index f6b83ebfd..33af5dbc6 100644 --- a/core/util.rs +++ b/core/util.rs @@ -82,27 +82,12 @@ pub fn parse_schema_rows( vtab.clone() } else { let mod_name = module_name_from_sql(sql)?; - if let Some(vmod) = syms.vtab_modules.get(mod_name) { - if let limbo_ext::VTabKind::VirtualTable = vmod.module_kind - { - crate::VirtualTable::from_args( - Some(name), - mod_name, - module_args_from_sql(sql)?, - syms, - vmod.module_kind, - None, - )? - } else { - return Err(LimboError::Corrupt("Table valued function: {name} registered as virtual table in schema".to_string())); - } - } else { - // the extension isn't loaded, so we emit a warning. - return Err(LimboError::ExtensionError(format!( - "Virtual table module '{}' not found\nPlease load extension", - &mod_name - ))); - } + crate::VirtualTable::table( + Some(name), + mod_name, + module_args_from_sql(sql)?, + syms, + )? }; schema.add_virtual_table(vtab); } else { diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 572e5a971..cf42e68ba 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -1040,23 +1040,8 @@ pub fn op_vcreate( "Failed to upgrade Connection".to_string(), )); }; - let mod_type = conn - .syms - .borrow() - .vtab_modules - .get(&module_name) - .ok_or_else(|| { - crate::LimboError::ExtensionError(format!("Module {} not found", module_name)) - })? - .module_kind; - let table = crate::VirtualTable::from_args( - Some(&table_name), - &module_name, - args, - &conn.syms.borrow(), - mod_type, - None, - )?; + let table = + crate::VirtualTable::table(Some(&table_name), &module_name, args, &conn.syms.borrow())?; { conn.syms .borrow_mut() @@ -1090,12 +1075,7 @@ pub fn op_vfilter( let cursor = cursor.as_virtual_mut(); let mut args = Vec::with_capacity(*arg_count); for i in 0..*arg_count { - args.push( - state.registers[args_reg + i] - .get_owned_value() - .clone() - .to_ffi(), - ); + args.push(state.registers[args_reg + i].get_owned_value().clone()); } let idx_str = if let Some(idx_str) = idx_str { Some(state.registers[*idx_str].get_owned_value().to_string()) diff --git a/core/vtab.rs b/core/vtab.rs index 30b67649a..115ce536d 100644 --- a/core/vtab.rs +++ b/core/vtab.rs @@ -1,5 +1,5 @@ use crate::schema::Column; -use crate::util::columns_from_create_table_body; +use crate::util::{columns_from_create_table_body, vtable_args}; use crate::{Connection, LimboError, SymbolTable, Value}; use fallible_iterator::FallibleIterator; use limbo_ext::{ConstraintInfo, IndexInfo, OrderByInfo, ResultCode, VTabKind, VTabModuleImpl}; @@ -8,18 +8,137 @@ use std::cell::RefCell; use std::ffi::c_void; use std::rc::{Rc, Weak}; +#[derive(Debug, Clone)] +enum VirtualTableType { + External(ExtVirtualTable), +} + #[derive(Clone, Debug)] pub struct VirtualTable { - pub name: String, - pub args: Option>, - pub implementation: Rc, - pub columns: Vec, - pub kind: VTabKind, + pub(crate) name: String, + pub(crate) args: Option>, + pub(crate) columns: Vec, + pub(crate) kind: VTabKind, + vtab_type: VirtualTableType, +} + +impl VirtualTable { + pub(crate) fn function( + name: &str, + args: Option>, + syms: &SymbolTable, + ) -> crate::Result> { + let ext_args = match args { + Some(ref args) => vtable_args(args), + None => vec![], + }; + let (vtab, columns) = + ExtVirtualTable::from_args(name, ext_args, syms, VTabKind::TableValuedFunction)?; + let vtab = VirtualTable { + name: name.to_owned(), + args, + columns, + kind: VTabKind::TableValuedFunction, + vtab_type: VirtualTableType::External(vtab), + }; + Ok(Rc::new(vtab)) + } + + pub fn table( + tbl_name: Option<&str>, + module_name: &str, + args: Vec, + syms: &SymbolTable, + ) -> crate::Result> { + let (table, columns) = + ExtVirtualTable::from_args(module_name, args, syms, VTabKind::VirtualTable)?; + let vtab = VirtualTable { + name: tbl_name.unwrap_or(module_name).to_owned(), + args: None, + columns, + kind: VTabKind::VirtualTable, + vtab_type: VirtualTableType::External(table), + }; + Ok(Rc::new(vtab)) + } + + pub(crate) fn open(&self, conn: Weak) -> crate::Result { + match &self.vtab_type { + VirtualTableType::External(table) => { + Ok(VirtualTableCursor::External(table.open(conn)?)) + } + } + } + + pub(crate) fn update(&self, args: &[Value]) -> crate::Result> { + match &self.vtab_type { + VirtualTableType::External(table) => table.update(args), + } + } + + pub(crate) fn destroy(&self) -> crate::Result<()> { + match &self.vtab_type { + VirtualTableType::External(table) => table.destroy(), + } + } + + pub(crate) fn best_index( + &self, + constraints: &[ConstraintInfo], + order_by: &[OrderByInfo], + ) -> IndexInfo { + match &self.vtab_type { + VirtualTableType::External(table) => table.best_index(constraints, order_by), + } + } +} + +pub enum VirtualTableCursor { + External(ExtVirtualTableCursor), +} + +impl VirtualTableCursor { + pub(crate) fn next(&mut self) -> crate::Result { + match self { + VirtualTableCursor::External(cursor) => cursor.next(), + } + } + + pub(crate) fn rowid(&self) -> i64 { + match self { + VirtualTableCursor::External(cursor) => cursor.rowid(), + } + } + + pub(crate) fn column(&self, column: usize) -> crate::Result { + match self { + VirtualTableCursor::External(cursor) => cursor.column(column), + } + } + + pub(crate) fn filter( + &mut self, + idx_num: i32, + idx_str: Option, + arg_count: usize, + args: Vec, + ) -> crate::Result { + match self { + VirtualTableCursor::External(cursor) => { + cursor.filter(idx_num, idx_str, arg_count, args) + } + } + } +} + +#[derive(Clone, Debug)] +struct ExtVirtualTable { + implementation: Rc, table_ptr: *const c_void, connection_ptr: RefCell>, } -impl Drop for VirtualTable { +impl Drop for ExtVirtualTable { fn drop(&mut self) { if let Some(conn) = self.connection_ptr.borrow_mut().take() { if conn.is_null() { @@ -34,12 +153,8 @@ impl Drop for VirtualTable { } } -impl VirtualTable { - pub(crate) fn best_index( - &self, - constraints: &[ConstraintInfo], - order_by: &[OrderByInfo], - ) -> IndexInfo { +impl ExtVirtualTable { + fn best_index(&self, constraints: &[ConstraintInfo], order_by: &[OrderByInfo]) -> IndexInfo { unsafe { IndexInfo::from_ffi((self.implementation.best_idx)( constraints.as_ptr(), @@ -51,14 +166,12 @@ impl VirtualTable { } /// takes ownership of the provided Args - pub(crate) fn from_args( - tbl_name: Option<&str>, + fn from_args( module_name: &str, args: Vec, syms: &SymbolTable, kind: VTabKind, - exprs: Option>, - ) -> crate::Result> { + ) -> crate::Result<(Self, Vec)> { let module = syms .vtab_modules .get(module_name) @@ -66,30 +179,28 @@ impl VirtualTable { "Virtual table module not found: {}", module_name )))?; - if let VTabKind::VirtualTable = kind { - if module.module_kind == VTabKind::TableValuedFunction { - return Err(LimboError::ExtensionError(format!( - "{} is not a virtual table module", - module_name - ))); - } - }; + if kind != module.module_kind { + let expected = match kind { + VTabKind::VirtualTable => "virtual table", + VTabKind::TableValuedFunction => "table-valued function", + }; + return Err(LimboError::ExtensionError(format!( + "{} is not a {} module", + module_name, expected + ))); + } let (schema, table_ptr) = module.implementation.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()), )? { let columns = columns_from_create_table_body(&body)?; - let vtab = Rc::new(VirtualTable { - name: tbl_name.unwrap_or(module_name).to_owned(), + let vtab = ExtVirtualTable { connection_ptr: RefCell::new(None), implementation: module.implementation.clone(), - columns, - args: exprs, - kind, table_ptr, - }); - return Ok(vtab); + }; + return Ok((vtab, columns)); } Err(LimboError::ParseError( "Failed to parse schema from virtual table module".to_string(), @@ -98,7 +209,7 @@ impl VirtualTable { /// Accepts a Weak pointer to the connection that owns the VTable, that the module /// can optionally use to query the other tables. - pub fn open(&self, conn: Weak) -> crate::Result { + fn open(&self, conn: Weak) -> crate::Result { // we need a Weak to upgrade and call from the extension. let weak_box: *mut Weak = Box::into_raw(Box::new(conn)); let conn = limbo_ext::Conn::new( @@ -111,10 +222,10 @@ impl VirtualTable { // store the leaked connection pointer on the table so it can be freed on drop *self.connection_ptr.borrow_mut() = Some(ext_conn_ptr); let cursor = unsafe { (self.implementation.open)(self.table_ptr, ext_conn_ptr) }; - VirtualTableCursor::new(cursor, self.implementation.clone()) + ExtVirtualTableCursor::new(cursor, self.implementation.clone()) } - pub fn update(&self, args: &[Value]) -> crate::Result> { + fn update(&self, args: &[Value]) -> crate::Result> { let arg_count = args.len(); let ext_args = args.iter().map(|arg| arg.to_ffi()).collect::>(); let newrowid = 0i64; @@ -138,7 +249,7 @@ impl VirtualTable { } } - pub fn destroy(&self) -> crate::Result<()> { + fn destroy(&self) -> crate::Result<()> { let rc = unsafe { (self.implementation.destroy)(self.table_ptr) }; match rc { ResultCode::OK => Ok(()), @@ -147,13 +258,13 @@ impl VirtualTable { } } -pub struct VirtualTableCursor { +pub struct ExtVirtualTableCursor { cursor: *const c_void, implementation: Rc, } -impl VirtualTableCursor { - pub fn new(cursor: *const c_void, implementation: Rc) -> crate::Result { +impl ExtVirtualTableCursor { + fn new(cursor: *const c_void, implementation: Rc) -> crate::Result { if cursor.is_null() { return Err(LimboError::InternalError( "VirtualTableCursor: cursor is null".into(), @@ -165,19 +276,20 @@ impl VirtualTableCursor { }) } - pub(crate) fn rowid(&self) -> i64 { + fn rowid(&self) -> i64 { unsafe { (self.implementation.rowid)(self.cursor) } } #[tracing::instrument(skip(self))] - pub fn filter( + fn filter( &self, idx_num: i32, idx_str: Option, arg_count: usize, - args: Vec, + args: Vec, ) -> crate::Result { tracing::trace!("xFilter"); + let ext_args = args.iter().map(|arg| arg.to_ffi()).collect::>(); let c_idx_str = idx_str .map(|s| std::ffi::CString::new(s).unwrap()) .map(|cstr| cstr.into_raw()) @@ -186,12 +298,12 @@ impl VirtualTableCursor { (self.implementation.filter)( self.cursor, arg_count as i32, - args.as_ptr(), + ext_args.as_ptr(), c_idx_str, idx_num, ) }; - for arg in args { + for arg in ext_args { unsafe { arg.__free_internal_type(); } @@ -203,12 +315,12 @@ impl VirtualTableCursor { } } - pub fn column(&self, column: usize) -> crate::Result { + fn column(&self, column: usize) -> crate::Result { let val = unsafe { (self.implementation.column)(self.cursor, column as u32) }; Value::from_ffi(val) } - pub fn next(&self) -> crate::Result { + fn next(&self) -> crate::Result { let rc = unsafe { (self.implementation.next)(self.cursor) }; match rc { ResultCode::OK => Ok(true), @@ -218,7 +330,7 @@ impl VirtualTableCursor { } } -impl Drop for VirtualTableCursor { +impl Drop for ExtVirtualTableCursor { fn drop(&mut self) { let result = unsafe { (self.implementation.close)(self.cursor) }; if !result.is_ok() {