clt_database/ext/
mod.rs

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

clt_database/ext/
mod.rs
