diff --git a/core/ext/mod.rs b/core/ext/mod.rs index b6fe6fc44..d89c164e4 100644 --- a/core/ext/mod.rs +++ b/core/ext/mod.rs @@ -1,5 +1,6 @@ #[cfg(feature = "fs")] mod dynamic; +mod vtab_xconnect; #[cfg(all(target_os = "linux", feature = "io_uring"))] use crate::UringIO; use crate::{function::ExternalFunc, Connection, Database, LimboError, IO}; @@ -14,6 +15,7 @@ use std::{ rc::Rc, sync::Arc, }; +pub use vtab_xconnect::{close, prepare_stmt}; type ExternAggFunc = (InitAggFunction, StepFunction, FinalizeFunction); #[derive(Clone)] diff --git a/core/lib.rs b/core/lib.rs index 9d7858c39..4f901bd4e 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -46,6 +46,7 @@ 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::rc::Weak; use std::{ borrow::Cow, cell::{Cell, RefCell, UnsafeCell}, @@ -772,6 +773,7 @@ pub struct VirtualTable { columns: Vec, kind: VTabKind, table_ptr: *const c_void, + connection_ptr: RefCell<*mut limbo_ext::Conn>, } impl VirtualTable { @@ -817,7 +819,7 @@ impl VirtualTable { ))); } }; - let (schema, table_ptr) = module.implementation.as_ref().create(args)?; + 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()), @@ -825,6 +827,7 @@ impl VirtualTable { let columns = columns_from_create_table_body(&body)?; let vtab = Rc::new(VirtualTable { name: tbl_name.unwrap_or(module_name).to_owned(), + connection_ptr: RefCell::new(std::ptr::null_mut()), implementation: module.implementation.clone(), columns, args: exprs, @@ -838,8 +841,14 @@ impl VirtualTable { )) } - pub fn open(&self) -> crate::Result { - let cursor = unsafe { (self.implementation.open)(self.table_ptr) }; + pub fn open(&self, conn: Weak) -> crate::Result { + let conn = conn.as_ptr() as *mut c_void; + // prepare a pointer to the Weak and function pointers to box up + // and send to the extension. + let conn = limbo_ext::Conn::new(conn, crate::ext::prepare_stmt, crate::ext::close); + let conn = Box::into_raw(Box::new(conn)) as *mut limbo_ext::Conn; + *self.connection_ptr.borrow_mut() = conn; + let cursor = unsafe { (self.implementation.open)(self.table_ptr, conn) }; VTabOpaqueCursor::new(cursor, self.implementation.close) } diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 5c343d96b..00bf09ec7 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -993,7 +993,7 @@ pub fn op_vopen( let CursorType::VirtualTable(virtual_table) = cursor_type else { panic!("VOpen on non-virtual table cursor"); }; - let cursor = virtual_table.open()?; + let cursor = virtual_table.open(program.connection.clone())?; state .cursors .borrow_mut() diff --git a/extensions/core/src/lib.rs b/extensions/core/src/lib.rs index 75ea1c77d..ff52e90ba 100644 --- a/extensions/core/src/lib.rs +++ b/extensions/core/src/lib.rs @@ -16,8 +16,9 @@ pub use types::{ResultCode, Value, ValueType}; pub use vfs_modules::{RegisterVfsFn, VfsExtension, VfsFile, VfsFileImpl, VfsImpl, VfsInterface}; use vtabs::RegisterModuleFn; pub use vtabs::{ - ConstraintInfo, ConstraintOp, ConstraintUsage, ExtIndexInfo, IndexInfo, OrderByInfo, - VTabCreateResult, VTabCursor, VTabKind, VTabModule, VTabModuleImpl, VTable, + Conn, Connection, ConstraintInfo, ConstraintOp, ConstraintUsage, ExtIndexInfo, IndexInfo, + OrderByInfo, Statement, Stmt, VTabCreateResult, VTabCursor, VTabKind, VTabModule, + VTabModuleImpl, VTable, }; pub type ExtResult = std::result::Result;