Skip to main content

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