turso_core/ext/
mod.rs

1#[cfg(feature = "fs")]
2mod dynamic;
3mod vtab_xconnect;
4use crate::index_method::backing_btree::BackingBtreeIndexMethod;
5#[cfg(all(feature = "fts", not(target_family = "wasm")))]
6use crate::index_method::fts::{FtsIndexMethod, FTS_INDEX_METHOD_NAME};
7use crate::index_method::toy_vector_sparse_ivf::VectorSparseInvertedIndexMethod;
8use crate::index_method::{
9    BACKING_BTREE_INDEX_METHOD_NAME, TOY_VECTOR_SPARSE_IVF_INDEX_METHOD_NAME,
10};
11use crate::schema::{Schema, Table};
12use crate::sync::atomic::{AtomicU64, Ordering};
13use crate::sync::Mutex;
14#[cfg(all(target_os = "linux", feature = "io_uring", not(miri)))]
15use crate::UringIO;
16#[cfg(all(target_os = "windows", feature = "experimental_win_iocp", not(miri)))]
17use crate::WindowsIOCP;
18
19use crate::{function::ExternalFunc, Connection, Database};
20use crate::{vtab::VirtualTable, SymbolTable};
21#[cfg(feature = "fs")]
22use crate::{LimboError, IO};
23#[cfg(feature = "fs")]
24pub use dynamic::{add_builtin_vfs_extensions, add_vfs_module, list_vfs_modules, VfsMod};
25use std::{
26    ffi::{c_char, c_void, CStr, CString},
27    sync::Arc,
28};
29use turso_ext::{
30    ContextDestructor, ExtensionApi, InitAggFunction, ResultCode, ScalarFunction, VTabKind,
31    VTabModuleImpl, ValueDestructor,
32};
33pub use turso_ext::{FinalizeFunction, StepFunction, Value as ExtValue, ValueType as ExtValueType};
34pub use vtab_xconnect::{execute, prepare_stmt};
35
36/// The context passed to extensions to register with Core
37/// along with the function pointers
38#[repr(C)]
39pub struct ExtensionCtx {
40    syms: *mut SymbolTable,
41    schema: *mut c_void,
42    /// We must bump the prepare context generation so prepared statements
43    /// know they need to be reprepared after extension registration.
44    prepare_context_generation: *const AtomicU64,
45}
46
47pub(crate) unsafe extern "C" fn register_vtab_module(
48    ctx: *mut c_void,
49    name: *const c_char,
50    module: VTabModuleImpl,
51    kind: VTabKind,
52) -> ResultCode {
53    if name.is_null() || ctx.is_null() {
54        return ResultCode::Error;
55    }
56
57    let c_str = unsafe { CString::from_raw(name as *mut c_char) };
58    let name_str = match c_str.to_str() {
59        Ok(s) => s.to_string(),
60        Err(_) => return ResultCode::Error,
61    };
62
63    let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) };
64    let module = Arc::new(module);
65    let vmodule = VTabImpl {
66        module_kind: kind,
67        implementation: module,
68    };
69
70    unsafe {
71        let syms = &mut *ext_ctx.syms;
72        syms.vtab_modules.insert(name_str.clone(), vmodule.into());
73        if !ext_ctx.prepare_context_generation.is_null() {
74            (*ext_ctx.prepare_context_generation).fetch_add(1, Ordering::Release);
75        }
76
77        if kind == VTabKind::TableValuedFunction {
78            if let Ok(vtab) = VirtualTable::function(&name_str, syms) {
79                let table = Arc::new(Table::Virtual(vtab));
80                let mutex = &*(ext_ctx.schema as *mut Mutex<Arc<Schema>>);
81                let mut guard = mutex.lock();
82                let Ok(schema) = Schema::try_make_mut(&mut guard) else {
83                    return ResultCode::Error;
84                };
85                schema.tables.insert(name_str, table);
86            } else {
87                return ResultCode::Error;
88            }
89        }
90    }
91    ResultCode::OK
92}
93
94#[derive(Clone)]
95pub struct VTabImpl {
96    pub module_kind: VTabKind,
97    pub implementation: Arc<VTabModuleImpl>,
98}
99
100pub(crate) unsafe fn register_scalar_function(
101    ctx: *mut c_void,
102    name: *const c_char,
103    func: ScalarFunction,
104) -> ResultCode {
105    unsafe { register_scalar_function_with_options(ctx, name, -1, false, 0, func, None, None) }
106}
107
108pub(crate) unsafe extern "C" fn register_scalar_function_with_options(
109    ctx: *mut c_void,
110    name: *const c_char,
111    argc: i32,
112    deterministic: bool,
113    context: usize,
114    callback: ScalarFunction,
115    context_destructor: Option<ContextDestructor>,
116    value_destructor: Option<ValueDestructor>,
117) -> ResultCode {
118    if ctx.is_null() || name.is_null() || argc < -1 {
119        return ResultCode::InvalidArgs;
120    }
121    let c_str = unsafe { CStr::from_ptr(name) };
122    let name_str = match c_str.to_str() {
123        Ok(s) => crate::util::normalize_ident(s),
124        Err(_) => return ResultCode::InvalidArgs,
125    };
126    let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) };
127    unsafe {
128        (*ext_ctx.syms).functions.insert(
129            name_str.clone(),
130            Arc::new(ExternalFunc::new_scalar(
131                name_str,
132                argc,
133                deterministic,
134                context,
135                callback,
136                context_destructor,
137                value_destructor,
138            )),
139        );
140        if !ext_ctx.prepare_context_generation.is_null() {
141            (*ext_ctx.prepare_context_generation).fetch_add(1, Ordering::Release);
142        }
143    }
144    ResultCode::OK
145}
146
147pub(crate) unsafe extern "C" fn unregister_function(
148    ctx: *mut c_void,
149    name: *const c_char,
150) -> ResultCode {
151    if ctx.is_null() || name.is_null() {
152        return ResultCode::InvalidArgs;
153    }
154    let c_str = unsafe { CStr::from_ptr(name) };
155    let name_str = match c_str.to_str() {
156        Ok(s) => crate::util::normalize_ident(s),
157        Err(_) => return ResultCode::InvalidArgs,
158    };
159    let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) };
160    unsafe {
161        if (*ext_ctx.syms).functions.remove(&name_str).is_none() {
162            return ResultCode::NotFound;
163        }
164        if !ext_ctx.prepare_context_generation.is_null() {
165            (*ext_ctx.prepare_context_generation).fetch_add(1, Ordering::Release);
166        }
167    }
168    ResultCode::OK
169}
170
171pub(crate) unsafe extern "C" fn register_aggregate_function(
172    ctx: *mut c_void,
173    name: *const c_char,
174    args: i32,
175    context: usize,
176    init_func: InitAggFunction,
177    step_func: StepFunction,
178    finalize_func: FinalizeFunction,
179    context_destructor: Option<ContextDestructor>,
180    aggregate_destructor: Option<ContextDestructor>,
181    value_destructor: Option<ValueDestructor>,
182) -> ResultCode {
183    if ctx.is_null() || name.is_null() || args < -1 {
184        return ResultCode::InvalidArgs;
185    }
186    let c_str = unsafe { CStr::from_ptr(name) };
187    let name_str = match c_str.to_str() {
188        Ok(s) => crate::util::normalize_ident(s),
189        Err(_) => return ResultCode::InvalidArgs,
190    };
191    let ext_ctx = unsafe { &mut *(ctx as *mut ExtensionCtx) };
192    unsafe {
193        (*ext_ctx.syms).functions.insert(
194            name_str.clone(),
195            Arc::new(ExternalFunc::new_aggregate(
196                name_str,
197                args,
198                context,
199                (init_func, step_func, finalize_func),
200                context_destructor,
201                aggregate_destructor,
202                value_destructor,
203            )),
204        );
205        if !ext_ctx.prepare_context_generation.is_null() {
206            (*ext_ctx.prepare_context_generation).fetch_add(1, Ordering::Release);
207        }
208    }
209    ResultCode::OK
210}
211
212impl Database {
213    #[cfg(feature = "fs")]
214    #[allow(clippy::arc_with_non_send_sync, dead_code)]
215    pub fn open_with_vfs(
216        &self,
217        path: &str,
218        vfs: &str,
219    ) -> crate::Result<(Arc<dyn IO>, Arc<Database>)> {
220        use crate::{MemoryIO, SyscallIO};
221        use dynamic::get_vfs_modules;
222
223        let io: Arc<dyn IO> = match vfs {
224            "memory" => Arc::new(MemoryIO::new()),
225            #[cfg(feature = "io_memory_yield")]
226            "memory_yield" => Arc::new(crate::MemoryYieldIO::new()),
227            "syscall" => Arc::new(SyscallIO::new()?),
228            #[cfg(all(target_os = "linux", feature = "io_uring", not(miri)))]
229            "io_uring" => Arc::new(UringIO::new()?),
230            #[cfg(all(target_os = "windows", feature = "experimental_win_iocp", not(miri)))]
231            "experimental_win_iocp" => Arc::new(WindowsIOCP::new()?),
232            other => match get_vfs_modules().iter().find(|v| v.0 == vfs) {
233                Some((_, vfs)) => vfs.clone(),
234                None => {
235                    return Err(LimboError::InvalidArgument(format!("no such VFS: {other}")));
236                }
237            },
238        };
239        let db = Self::open_file(io.clone(), path)?;
240        Ok((io, db))
241    }
242
243    /// Register any built-in extensions that can be stored on the Database so we do not have
244    /// to register these once-per-connection, and the connection can just extend its symbol table
245    pub fn register_global_builtin_extensions(&self) -> Result<(), String> {
246        {
247            let mut syms = self.builtin_syms.write();
248            syms.index_methods.insert(
249                TOY_VECTOR_SPARSE_IVF_INDEX_METHOD_NAME.to_string(),
250                Arc::new(VectorSparseInvertedIndexMethod),
251            );
252            syms.index_methods.insert(
253                BACKING_BTREE_INDEX_METHOD_NAME.to_string(),
254                Arc::new(BackingBtreeIndexMethod),
255            );
256            #[cfg(all(feature = "fts", not(target_family = "wasm")))]
257            syms.index_methods
258                .insert(FTS_INDEX_METHOD_NAME.to_string(), Arc::new(FtsIndexMethod));
259        }
260        let syms = self.builtin_syms.data_ptr();
261        // Pass the mutex pointer and the appropriate handler
262        let schema_mutex_ptr =
263            &*self.schema as *const Mutex<Arc<Schema>> as *mut Mutex<Arc<Schema>>;
264        let ctx = Box::into_raw(Box::new(ExtensionCtx {
265            syms,
266            schema: schema_mutex_ptr as *mut c_void,
267            prepare_context_generation: std::ptr::null(),
268        }));
269        #[allow(unused)]
270        let mut ext_api = ExtensionApi {
271            ctx: ctx as *mut c_void,
272            register_scalar_function: register_scalar_function_with_options,
273            register_aggregate_function,
274            unregister_function,
275            register_vtab_module,
276            #[cfg(feature = "fs")]
277            vfs_interface: turso_ext::VfsInterface {
278                register_vfs: dynamic::register_vfs,
279                builtin_vfs: std::ptr::null_mut(),
280                builtin_vfs_count: 0,
281            },
282        };
283
284        #[cfg(feature = "uuid")]
285        crate::uuid::register_extension(&mut ext_api);
286        #[cfg(feature = "series")]
287        crate::series::register_extension(&mut ext_api);
288        #[cfg(feature = "time")]
289        crate::time::register_extension(&mut ext_api);
290        #[cfg(feature = "percentile")]
291        crate::percentile::register_extension(&mut ext_api);
292        crate::regexp::register_extension(&mut ext_api);
293        #[cfg(feature = "fs")]
294        {
295            let vfslist = add_builtin_vfs_extensions(Some(ext_api)).map_err(|e| e.to_string())?;
296            for (name, vfs) in vfslist {
297                add_vfs_module(name, vfs);
298            }
299        }
300        let _ = unsafe { Box::from_raw(ctx) };
301        Ok(())
302    }
303}
304
305impl Connection {
306    /// Build the connection's extension api context for manually registering an extension.
307    /// you probably want to use `Connection::load_extension(path)`.
308    ///
309    /// # Safety
310    /// Only to be used when registering a staticly linked extension manually.
311    /// You should only ever call this method on your applications startup,
312    /// The caller is responsible for calling `_free_extension_ctx` after registering the
313    /// extension.
314    ///
315    /// usage:
316    /// ```ignore
317    /// let ext_api = conn._build_turso_ext();
318    /// unsafe {
319    ///     my_extension::register_extension(&mut ext_api);
320    ///     conn._free_extension_ctx(ext_api);
321    /// }
322    ///```
323    pub unsafe fn _build_turso_ext(&self) -> ExtensionApi {
324        let schema_mutex_ptr =
325            &*self.db.schema as *const Mutex<Arc<Schema>> as *mut Mutex<Arc<Schema>>;
326        let ctx = ExtensionCtx {
327            syms: self.syms.data_ptr(),
328            schema: schema_mutex_ptr as *mut c_void,
329            prepare_context_generation: &self.prepare_context_generation as *const _,
330        };
331        let ctx = Box::into_raw(Box::new(ctx)) as *mut c_void;
332        ExtensionApi {
333            ctx,
334            register_scalar_function: register_scalar_function_with_options,
335            register_aggregate_function,
336            unregister_function,
337            register_vtab_module,
338            #[cfg(feature = "fs")]
339            vfs_interface: turso_ext::VfsInterface {
340                register_vfs: dynamic::register_vfs,
341                builtin_vfs: std::ptr::null_mut(),
342                builtin_vfs_count: 0,
343            },
344        }
345    }
346
347    /// Free the connection's extension libary context after registering an extension manually.
348    /// # Safety
349    /// Only to be used if you have previously called Connection::build_turso_ext
350    pub unsafe fn _free_extension_ctx(&self, api: ExtensionApi) {
351        if api.ctx.is_null() {
352            return;
353        }
354        let _ = unsafe { Box::from_raw(api.ctx as *mut ExtensionCtx) };
355    }
356}
turso_core/ext/mod.rs

turso_core/ext/
mod.rs
