Introduce virtual table types

This commit is contained in:
Piotr Rzysko
2025-05-28 09:47:08 +02:00
parent b291179554
commit 4d35e36b77
4 changed files with 170 additions and 104 deletions

View File

@@ -14,7 +14,7 @@ use crate::{
function::Func,
schema::{Schema, Table},
translate::expr::walk_expr_mut,
util::{exprs_are_equivalent, normalize_ident, vtable_args},
util::{exprs_are_equivalent, normalize_ident},
vdbe::{builder::TableRefIdCounter, BranchOffset},
Result,
};
@@ -344,18 +344,7 @@ fn parse_from_clause_table<'a>(
}
ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => {
let normalized_name = &normalize_ident(qualified_name.name.0.as_str());
let args = match maybe_args {
Some(ref args) => vtable_args(args),
None => vec![],
};
let vtab = crate::VirtualTable::from_args(
None,
normalized_name,
args,
syms,
limbo_ext::VTabKind::TableValuedFunction,
maybe_args,
)?;
let vtab = crate::VirtualTable::function(normalized_name, maybe_args, syms)?;
let alias = maybe_alias
.as_ref()
.map(|a| match a {

View File

@@ -82,27 +82,12 @@ pub fn parse_schema_rows(
vtab.clone()
} else {
let mod_name = module_name_from_sql(sql)?;
if let Some(vmod) = syms.vtab_modules.get(mod_name) {
if let limbo_ext::VTabKind::VirtualTable = vmod.module_kind
{
crate::VirtualTable::from_args(
Some(name),
mod_name,
module_args_from_sql(sql)?,
syms,
vmod.module_kind,
None,
)?
} else {
return Err(LimboError::Corrupt("Table valued function: {name} registered as virtual table in schema".to_string()));
}
} else {
// the extension isn't loaded, so we emit a warning.
return Err(LimboError::ExtensionError(format!(
"Virtual table module '{}' not found\nPlease load extension",
&mod_name
)));
}
crate::VirtualTable::table(
Some(name),
mod_name,
module_args_from_sql(sql)?,
syms,
)?
};
schema.add_virtual_table(vtab);
} else {

View File

@@ -1040,23 +1040,8 @@ pub fn op_vcreate(
"Failed to upgrade Connection".to_string(),
));
};
let mod_type = conn
.syms
.borrow()
.vtab_modules
.get(&module_name)
.ok_or_else(|| {
crate::LimboError::ExtensionError(format!("Module {} not found", module_name))
})?
.module_kind;
let table = crate::VirtualTable::from_args(
Some(&table_name),
&module_name,
args,
&conn.syms.borrow(),
mod_type,
None,
)?;
let table =
crate::VirtualTable::table(Some(&table_name), &module_name, args, &conn.syms.borrow())?;
{
conn.syms
.borrow_mut()
@@ -1090,12 +1075,7 @@ pub fn op_vfilter(
let cursor = cursor.as_virtual_mut();
let mut args = Vec::with_capacity(*arg_count);
for i in 0..*arg_count {
args.push(
state.registers[args_reg + i]
.get_owned_value()
.clone()
.to_ffi(),
);
args.push(state.registers[args_reg + i].get_owned_value().clone());
}
let idx_str = if let Some(idx_str) = idx_str {
Some(state.registers[*idx_str].get_owned_value().to_string())

View File

@@ -1,5 +1,5 @@
use crate::schema::Column;
use crate::util::columns_from_create_table_body;
use crate::util::{columns_from_create_table_body, vtable_args};
use crate::{Connection, LimboError, SymbolTable, Value};
use fallible_iterator::FallibleIterator;
use limbo_ext::{ConstraintInfo, IndexInfo, OrderByInfo, ResultCode, VTabKind, VTabModuleImpl};
@@ -8,18 +8,137 @@ use std::cell::RefCell;
use std::ffi::c_void;
use std::rc::{Rc, Weak};
#[derive(Debug, Clone)]
enum VirtualTableType {
External(ExtVirtualTable),
}
#[derive(Clone, Debug)]
pub struct VirtualTable {
pub name: String,
pub args: Option<Vec<ast::Expr>>,
pub implementation: Rc<VTabModuleImpl>,
pub columns: Vec<Column>,
pub kind: VTabKind,
pub(crate) name: String,
pub(crate) args: Option<Vec<ast::Expr>>,
pub(crate) columns: Vec<Column>,
pub(crate) kind: VTabKind,
vtab_type: VirtualTableType,
}
impl VirtualTable {
pub(crate) fn function(
name: &str,
args: Option<Vec<ast::Expr>>,
syms: &SymbolTable,
) -> crate::Result<Rc<VirtualTable>> {
let ext_args = match args {
Some(ref args) => vtable_args(args),
None => vec![],
};
let (vtab, columns) =
ExtVirtualTable::from_args(name, ext_args, syms, VTabKind::TableValuedFunction)?;
let vtab = VirtualTable {
name: name.to_owned(),
args,
columns,
kind: VTabKind::TableValuedFunction,
vtab_type: VirtualTableType::External(vtab),
};
Ok(Rc::new(vtab))
}
pub fn table(
tbl_name: Option<&str>,
module_name: &str,
args: Vec<limbo_ext::Value>,
syms: &SymbolTable,
) -> crate::Result<Rc<VirtualTable>> {
let (table, columns) =
ExtVirtualTable::from_args(module_name, args, syms, VTabKind::VirtualTable)?;
let vtab = VirtualTable {
name: tbl_name.unwrap_or(module_name).to_owned(),
args: None,
columns,
kind: VTabKind::VirtualTable,
vtab_type: VirtualTableType::External(table),
};
Ok(Rc::new(vtab))
}
pub(crate) fn open(&self, conn: Weak<Connection>) -> crate::Result<VirtualTableCursor> {
match &self.vtab_type {
VirtualTableType::External(table) => {
Ok(VirtualTableCursor::External(table.open(conn)?))
}
}
}
pub(crate) fn update(&self, args: &[Value]) -> crate::Result<Option<i64>> {
match &self.vtab_type {
VirtualTableType::External(table) => table.update(args),
}
}
pub(crate) fn destroy(&self) -> crate::Result<()> {
match &self.vtab_type {
VirtualTableType::External(table) => table.destroy(),
}
}
pub(crate) fn best_index(
&self,
constraints: &[ConstraintInfo],
order_by: &[OrderByInfo],
) -> IndexInfo {
match &self.vtab_type {
VirtualTableType::External(table) => table.best_index(constraints, order_by),
}
}
}
pub enum VirtualTableCursor {
External(ExtVirtualTableCursor),
}
impl VirtualTableCursor {
pub(crate) fn next(&mut self) -> crate::Result<bool> {
match self {
VirtualTableCursor::External(cursor) => cursor.next(),
}
}
pub(crate) fn rowid(&self) -> i64 {
match self {
VirtualTableCursor::External(cursor) => cursor.rowid(),
}
}
pub(crate) fn column(&self, column: usize) -> crate::Result<Value> {
match self {
VirtualTableCursor::External(cursor) => cursor.column(column),
}
}
pub(crate) fn filter(
&mut self,
idx_num: i32,
idx_str: Option<String>,
arg_count: usize,
args: Vec<Value>,
) -> crate::Result<bool> {
match self {
VirtualTableCursor::External(cursor) => {
cursor.filter(idx_num, idx_str, arg_count, args)
}
}
}
}
#[derive(Clone, Debug)]
struct ExtVirtualTable {
implementation: Rc<VTabModuleImpl>,
table_ptr: *const c_void,
connection_ptr: RefCell<Option<*mut limbo_ext::Conn>>,
}
impl Drop for VirtualTable {
impl Drop for ExtVirtualTable {
fn drop(&mut self) {
if let Some(conn) = self.connection_ptr.borrow_mut().take() {
if conn.is_null() {
@@ -34,12 +153,8 @@ impl Drop for VirtualTable {
}
}
impl VirtualTable {
pub(crate) fn best_index(
&self,
constraints: &[ConstraintInfo],
order_by: &[OrderByInfo],
) -> IndexInfo {
impl ExtVirtualTable {
fn best_index(&self, constraints: &[ConstraintInfo], order_by: &[OrderByInfo]) -> IndexInfo {
unsafe {
IndexInfo::from_ffi((self.implementation.best_idx)(
constraints.as_ptr(),
@@ -51,14 +166,12 @@ impl VirtualTable {
}
/// takes ownership of the provided Args
pub(crate) fn from_args(
tbl_name: Option<&str>,
fn from_args(
module_name: &str,
args: Vec<limbo_ext::Value>,
syms: &SymbolTable,
kind: VTabKind,
exprs: Option<Vec<ast::Expr>>,
) -> crate::Result<Rc<Self>> {
) -> crate::Result<(Self, Vec<Column>)> {
let module = syms
.vtab_modules
.get(module_name)
@@ -66,30 +179,28 @@ impl VirtualTable {
"Virtual table module not found: {}",
module_name
)))?;
if let VTabKind::VirtualTable = kind {
if module.module_kind == VTabKind::TableValuedFunction {
return Err(LimboError::ExtensionError(format!(
"{} is not a virtual table module",
module_name
)));
}
};
if kind != module.module_kind {
let expected = match kind {
VTabKind::VirtualTable => "virtual table",
VTabKind::TableValuedFunction => "table-valued function",
};
return Err(LimboError::ExtensionError(format!(
"{} is not a {} module",
module_name, expected
)));
}
let (schema, table_ptr) = module.implementation.create(args)?;
let mut parser = Parser::new(schema.as_bytes());
if let ast::Cmd::Stmt(ast::Stmt::CreateTable { body, .. }) = parser.next()?.ok_or(
LimboError::ParseError("Failed to parse schema from virtual table module".to_string()),
)? {
let columns = columns_from_create_table_body(&body)?;
let vtab = Rc::new(VirtualTable {
name: tbl_name.unwrap_or(module_name).to_owned(),
let vtab = ExtVirtualTable {
connection_ptr: RefCell::new(None),
implementation: module.implementation.clone(),
columns,
args: exprs,
kind,
table_ptr,
});
return Ok(vtab);
};
return Ok((vtab, columns));
}
Err(LimboError::ParseError(
"Failed to parse schema from virtual table module".to_string(),
@@ -98,7 +209,7 @@ impl VirtualTable {
/// Accepts a Weak pointer to the connection that owns the VTable, that the module
/// can optionally use to query the other tables.
pub fn open(&self, conn: Weak<Connection>) -> crate::Result<VirtualTableCursor> {
fn open(&self, conn: Weak<Connection>) -> crate::Result<ExtVirtualTableCursor> {
// we need a Weak<Connection> to upgrade and call from the extension.
let weak_box: *mut Weak<Connection> = Box::into_raw(Box::new(conn));
let conn = limbo_ext::Conn::new(
@@ -111,10 +222,10 @@ impl VirtualTable {
// store the leaked connection pointer on the table so it can be freed on drop
*self.connection_ptr.borrow_mut() = Some(ext_conn_ptr);
let cursor = unsafe { (self.implementation.open)(self.table_ptr, ext_conn_ptr) };
VirtualTableCursor::new(cursor, self.implementation.clone())
ExtVirtualTableCursor::new(cursor, self.implementation.clone())
}
pub fn update(&self, args: &[Value]) -> crate::Result<Option<i64>> {
fn update(&self, args: &[Value]) -> crate::Result<Option<i64>> {
let arg_count = args.len();
let ext_args = args.iter().map(|arg| arg.to_ffi()).collect::<Vec<_>>();
let newrowid = 0i64;
@@ -138,7 +249,7 @@ impl VirtualTable {
}
}
pub fn destroy(&self) -> crate::Result<()> {
fn destroy(&self) -> crate::Result<()> {
let rc = unsafe { (self.implementation.destroy)(self.table_ptr) };
match rc {
ResultCode::OK => Ok(()),
@@ -147,13 +258,13 @@ impl VirtualTable {
}
}
pub struct VirtualTableCursor {
pub struct ExtVirtualTableCursor {
cursor: *const c_void,
implementation: Rc<VTabModuleImpl>,
}
impl VirtualTableCursor {
pub fn new(cursor: *const c_void, implementation: Rc<VTabModuleImpl>) -> crate::Result<Self> {
impl ExtVirtualTableCursor {
fn new(cursor: *const c_void, implementation: Rc<VTabModuleImpl>) -> crate::Result<Self> {
if cursor.is_null() {
return Err(LimboError::InternalError(
"VirtualTableCursor: cursor is null".into(),
@@ -165,19 +276,20 @@ impl VirtualTableCursor {
})
}
pub(crate) fn rowid(&self) -> i64 {
fn rowid(&self) -> i64 {
unsafe { (self.implementation.rowid)(self.cursor) }
}
#[tracing::instrument(skip(self))]
pub fn filter(
fn filter(
&self,
idx_num: i32,
idx_str: Option<String>,
arg_count: usize,
args: Vec<limbo_ext::Value>,
args: Vec<Value>,
) -> crate::Result<bool> {
tracing::trace!("xFilter");
let ext_args = args.iter().map(|arg| arg.to_ffi()).collect::<Vec<_>>();
let c_idx_str = idx_str
.map(|s| std::ffi::CString::new(s).unwrap())
.map(|cstr| cstr.into_raw())
@@ -186,12 +298,12 @@ impl VirtualTableCursor {
(self.implementation.filter)(
self.cursor,
arg_count as i32,
args.as_ptr(),
ext_args.as_ptr(),
c_idx_str,
idx_num,
)
};
for arg in args {
for arg in ext_args {
unsafe {
arg.__free_internal_type();
}
@@ -203,12 +315,12 @@ impl VirtualTableCursor {
}
}
pub fn column(&self, column: usize) -> crate::Result<Value> {
fn column(&self, column: usize) -> crate::Result<Value> {
let val = unsafe { (self.implementation.column)(self.cursor, column as u32) };
Value::from_ffi(val)
}
pub fn next(&self) -> crate::Result<bool> {
fn next(&self) -> crate::Result<bool> {
let rc = unsafe { (self.implementation.next)(self.cursor) };
match rc {
ResultCode::OK => Ok(true),
@@ -218,7 +330,7 @@ impl VirtualTableCursor {
}
}
impl Drop for VirtualTableCursor {
impl Drop for ExtVirtualTableCursor {
fn drop(&mut self) {
let result = unsafe { (self.implementation.close)(self.cursor) };
if !result.is_ok() {