From 5ebd3f7271ccaf20f7fde7426517edb64d2671fb Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Mon, 21 Jul 2025 19:08:27 -0400 Subject: [PATCH] Change api of extension api context to support static extensions held on db --- core/ext/mod.rs | 275 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 188 insertions(+), 87 deletions(-) diff --git a/core/ext/mod.rs b/core/ext/mod.rs index 89c3e1a61..da723e8dd 100644 --- a/core/ext/mod.rs +++ b/core/ext/mod.rs @@ -1,12 +1,14 @@ #[cfg(feature = "fs")] mod dynamic; mod vtab_xconnect; -use crate::vtab::VirtualTable; +use crate::schema::{Schema, Table}; #[cfg(all(target_os = "linux", feature = "io_uring"))] use crate::UringIO; use crate::{function::ExternalFunc, Connection, Database, LimboError, IO}; +use crate::{vtab::VirtualTable, SymbolTable}; #[cfg(feature = "fs")] pub use dynamic::{add_builtin_vfs_extensions, add_vfs_module, list_vfs_modules, VfsMod}; +use parking_lot::Mutex; use std::{ ffi::{c_char, c_void, CStr, CString}, rc::Rc, @@ -17,7 +19,113 @@ use turso_ext::{ }; pub use turso_ext::{FinalizeFunction, StepFunction, Value as ExtValue, ValueType as ExtValueType}; pub use vtab_xconnect::{close, execute, prepare_stmt}; -type ExternAggFunc = (InitAggFunction, StepFunction, FinalizeFunction); + +/// The context passed to extensions to register with Core +/// along with the function pointers +#[repr(C)] +pub struct ExtensionCtx { + syms: *mut SymbolTable, + schema_data: *mut c_void, + schema_handler: SchemaHandler, +} + +type SchemaHandler = unsafe extern "C" fn( + schema_data: *mut c_void, + table_name: *const c_char, + table: *mut c_void, +) -> ResultCode; + +/// Handler for our Connection that has direct Arc access +/// to register a table from an extension. +unsafe extern "C" fn handle_schema_insert_connection( + schema_data: *mut c_void, + table_name: *const c_char, + table: *mut c_void, +) -> ResultCode { + let schema = &mut *(schema_data as *mut Schema); + let c_str = CStr::from_ptr(table_name); + let name = match c_str.to_str() { + Ok(s) => s.to_string(), + Err(_) => return ResultCode::InvalidArgs, + }; + let table = Box::from_raw(table as *mut Arc); + schema.tables.insert(name, *table); + ResultCode::OK +} + +/// Handler for Database with Mutex> access to +/// register a table from an extension. +unsafe extern "C" fn handle_schema_insert_database( + schema_data: *mut c_void, + table_name: *const c_char, + table: *mut c_void, +) -> ResultCode { + let mutex = &*(schema_data as *mut Mutex>); + let mut guard = mutex.lock(); + let schema = Arc::make_mut(&mut *guard); + + let c_str = CStr::from_ptr(table_name); + let name = match c_str.to_str() { + Ok(s) => s.to_string(), + Err(_) => return ResultCode::InvalidArgs, + }; + let table = Box::from_raw(table as *mut Arc
); + schema.tables.insert(name, *table); + ResultCode::OK +} + +pub(crate) unsafe extern "C" fn register_vtab_module( + ctx: *mut c_void, + name: *const c_char, + module: VTabModuleImpl, + kind: VTabKind, +) -> ResultCode { + if name.is_null() || ctx.is_null() { + return ResultCode::Error; + } + + let c_str = unsafe { CString::from_raw(name as *mut c_char) }; + let name_str = match c_str.to_str() { + Ok(s) => s.to_string(), + Err(_) => return ResultCode::Error, + }; + + let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) }; + let module = Rc::new(module); + let vmodule = VTabImpl { + module_kind: kind, + implementation: module, + }; + + unsafe { + let syms = &mut *ext_ctx.syms; + syms.vtab_modules.insert(name_str.clone(), vmodule.into()); + + if kind == VTabKind::TableValuedFunction { + if let Ok(vtab) = VirtualTable::function(&name_str, syms) { + // Use the schema handler to insert the table + let table = Box::into_raw(Box::new(Arc::new(Table::Virtual(vtab)))); + let c_name = match CString::new(name_str) { + Ok(s) => s, + Err(_) => return ResultCode::Error, + }; + + let result = (ext_ctx.schema_handler)( + ext_ctx.schema_data, + c_name.as_ptr(), + table as *mut c_void, + ); + if result != ResultCode::OK { + let _ = Box::from_raw(table); + return result; + } + } else { + return ResultCode::Error; + } + } + } + ResultCode::OK +} #[derive(Clone)] pub struct VTabImpl { @@ -38,8 +146,14 @@ pub(crate) unsafe extern "C" fn register_scalar_function( if ctx.is_null() { return ResultCode::Error; } - let conn = unsafe { &*(ctx as *const Connection) }; - conn.register_scalar_function_impl(&name_str, func) + let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) }; + unsafe { + (*ext_ctx.syms).functions.insert( + name_str.clone(), + Rc::new(ExternalFunc::new_scalar(name_str, func)), + ); + } + ResultCode::OK } pub(crate) unsafe extern "C" fn register_aggregate_function( @@ -58,30 +172,18 @@ pub(crate) unsafe extern "C" fn register_aggregate_function( if ctx.is_null() { return ResultCode::Error; } - let conn = unsafe { &*(ctx as *const Connection) }; - conn.register_aggregate_function_impl(&name_str, args, (init_func, step_func, finalize_func)) -} - -pub(crate) unsafe extern "C" fn register_vtab_module( - ctx: *mut c_void, - name: *const c_char, - module: VTabModuleImpl, - kind: VTabKind, -) -> ResultCode { - if name.is_null() || ctx.is_null() { - return ResultCode::Error; + let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) }; + unsafe { + (*ext_ctx.syms).functions.insert( + name_str.clone(), + Rc::new(ExternalFunc::new_aggregate( + name_str, + args, + (init_func, step_func, finalize_func), + )), + ); } - let c_str = unsafe { CString::from_raw(name as *mut _) }; - let name_str = match c_str.to_str() { - Ok(s) => s.to_string(), - Err(_) => return ResultCode::Error, - }; - if ctx.is_null() { - return ResultCode::Error; - } - let conn = unsafe { &mut *(ctx as *mut Connection) }; - - conn.register_vtab_module_impl(&name_str, module, kind) + ResultCode::OK } impl Database { @@ -110,58 +212,64 @@ impl Database { let db = Self::open_file(io.clone(), path, false, false)?; Ok((io, db)) } + + /// Register any built-in extensions that can be stored on the Database so we do not have + /// to register these once-per-connection, and the connection can just extend its symbol table + pub fn register_global_builtin_extensions(&self) -> Result<(), String> { + let syms = self.builtin_syms.as_ptr(); + // Pass the mutex pointer and the appropriate handler + let schema_mutex_ptr = &self.schema as *const Mutex> as *mut Mutex>; + let ctx = Box::into_raw(Box::new(ExtensionCtx { + syms, + schema_data: schema_mutex_ptr as *mut c_void, + schema_handler: handle_schema_insert_database, + })); + let mut ext_api = ExtensionApi { + ctx: ctx as *mut c_void, + register_scalar_function, + register_aggregate_function, + register_vtab_module, + #[cfg(feature = "fs")] + vfs_interface: turso_ext::VfsInterface { + register_vfs: dynamic::register_vfs, + builtin_vfs: std::ptr::null_mut(), + builtin_vfs_count: 0, + }, + }; + + #[cfg(feature = "uuid")] + crate::uuid::register_extension(&mut ext_api); + #[cfg(feature = "series")] + crate::series::register_extension(&mut ext_api); + #[cfg(feature = "fs")] + { + let vfslist = add_builtin_vfs_extensions(Some(ext_api)).map_err(|e| e.to_string())?; + for (name, vfs) in vfslist { + add_vfs_module(name, vfs); + } + } + let _ = unsafe { Box::from_raw(ctx) }; + Ok(()) + } } impl Connection { - fn register_scalar_function_impl(&self, name: &str, func: ScalarFunction) -> ResultCode { - self.syms.borrow_mut().functions.insert( - name.to_string(), - Rc::new(ExternalFunc::new_scalar(name.to_string(), func)), - ); - ResultCode::OK - } - - fn register_aggregate_function_impl( - &self, - name: &str, - args: i32, - func: ExternAggFunc, - ) -> ResultCode { - self.syms.borrow_mut().functions.insert( - name.to_string(), - Rc::new(ExternalFunc::new_aggregate(name.to_string(), args, func)), - ); - ResultCode::OK - } - - fn register_vtab_module_impl( - &mut self, - name: &str, - module: VTabModuleImpl, - kind: VTabKind, - ) -> ResultCode { - let module = Rc::new(module); - let vmodule = VTabImpl { - module_kind: kind, - implementation: module, + /// # Safety: + /// Only to be used when registering a staticly linked extension manually. + /// you probably want to use `Connection::load_extension(path)`. + /// Do not call if you have multiple connection open. + pub fn _build_turso_ext(&self) -> ExtensionApi { + let schema_ptr = self.schema.as_ptr(); + let schema_direct = unsafe { Arc::as_ptr(&*schema_ptr) as *mut Schema }; + let ctx = ExtensionCtx { + syms: self.syms.as_ptr(), + schema_data: schema_direct as *mut c_void, + schema_handler: handle_schema_insert_connection, }; - self.syms - .borrow_mut() - .vtab_modules - .insert(name.to_string(), vmodule.into()); - if kind == VTabKind::TableValuedFunction { - if let Ok(vtab) = VirtualTable::function(name, &self.syms.borrow()) { - self.with_schema_mut(|schema| schema.add_virtual_table(vtab)); - } else { - return ResultCode::Error; - } - } - ResultCode::OK - } - pub fn build_turso_ext(&self) -> ExtensionApi { + let ctx = Box::into_raw(Box::new(ctx)) as *mut c_void; ExtensionApi { - ctx: self as *const _ as *mut c_void, + ctx, register_scalar_function, register_aggregate_function, register_vtab_module, @@ -174,20 +282,13 @@ impl Connection { } } - pub fn register_builtins(&self) -> Result<(), String> { - #[allow(unused_variables)] - let mut ext_api = self.build_turso_ext(); - #[cfg(feature = "uuid")] - crate::uuid::register_extension(&mut ext_api); - #[cfg(feature = "series")] - crate::series::register_extension(&mut ext_api); - #[cfg(feature = "fs")] - { - let vfslist = add_builtin_vfs_extensions(Some(ext_api)).map_err(|e| e.to_string())?; - for (name, vfs) in vfslist { - add_vfs_module(name, vfs); - } + /// # Safety: + /// Only to be used if you have previously called Connection::build_turso_ext + /// before registering an extension manually. + pub unsafe fn _free_extension_ctx(&self, api: ExtensionApi) { + if api.ctx.is_null() { + return; } - Ok(()) + let _ = unsafe { Box::from_raw(api.ctx as *mut ExtensionCtx) }; } }