limbo_ext/
vtabs.rs

1use crate::{types::StepResult, ExtResult, ResultCode, Value};
2use std::{
3    ffi::{c_char, c_void, CStr, CString},
4    num::NonZeroUsize,
5    sync::Arc,
6};
7
8pub type RegisterModuleFn = unsafe extern "C" fn(
9    ctx: *mut c_void,
10    name: *const c_char,
11    module: VTabModuleImpl,
12    kind: VTabKind,
13) -> ResultCode;
14
15#[repr(C)]
16#[derive(Clone, Debug)]
17pub struct VTabModuleImpl {
18    pub name: *const c_char,
19    pub create: VtabFnCreate,
20    pub open: VtabFnOpen,
21    pub close: VtabFnClose,
22    pub filter: VtabFnFilter,
23    pub column: VtabFnColumn,
24    pub next: VtabFnNext,
25    pub eof: VtabFnEof,
26    pub update: VtabFnUpdate,
27    pub rowid: VtabRowIDFn,
28    pub destroy: VtabFnDestroy,
29    pub best_idx: BestIdxFn,
30}
31
32#[repr(C)]
33pub struct VTabCreateResult {
34    pub code: ResultCode,
35    pub schema: *const c_char,
36    pub table: *const c_void,
37}
38
39#[cfg(feature = "core_only")]
40impl VTabModuleImpl {
41    pub fn create(&self, args: Vec<Value>) -> crate::ExtResult<(String, *const c_void)> {
42        let result = unsafe { (self.create)(args.as_ptr(), args.len() as i32) };
43        for arg in args {
44            unsafe { arg.__free_internal_type() };
45        }
46        if !result.code.is_ok() {
47            return Err(result.code);
48        }
49        let schema = unsafe { std::ffi::CString::from_raw(result.schema as *mut _) };
50        Ok((schema.to_string_lossy().to_string(), result.table))
51    }
52
53    // TODO: This function is temporary and should eventually be removed.
54    //       The only difference from `create` is that it takes ownership of the table instance.
55    //       Currently, it is used to generate virtual table column names that are stored in
56    //       `sqlite_schema` alongside the table's schema.
57    //       However, storing column names is not necessary to match SQLite's behavior.
58    //       SQLite computes the list of columns dynamically each time the `.schema` command
59    //       is executed, using the `shell_add_schema` UDF function.
60    pub fn create_schema(&self, args: Vec<Value>) -> crate::ExtResult<String> {
61        self.create(args).and_then(|(schema, table)| {
62            // Drop the allocated table instance to avoid a memory leak.
63            let result = unsafe { (self.destroy)(table) };
64            if result.is_ok() {
65                Ok(schema)
66            } else {
67                Err(result)
68            }
69        })
70    }
71}
72
73pub type VtabFnCreate = unsafe extern "C" fn(args: *const Value, argc: i32) -> VTabCreateResult;
74
75pub type VtabFnOpen = unsafe extern "C" fn(table: *const c_void, conn: *mut Conn) -> *const c_void;
76
77pub type VtabFnClose = unsafe extern "C" fn(cursor: *const c_void) -> ResultCode;
78
79pub type VtabFnFilter = unsafe extern "C" fn(
80    cursor: *const c_void,
81    argc: i32,
82    argv: *const Value,
83    idx_str: *const c_char,
84    idx_num: i32,
85) -> ResultCode;
86
87pub type VtabFnColumn = unsafe extern "C" fn(cursor: *const c_void, idx: u32) -> Value;
88
89pub type VtabFnNext = unsafe extern "C" fn(cursor: *const c_void) -> ResultCode;
90
91pub type VtabFnEof = unsafe extern "C" fn(cursor: *const c_void) -> bool;
92
93pub type VtabRowIDFn = unsafe extern "C" fn(cursor: *const c_void) -> i64;
94
95pub type VtabFnUpdate = unsafe extern "C" fn(
96    table: *const c_void,
97    argc: i32,
98    argv: *const Value,
99    p_out_rowid: *mut i64,
100) -> ResultCode;
101
102pub type VtabFnDestroy = unsafe extern "C" fn(table: *const c_void) -> ResultCode;
103
104pub type BestIdxFn = unsafe extern "C" fn(
105    constraints: *const ConstraintInfo,
106    constraint_len: i32,
107    order_by: *const OrderByInfo,
108    order_by_len: i32,
109) -> ExtIndexInfo;
110
111#[repr(C)]
112#[derive(Clone, Copy, Debug, PartialEq)]
113pub enum VTabKind {
114    VirtualTable,
115    TableValuedFunction,
116}
117
118pub trait VTabModule: 'static {
119    type Table: VTable;
120    const VTAB_KIND: VTabKind;
121    const NAME: &'static str;
122
123    /// Creates a new instance of a virtual table.
124    /// Returns a tuple where the first element is the table's schema.
125    fn create(args: &[Value]) -> Result<(String, Self::Table), ResultCode>;
126}
127
128pub trait VTable {
129    type Cursor: VTabCursor<Error = Self::Error>;
130    type Error: std::fmt::Display;
131
132    /// 'conn' is an Option to allow for testing. Otherwise a valid connection to the core database
133    /// that created the virtual table will be available to use in your extension here.
134    fn open(&self, _conn: Option<Arc<Connection>>) -> Result<Self::Cursor, Self::Error>;
135    fn update(&mut self, _rowid: i64, _args: &[Value]) -> Result<(), Self::Error> {
136        Ok(())
137    }
138    fn insert(&mut self, _args: &[Value]) -> Result<i64, Self::Error> {
139        Ok(0)
140    }
141    fn delete(&mut self, _rowid: i64) -> Result<(), Self::Error> {
142        Ok(())
143    }
144    fn destroy(&mut self) -> Result<(), Self::Error> {
145        Ok(())
146    }
147    fn best_index(_constraints: &[ConstraintInfo], _order_by: &[OrderByInfo]) -> IndexInfo {
148        IndexInfo {
149            idx_num: 0,
150            idx_str: None,
151            order_by_consumed: false,
152            estimated_cost: 1_000_000.0,
153            estimated_rows: u32::MAX,
154            constraint_usages: _constraints
155                .iter()
156                .map(|_| ConstraintUsage {
157                    argv_index: Some(0),
158                    omit: false,
159                })
160                .collect(),
161        }
162    }
163}
164
165pub trait VTabCursor: Sized {
166    type Error: std::fmt::Display;
167    fn filter(&mut self, args: &[Value], idx_info: Option<(&str, i32)>) -> ResultCode;
168    fn rowid(&self) -> i64;
169    fn column(&self, idx: u32) -> Result<Value, Self::Error>;
170    fn eof(&self) -> bool;
171    fn next(&mut self) -> ResultCode;
172    fn close(&self) -> ResultCode {
173        ResultCode::OK
174    }
175}
176
177#[repr(u8)]
178#[derive(Copy, Clone, Debug, PartialEq, Eq)]
179pub enum ConstraintOp {
180    Eq = 2,
181    Lt = 4,
182    Le = 8,
183    Gt = 16,
184    Ge = 32,
185    Match = 64,
186    Like = 65,
187    Glob = 66,
188    Regexp = 67,
189    Ne = 68,
190    IsNot = 69,
191    IsNotNull = 70,
192    IsNull = 71,
193    Is = 72,
194    In = 73,
195}
196
197#[repr(C)]
198#[derive(Copy, Clone)]
199/// Describes an ORDER BY clause in a query involving a virtual table.
200/// Passed along with the constraints to xBestIndex.
201pub struct OrderByInfo {
202    /// The index of the column referenced in the ORDER BY clause.
203    pub column_index: u32,
204    /// Whether or not the clause is in descending order.
205    pub desc: bool,
206}
207
208/// The internal (core) representation of an 'index' on a virtual table.
209/// Returned from xBestIndex and then processed and passed to VFilter.
210#[derive(Debug, Clone)]
211pub struct IndexInfo {
212    /// The index number, used to identify the index internally by the VTab
213    pub idx_num: i32,
214    /// Optional index name. these are passed to vfilter in a tuple (idx_num, idx_str)
215    pub idx_str: Option<String>,
216    /// Whether the index is used for order by
217    pub order_by_consumed: bool,
218    /// TODO: for eventual cost based query planning
219    pub estimated_cost: f64,
220    /// Estimated number of rows that the query will return
221    pub estimated_rows: u32,
222    /// List of constraints that can be used to optimize the query.
223    pub constraint_usages: Vec<ConstraintUsage>,
224}
225impl Default for IndexInfo {
226    fn default() -> Self {
227        Self {
228            idx_num: 0,
229            idx_str: None,
230            order_by_consumed: false,
231            estimated_cost: 1_000_000.0,
232            estimated_rows: u32::MAX,
233            constraint_usages: Vec::new(),
234        }
235    }
236}
237
238impl IndexInfo {
239    ///
240    /// Converts IndexInfo to an FFI-safe `ExtIndexInfo`.
241    /// This method transfers ownership of `constraint_usages` and `idx_str`,
242    /// which must later be reclaimed using `from_ffi` to prevent leaks.
243    pub fn to_ffi(self) -> ExtIndexInfo {
244        let len = self.constraint_usages.len();
245        let ptr = Box::into_raw(self.constraint_usages.into_boxed_slice()) as *mut ConstraintUsage;
246        let idx_str_len = self.idx_str.as_ref().map(|s| s.len()).unwrap_or(0);
247        let c_idx_str = self
248            .idx_str
249            .and_then(|s| std::ffi::CString::new(s).ok())
250            .map(|cs| cs.into_raw())
251            .unwrap_or(std::ptr::null_mut());
252        ExtIndexInfo {
253            idx_num: self.idx_num,
254            estimated_cost: self.estimated_cost,
255            estimated_rows: self.estimated_rows,
256            order_by_consumed: self.order_by_consumed,
257            constraint_usages_ptr: ptr,
258            constraint_usage_len: len,
259            idx_str: c_idx_str as *mut _,
260            idx_str_len,
261        }
262    }
263
264    /// Reclaims ownership of `constraint_usages` and `idx_str` from an FFI-safe `ExtIndexInfo`.
265    /// # Safety
266    /// This method is unsafe because it can cause memory leaks if not used correctly.
267    /// to_ffi and from_ffi are meant to send index info across ffi bounds then immediately reclaim it.
268    pub unsafe fn from_ffi(ffi: ExtIndexInfo) -> Self {
269        let constraint_usages = unsafe {
270            Box::from_raw(std::slice::from_raw_parts_mut(
271                ffi.constraint_usages_ptr,
272                ffi.constraint_usage_len,
273            ))
274            .to_vec()
275        };
276        let idx_str = if ffi.idx_str.is_null() {
277            None
278        } else {
279            Some(unsafe {
280                std::ffi::CString::from_raw(ffi.idx_str as *mut _)
281                    .to_string_lossy()
282                    .into_owned()
283            })
284        };
285        Self {
286            idx_num: ffi.idx_num,
287            idx_str,
288            order_by_consumed: ffi.order_by_consumed,
289            estimated_cost: ffi.estimated_cost,
290            estimated_rows: ffi.estimated_rows,
291            constraint_usages,
292        }
293    }
294}
295
296#[repr(C)]
297#[derive(Clone, Debug)]
298/// FFI representation of IndexInfo.
299pub struct ExtIndexInfo {
300    pub idx_num: i32,
301    pub idx_str: *const u8,
302    pub idx_str_len: usize,
303    pub order_by_consumed: bool,
304    pub estimated_cost: f64,
305    pub estimated_rows: u32,
306    pub constraint_usages_ptr: *mut ConstraintUsage,
307    pub constraint_usage_len: usize,
308}
309
310/// Returned from xBestIndex to describe how the virtual table
311/// can use the constraints in the WHERE clause of a query.
312#[derive(Debug, Clone, Copy)]
313pub struct ConstraintUsage {
314    /// 1 based index of the argument passed
315    pub argv_index: Option<u32>,
316    /// If true, core can omit this constraint in the vdbe layer.
317    pub omit: bool,
318}
319
320#[derive(Clone, Copy, Debug)]
321#[repr(C)]
322/// The primary argument to xBestIndex, which describes a constraint
323/// in a query involving a virtual table.
324pub struct ConstraintInfo {
325    /// The index of the column referenced in the WHERE clause.
326    pub column_index: u32,
327    /// The operator used in the clause.
328    pub op: ConstraintOp,
329    /// Whether or not constraint is garaunteed to be enforced.
330    pub usable: bool,
331    /// packed integer with the index of the constraint in the planner,
332    /// and the side of the binary expr that the relevant column is on.
333    pub plan_info: u32,
334}
335
336impl ConstraintInfo {
337    #[inline(always)]
338    pub fn pack_plan_info(pred_idx: u32, is_right_side: bool) -> u32 {
339        ((pred_idx) << 1) | (is_right_side as u32)
340    }
341    #[inline(always)]
342    pub fn unpack_plan_info(&self) -> (usize, bool) {
343        ((self.plan_info >> 1) as usize, (self.plan_info & 1) != 0)
344    }
345}
346
347pub type PrepareStmtFn = unsafe extern "C" fn(api: *mut Conn, sql: *const c_char) -> *mut Stmt;
348pub type ExecuteFn = unsafe extern "C" fn(
349    ctx: *mut Conn,
350    sql: *const c_char,
351    args: *mut Value,
352    arg_count: i32,
353    last_insert_rowid: *mut i64,
354) -> ResultCode;
355pub type GetColumnNamesFn =
356    unsafe extern "C" fn(ctx: *mut Stmt, count: *mut i32) -> *mut *mut c_char;
357pub type BindArgsFn = unsafe extern "C" fn(ctx: *mut Stmt, idx: i32, arg: Value) -> ResultCode;
358pub type StmtStepFn = unsafe extern "C" fn(ctx: *mut Stmt) -> ResultCode;
359pub type StmtGetRowValuesFn = unsafe extern "C" fn(ctx: *mut Stmt);
360pub type FreeCurrentRowFn = unsafe extern "C" fn(ctx: *mut Stmt);
361pub type CloseConnectionFn = unsafe extern "C" fn(ctx: *mut c_void);
362pub type CloseStmtFn = unsafe extern "C" fn(ctx: *mut Stmt);
363
364/// core database connection
365/// public fields for core only
366#[repr(C)]
367#[derive(Debug, Clone)]
368pub struct Conn {
369    // boxed Rc::Weak from core::Connection
370    pub _ctx: *mut c_void,
371    pub _prepare_stmt: PrepareStmtFn,
372    pub _execute: ExecuteFn,
373    pub _close: CloseConnectionFn,
374}
375
376impl Conn {
377    pub fn new(
378        ctx: *mut c_void,
379        prepare_stmt: PrepareStmtFn,
380        exec_fn: ExecuteFn,
381        close: CloseConnectionFn,
382    ) -> Self {
383        Conn {
384            _ctx: ctx,
385            _prepare_stmt: prepare_stmt,
386            _execute: exec_fn,
387            _close: close,
388        }
389    }
390
391    /// # Safety
392    /// Dereferences a null pointer with a null check
393    pub unsafe fn from_ptr(ptr: *mut Conn) -> crate::ExtResult<&'static mut Self> {
394        if ptr.is_null() {
395            return Err(ResultCode::Error);
396        }
397        Ok(unsafe { &mut *(ptr) })
398    }
399
400    pub fn close(&mut self) {
401        if self._ctx.is_null() {
402            return;
403        }
404        unsafe { (self._close)(self._ctx) };
405        self._ctx = std::ptr::null_mut();
406    }
407
408    /// execute a SQL statement with the given arguments.
409    /// optionally returns the last inserted rowid for the query
410    pub fn execute(&self, sql: &str, args: &[Value]) -> crate::ExtResult<Option<usize>> {
411        let Ok(sql) = CString::new(sql) else {
412            return Err(ResultCode::Error);
413        };
414        let arg_count = args.len() as i32;
415        let args = args.as_ptr();
416        let last_insert_rowid = 0;
417        if let ResultCode::OK = unsafe {
418            (self._execute)(
419                self as *const _ as *mut Conn,
420                sql.as_ptr(),
421                args as *mut Value,
422                arg_count,
423                &last_insert_rowid as *const _ as *mut i64,
424            )
425        } {
426            return Ok(Some(last_insert_rowid as usize));
427        }
428        Err(ResultCode::Error)
429    }
430
431    pub fn prepare_stmt(&self, sql: &str) -> *mut Stmt {
432        let Ok(sql) = CString::new(sql) else {
433            return std::ptr::null_mut();
434        };
435        unsafe { (self._prepare_stmt)(self as *const _ as *mut Conn, sql.as_ptr()) }
436    }
437}
438
439/// Prepared statement for querying a core database connection public API for extensions
440/// Statements can be manually closed.
441#[derive(Debug)]
442#[repr(transparent)]
443pub struct Statement(*mut Stmt);
444
445impl Drop for Statement {
446    fn drop(&mut self) {
447        if self.0.is_null() {
448            return;
449        }
450        unsafe { (*self.0).close() }
451    }
452}
453
454/// Public API for methods to allow extensions to query other tables for
455/// the connection that opened the VTable. This value and its resources are cleaned up when
456/// the VTable is dropped, so there is no need to manually close the connection.
457#[derive(Debug)]
458#[repr(transparent)]
459pub struct Connection(*mut Conn);
460
461impl Connection {
462    pub fn new(ctx: *mut Conn) -> Self {
463        Connection(ctx)
464    }
465
466    /// From the included SQL string, prepare a statement for execution.
467    pub fn prepare(self: &Arc<Self>, sql: &str) -> ExtResult<Statement> {
468        let stmt = unsafe { (*self.0).prepare_stmt(sql) };
469        if stmt.is_null() {
470            return Err(ResultCode::Error);
471        }
472        Ok(Statement(stmt))
473    }
474
475    /// Execute a SQL statement with the given arguments.
476    /// Optionally returns the last inserted rowid for the query.
477    pub fn execute(self: &Arc<Self>, sql: &str, args: &[Value]) -> crate::ExtResult<Option<usize>> {
478        if self.0.is_null() {
479            return Err(ResultCode::Error);
480        }
481        unsafe { (*self.0).execute(sql, args) }
482    }
483}
484
485impl Statement {
486    /// Bind a value to a parameter in the prepared statement.
487    ///```text
488    /// let stmt = conn.prepare_stmt("select * from users where name = ?");
489    /// stmt.bind_at(1, Value::from_text("test".into()));
490    ///```
491    pub fn bind_at(&self, idx: NonZeroUsize, arg: Value) {
492        unsafe {
493            (*self.0).bind_args(idx, arg);
494        }
495    }
496
497    /// Execute the statement and return the next row
498    ///```text
499    /// while stmt.step() == StepResult::Row {
500    ///     let row = stmt.get_row();
501    ///     println!("row: {:?}", row);
502    /// }
503    /// ```
504    pub fn step(&self) -> StepResult {
505        unsafe { (*self.0).step() }
506    }
507
508    // Get the current row values
509    ///```text
510    /// while stmt.step() == StepResult::Row {
511    ///    let row = stmt.get_row();
512    ///    println!("row: {:?}", row);
513    ///```
514    pub fn get_row(&mut self) -> &[Value] {
515        unsafe { (*self.0).get_row() }
516    }
517
518    /// Get the result column names for the prepared statement
519    pub fn get_column_names(&self) -> Vec<String> {
520        unsafe { (*self.0).get_column_names() }
521    }
522
523    /// Close the statement and clean up resources.
524    pub fn close(self) {
525        if self.0.is_null() {
526            return;
527        }
528        unsafe { (*self.0).close() }
529    }
530}
531
532/// Internal/core use _only_
533/// Extensions should not import or use this type directly
534#[repr(C)]
535pub struct Stmt {
536    // Rc::into_raw from core::Connection
537    pub _conn: *mut c_void,
538    // Rc::into_raw from core::Statement
539    pub _ctx: *mut c_void,
540    pub _bind_args_fn: BindArgsFn,
541    pub _step: StmtStepFn,
542    pub _get_row_values: StmtGetRowValuesFn,
543    pub _get_column_names: GetColumnNamesFn,
544    pub _free_current_row: FreeCurrentRowFn,
545    pub _close: CloseStmtFn,
546    pub current_row: *mut Value,
547    pub current_row_len: i32,
548}
549
550impl Stmt {
551    #[allow(clippy::too_many_arguments)]
552    pub fn new(
553        conn: *mut c_void,
554        ctx: *mut c_void,
555        bind: BindArgsFn,
556        step: StmtStepFn,
557        rows: StmtGetRowValuesFn,
558        names: GetColumnNamesFn,
559        free_row: FreeCurrentRowFn,
560        close: CloseStmtFn,
561    ) -> Self {
562        Stmt {
563            _conn: conn,
564            _ctx: ctx,
565            _bind_args_fn: bind,
566            _step: step,
567            _get_row_values: rows,
568            _get_column_names: names,
569            _free_current_row: free_row,
570            _close: close,
571            current_row: std::ptr::null_mut(),
572            current_row_len: -1,
573        }
574    }
575
576    /// Close the statement
577    pub fn close(&mut self) {
578        // null check to prevent double free
579        if self._ctx.is_null() {
580            return;
581        }
582        unsafe { (self._close)(self as *const Stmt as *mut Stmt) };
583        self._ctx = std::ptr::null_mut();
584    }
585
586    /// # Safety
587    /// Derefs a null ptr, does a null check first
588    pub unsafe fn from_ptr(ptr: *mut Stmt) -> ExtResult<&'static mut Self> {
589        if ptr.is_null() {
590            return Err(ResultCode::Error);
591        }
592        Ok(unsafe { &mut *(ptr) })
593    }
594
595    /// Returns the pointer to the statement.
596    pub fn to_ptr(&self) -> *mut Stmt {
597        self as *const Stmt as *mut Stmt
598    }
599
600    /// Bind a value to a parameter in the prepared statement
601    /// Own the value so it can be freed in core
602    fn bind_args(&self, idx: NonZeroUsize, arg: Value) {
603        unsafe {
604            (self._bind_args_fn)(self.to_ptr(), idx.get() as i32, arg);
605        };
606    }
607
608    /// Execute the statement to attempt to retrieve the next result row.
609    fn step(&self) -> StepResult {
610        unsafe { (self._step)(self.to_ptr()) }.into()
611    }
612
613    /// Free the memory for the values obtained from the `get_row` method.
614    /// This is easier done on core side because __free_internal_type is 'core_only'
615    /// feature to prevent extensions causing memory issues.
616    /// # Safety
617    /// This fn is unsafe because it derefs a raw pointer after null and
618    /// length checks. This fn should only be called with the pointer returned from get_row.
619    pub unsafe fn free_current_row(&mut self) {
620        if self.current_row.is_null() || self.current_row_len <= 0 {
621            return;
622        }
623        // free from the core side so we don't have to expose `__free_internal_type`
624        (self._free_current_row)(self.to_ptr());
625        self.current_row = std::ptr::null_mut();
626        self.current_row_len = -1;
627    }
628
629    /// Returns the values from the current row in the prepared statement, should
630    /// be called after the step() method returns `StepResult::Row`
631    pub fn get_row(&self) -> &[Value] {
632        unsafe { (self._get_row_values)(self.to_ptr()) };
633        if self.current_row.is_null() || self.current_row_len < 1 {
634            return &[];
635        }
636        let col_count = self.current_row_len;
637        unsafe { std::slice::from_raw_parts(self.current_row, col_count as usize) }
638    }
639
640    /// Returns the names of the result columns for the prepared statement.
641    pub fn get_column_names(&self) -> Vec<String> {
642        let mut count_value: i32 = 0;
643        let count: *mut i32 = &mut count_value;
644        let col_names = unsafe { (self._get_column_names)(self.to_ptr(), count) };
645        if col_names.is_null() || count_value == 0 {
646            return Vec::new();
647        }
648        let mut names = Vec::new();
649        let slice = unsafe { std::slice::from_raw_parts(col_names, count_value as usize) };
650        for x in slice {
651            let name = unsafe { CStr::from_ptr(*x) };
652            if let Ok(s) = name.to_str() {
653                names.push(s.to_string());
654            }
655        }
656        unsafe { free_column_names(col_names, count_value) };
657        names
658    }
659}
660
661/// Free the column names returned from get_column_names
662/// # Safety
663/// This function is unsafe because it derefs a raw pointer, this fn
664/// should only be called with the pointer returned from get_column_names
665/// only when they will no longer be used.
666pub unsafe fn free_column_names(names: *mut *mut c_char, count: i32) {
667    if names.is_null() || count < 1 {
668        return;
669    }
670    let slice = std::slice::from_raw_parts_mut(names, count as usize);
671
672    for name in slice {
673        if !name.is_null() {
674            let _ = CString::from_raw(*name);
675        }
676    }
677    let _ = Box::from_raw(names);
678}
limbo_ext/vtabs.rs

limbo_ext/
vtabs.rs
