From 30ae6538ee34f247fabcf087465394a3ef86d860 Mon Sep 17 00:00:00 2001 From: Piotr Rzysko Date: Tue, 10 Jun 2025 08:54:24 +0200 Subject: [PATCH] Treat table-valued functions as tables With this change, the following two queries are considered equivalent: ```sql SELECT value FROM generate_series(5, 50); SELECT value FROM generate_series WHERE start = 5 AND stop = 50; ``` Arguments passed in parentheses to the virtual table name are now matched to hidden columns. Column references are still not supported as table-valued function arguments. The only difference is that previously, a query like: ```sql SELECT one.value, series.value FROM (SELECT 1 AS value) one, generate_series(one.value, 3) series; ``` would cause a panic. Now, it returns a proper error message instead. Adding support for column references is more nuanced for two main reasons: - We need to ensure that in joins where a TVF depends on other tables, those other tables are processed first. For example, in: ```sql SELECT one.value, series.value FROM generate_series(one.value, 3) series, (SELECT 1 AS value) one; ``` the one table must be processed by the top-level loop, and series must be nested. - For outer joins involving TVFs, the arguments must be treated as ON predicates, not WHERE predicates. --- core/ext/mod.rs | 8 ++ core/pragma.rs | 101 +++++++++-------- core/schema.rs | 3 + core/series.rs | 56 ++++++---- core/translate/main_loop.rs | 181 ++++++++++++++----------------- core/translate/planner.rs | 121 ++++++++++++++++----- core/translate/pragma.rs | 2 +- core/util.rs | 29 ----- core/vtab.rs | 36 +++--- extensions/completion/src/lib.rs | 35 +++--- testing/cli_tests/extensions.py | 46 +++++++- testing/pragma.test | 34 +++++- 12 files changed, 385 insertions(+), 267 deletions(-) diff --git a/core/ext/mod.rs b/core/ext/mod.rs index 8732047b2..b6ab521e9 100644 --- a/core/ext/mod.rs +++ b/core/ext/mod.rs @@ -1,6 +1,7 @@ #[cfg(feature = "fs")] mod dynamic; mod vtab_xconnect; +use crate::vtab::VirtualTable; #[cfg(all(target_os = "linux", feature = "io_uring"))] use crate::UringIO; use crate::{function::ExternalFunc, Connection, Database, LimboError, IO}; @@ -148,6 +149,13 @@ impl Connection { .borrow_mut() .vtab_modules .insert(name.to_string(), vmodule.into()); + if kind == VTabKind::TableValuedFunction { + if let Ok(vtab) = VirtualTable::function(name, &self.syms.borrow()) { + self.schema.borrow_mut().add_virtual_table(vtab); + } else { + return ResultCode::Error; + } + } ResultCode::OK } diff --git a/core/pragma.rs b/core/pragma.rs index 688ccebb8..2b112ed88 100644 --- a/core/pragma.rs +++ b/core/pragma.rs @@ -1,8 +1,8 @@ use crate::{Connection, LimboError, Statement, StepResult, Value}; use bitflags::bitflags; -use std::str::FromStr; -use turso_ext::{ConstraintInfo, ConstraintOp, ConstraintUsage, IndexInfo}; use std::sync::Arc; +use strum::IntoEnumIterator; +use turso_ext::{ConstraintInfo, ConstraintOp, ConstraintUsage, IndexInfo}; use turso_sqlite3_parser::ast::PragmaName; bitflags! { @@ -31,7 +31,7 @@ impl Pragma { } } -pub fn pragma_for(pragma: PragmaName) -> Pragma { +pub fn pragma_for(pragma: &PragmaName) -> Pragma { use PragmaName::*; match pragma { @@ -87,66 +87,65 @@ pub fn pragma_for(pragma: PragmaName) -> Pragma { #[derive(Debug, Clone)] pub(crate) struct PragmaVirtualTable { - pragma_name: String, + pub(crate) pragma_name: String, visible_column_count: usize, max_arg_count: usize, has_pragma_arg: bool, } impl PragmaVirtualTable { - pub(crate) fn create(pragma_name: &str) -> crate::Result<(Self, String)> { - if let Ok(pragma) = PragmaName::from_str(pragma_name) { - if pragma == PragmaName::LegacyFileFormat { - return Err(Self::no_such_pragma(pragma_name)); - } - let pragma = pragma_for(pragma); - if pragma - .flags - .intersects(PragmaFlags::Result0 | PragmaFlags::Result1) - { - let mut max_arg_count = 0; - let mut has_pragma_arg = false; - - let mut sql = String::from("CREATE TABLE x("); - let col_defs = pragma - .columns - .iter() - .map(|col| format!("\"{col}\"")) - .collect::>() - .join(", "); - sql.push_str(&col_defs); - if pragma.flags.contains(PragmaFlags::Result1) { - sql.push_str(", arg HIDDEN"); - max_arg_count += 1; - has_pragma_arg = true; - } + pub(crate) fn functions() -> Vec<(PragmaVirtualTable, String)> { + PragmaName::iter() + .filter(|name| *name != PragmaName::LegacyFileFormat) + .filter_map(|name| { + let pragma = pragma_for(&name); if pragma .flags - .intersects(PragmaFlags::SchemaOpt | PragmaFlags::SchemaReq) + .intersects(PragmaFlags::Result0 | PragmaFlags::Result1) { - sql.push_str(", schema HIDDEN"); - max_arg_count += 1; + Some(Self::create(name.to_string(), pragma)) + } else { + None } - sql.push(')'); - - return Ok(( - PragmaVirtualTable { - pragma_name: pragma_name.to_owned(), - visible_column_count: pragma.columns.len(), - max_arg_count, - has_pragma_arg, - }, - sql, - )); - } - } - Err(Self::no_such_pragma(pragma_name)) + }) + .collect() } - fn no_such_pragma(pragma_name: &str) -> LimboError { - LimboError::ParseError(format!( - "No such table-valued function: pragma_{pragma_name}" - )) + fn create(pragma_name: String, pragma: Pragma) -> (Self, String) { + let mut max_arg_count = 0; + let mut has_pragma_arg = false; + + let mut sql = String::from("CREATE TABLE x("); + let col_defs = pragma + .columns + .iter() + .map(|col| format!("\"{col}\"")) + .collect::>() + .join(", "); + sql.push_str(&col_defs); + if pragma.flags.contains(PragmaFlags::Result1) { + sql.push_str(", arg HIDDEN"); + max_arg_count += 1; + has_pragma_arg = true; + } + if pragma + .flags + .intersects(PragmaFlags::SchemaOpt | PragmaFlags::SchemaReq) + { + sql.push_str(", schema HIDDEN"); + max_arg_count += 1; + } + sql.push(')'); + + ( + PragmaVirtualTable { + pragma_name, + visible_column_count: pragma.columns.len(), + max_arg_count, + has_pragma_arg, + }, + sql, + ) } pub(crate) fn open(&self, conn: Arc) -> crate::Result { diff --git a/core/schema.rs b/core/schema.rs index 90d7a26ff..7cb2011d0 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -43,6 +43,9 @@ impl Schema { SCHEMA_TABLE_NAME.to_string(), Arc::new(Table::BTree(sqlite_schema_table().into())), ); + for function in VirtualTable::builtin_functions() { + tables.insert(function.name.to_owned(), Arc::new(Table::Virtual(function))); + } Self { tables, indexes, diff --git a/core/series.rs b/core/series.rs index d93de7c56..8caec6609 100644 --- a/core/series.rs +++ b/core/series.rs @@ -12,12 +12,12 @@ pub fn register_extension(ext_api: &mut ExtensionApi) { } } -macro_rules! try_option { - ($expr:expr, $err:expr) => { - match $expr { - Some(val) => val, - None => return $err, - } +macro_rules! extract_arg_integer { + ($args:expr, $idx:expr, $unknown_type_default:expr) => { + $args + .get($idx) + .map(|v| v.to_integer().unwrap_or($unknown_type_default)) + .unwrap_or(-1) }; } @@ -153,20 +153,36 @@ impl GenerateSeriesCursor { impl VTabCursor for GenerateSeriesCursor { type Error = ResultCode; - fn filter(&mut self, args: &[Value], _: Option<(&str, i32)>) -> ResultCode { - // args are the start, stop, and step - if args.is_empty() || args.len() > 3 { + fn filter(&mut self, args: &[Value], idx_info: Option<(&str, i32)>) -> ResultCode { + let mut start = -1; + let mut stop = -1; + let mut step = 1; + + if let Some((_, idx_num)) = idx_info { + let mut arg_idx = 0; + // For the semantics of `idx_num`, see the comment in the `best_index` method. + if idx_num & 1 != 0 { + start = extract_arg_integer!(args, arg_idx, -1); + arg_idx += 1; + } + if idx_num & 2 != 0 { + stop = extract_arg_integer!(args, arg_idx, i64::MAX); + arg_idx += 1; + } + if idx_num & 4 != 0 { + step = args + .get(arg_idx) + .map(|v| v.to_integer().unwrap_or(1)) + .unwrap_or(1); + } + } + + if start == -1 { return ResultCode::InvalidArgs; } - let start = try_option!(args[0].to_integer(), ResultCode::InvalidArgs); - let stop = try_option!( - args.get(1).map(|v| v.to_integer().unwrap_or(i64::MAX)), - ResultCode::EOF // Sqlite returns an empty series for wacky args - ); - let mut step = args - .get(2) - .map(|v| v.to_integer().unwrap_or(1)) - .unwrap_or(1); + if stop == -1 { + return ResultCode::EOF; // Sqlite returns an empty series for wacky args + } // Convert zero step to 1, matching SQLite behavior if step == 0 { @@ -300,7 +316,7 @@ mod tests { ]; // Initialize cursor through filter - match cursor.filter(&args, None) { + match cursor.filter(&args, Some(("idx", 1 | 2 | 4))) { ResultCode::OK => (), ResultCode::EOF => return Ok(vec![]), err => return Err(err), @@ -586,7 +602,7 @@ mod tests { ]; // Initialize cursor through filter - cursor.filter(&args, None); + cursor.filter(&args, Some(("idx", 1 | 2 | 4))); let mut rowids = vec![]; while !cursor.eof() { diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index ed9a5bf20..b4868fa50 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -1,4 +1,3 @@ -use turso_ext::VTabKind; use turso_sqlite3_parser::ast::{self, SortOrder}; use std::sync::Arc; @@ -442,116 +441,94 @@ pub fn open_loop( program.preassign_label_to_next_insn(loop_start); } Table::Virtual(vtab) => { - let (start_reg, count, maybe_idx_str, maybe_idx_int) = - if vtab.kind.eq(&VTabKind::VirtualTable) { - // Virtual‑table (non‑TVF) modules can receive constraints via xBestIndex. - // They return information with which to pass to VFilter operation. - // We forward every predicate that touches vtab columns. - // - // vtab.col = literal (always usable) - // vtab.col = outer_table.col (usable, because outer_table is already positioned) - // vtab.col = later_table.col (forwarded with usable = false) - // - // xBestIndex decides which ones it wants by setting argvIndex and whether the - // core layer may omit them (omit = true). - // We then materialise the RHS/LHS into registers before issuing VFilter. - let converted_constraints = predicates - .iter() - .enumerate() - .filter(|(_, p)| p.should_eval_at_loop(join_index, join_order)) - .filter_map(|(i, p)| { - // Build ConstraintInfo from the predicates - convert_where_to_vtab_constraint( - p, - joined_table_index, - i, - join_order, - ) - .unwrap_or(None) - }) - .collect::>(); - // TODO: get proper order_by information to pass to the vtab. - // maybe encode more info on t_ctx? we need: [col_idx, is_descending] - let index_info = vtab.best_index(&converted_constraints, &[]); + let (start_reg, count, maybe_idx_str, maybe_idx_int) = { + // Virtual‑table modules can receive constraints via xBestIndex. + // They return information with which to pass to VFilter operation. + // We forward every predicate that touches vtab columns. + // + // vtab.col = literal (always usable) + // vtab.col = outer_table.col (usable, because outer_table is already positioned) + // vtab.col = later_table.col (forwarded with usable = false) + // + // xBestIndex decides which ones it wants by setting argvIndex and whether the + // core layer may omit them (omit = true). + // We then materialise the RHS/LHS into registers before issuing VFilter. + let converted_constraints = predicates + .iter() + .enumerate() + .filter(|(_, p)| p.should_eval_at_loop(join_index, join_order)) + .filter_map(|(i, p)| { + // Build ConstraintInfo from the predicates + convert_where_to_vtab_constraint( + p, + joined_table_index, + i, + join_order, + ) + .unwrap_or(None) + }) + .collect::>(); + // TODO: get proper order_by information to pass to the vtab. + // maybe encode more info on t_ctx? we need: [col_idx, is_descending] + let index_info = vtab.best_index(&converted_constraints, &[]); - // Determine the number of VFilter arguments (constraints with an argv_index). - let args_needed = index_info - .constraint_usages - .iter() - .filter(|u| u.argv_index.is_some()) - .count(); - let start_reg = program.alloc_registers(args_needed); + // Determine the number of VFilter arguments (constraints with an argv_index). + let args_needed = index_info + .constraint_usages + .iter() + .filter(|u| u.argv_index.is_some()) + .count(); + let start_reg = program.alloc_registers(args_needed); - // For each constraint used by best_index, translate the opposite side. - for (i, usage) in index_info.constraint_usages.iter().enumerate() { - if let Some(argv_index) = usage.argv_index { - if let Some(cinfo) = converted_constraints.get(i) { - let (pred_idx, is_rhs) = cinfo.unpack_plan_info(); - if let ast::Expr::Binary(lhs, _, rhs) = - &predicates[pred_idx].expr - { - // translate the opposite side of the referenced vtab column - let expr = if is_rhs { lhs } else { rhs }; - // argv_index is 1-based; adjust to get the proper register offset. - if argv_index == 0 { - // invalid since argv_index is 1-based - continue; - } - let target_reg = - start_reg + (argv_index - 1) as usize; - translate_expr( - program, - Some(table_references), - expr, - target_reg, - &t_ctx.resolver, - )?; - if cinfo.usable && usage.omit { - predicates[pred_idx].consumed.set(true); - } + // For each constraint used by best_index, translate the opposite side. + for (i, usage) in index_info.constraint_usages.iter().enumerate() { + if let Some(argv_index) = usage.argv_index { + if let Some(cinfo) = converted_constraints.get(i) { + let (pred_idx, is_rhs) = cinfo.unpack_plan_info(); + if let ast::Expr::Binary(lhs, _, rhs) = + &predicates[pred_idx].expr + { + // translate the opposite side of the referenced vtab column + let expr = if is_rhs { lhs } else { rhs }; + // argv_index is 1-based; adjust to get the proper register offset. + if argv_index == 0 { + // invalid since argv_index is 1-based + continue; + } + let target_reg = start_reg + (argv_index - 1) as usize; + translate_expr( + program, + Some(table_references), + expr, + target_reg, + &t_ctx.resolver, + )?; + if cinfo.usable && usage.omit { + predicates[pred_idx].consumed.set(true); } } } } + } - // If best_index provided an idx_str, translate it. - let maybe_idx_str = if let Some(idx_str) = index_info.idx_str { - let reg = program.alloc_register(); - program.emit_insn(Insn::String8 { - dest: reg, - value: idx_str, - }); - Some(reg) - } else { - None - }; - ( - start_reg, - args_needed, - maybe_idx_str, - Some(index_info.idx_num), - ) + // If best_index provided an idx_str, translate it. + let maybe_idx_str = if let Some(idx_str) = index_info.idx_str { + let reg = program.alloc_register(); + program.emit_insn(Insn::String8 { + dest: reg, + value: idx_str, + }); + Some(reg) } else { - // For table-valued functions: translate the table args. - let args = match vtab.args.as_ref() { - Some(args) => args, - None => &vec![], - }; - let start_reg = program.alloc_registers(args.len()); - let mut cur_reg = start_reg; - for arg in args { - let reg = cur_reg; - cur_reg += 1; - let _ = translate_expr( - program, - Some(table_references), - arg, - reg, - &t_ctx.resolver, - )?; - } - (start_reg, args.len(), None, None) + None }; + ( + start_reg, + args_needed, + maybe_idx_str, + Some(index_info.idx_num), + ) + }; // Emit VFilter with the computed arguments. program.emit_insn(Insn::VFilter { diff --git a/core/translate/planner.rs b/core/translate/planner.rs index ea84062e0..8e2f8a501 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -19,6 +19,7 @@ use crate::{ vdbe::{builder::TableRefIdCounter, BranchOffset}, Result, }; +use turso_sqlite3_parser::ast::Literal::Null; use turso_sqlite3_parser::ast::{ self, As, Expr, FromClause, JoinType, Limit, Materialized, QualifiedName, TableInternalId, UnaryOperator, With, @@ -243,6 +244,7 @@ fn parse_from_clause_table( schema: &Schema, table: ast::SelectTable, table_references: &mut TableReferences, + out_where_clause: &mut Vec, ctes: &mut Vec, syms: &SymbolTable, table_ref_counter: &mut TableRefIdCounter, @@ -253,8 +255,10 @@ fn parse_from_clause_table( table_references, ctes, table_ref_counter, + out_where_clause, qualified_name, maybe_alias, + None, ), ast::SelectTable::Select(subselect, maybe_alias) => { let Plan::Select(subplan) = prepare_select_plan( @@ -286,42 +290,30 @@ fn parse_from_clause_table( )); Ok(()) } - ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => { - let normalized_name = &normalize_ident(qualified_name.name.0.as_str()); - let vtab = crate::VirtualTable::function(normalized_name, maybe_args, syms)?; - let alias = maybe_alias - .as_ref() - .map(|a| match a { - ast::As::As(id) => id.0.clone(), - ast::As::Elided(id) => id.0.clone(), - }) - .unwrap_or(normalized_name.to_string()); - - table_references.add_joined_table(JoinedTable { - op: Operation::Scan { - iter_dir: IterationDirection::Forwards, - index: None, - }, - join_info: None, - table: Table::Virtual(vtab), - identifier: alias, - internal_id: table_ref_counter.next(), - col_used_mask: ColumnUsedMask::default(), - }); - - Ok(()) - } + ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => parse_table( + schema, + table_references, + ctes, + table_ref_counter, + out_where_clause, + qualified_name, + maybe_alias, + maybe_args, + ), _ => todo!(), } } +#[allow(clippy::too_many_arguments)] fn parse_table( schema: &Schema, table_references: &mut TableReferences, ctes: &mut Vec, table_ref_counter: &mut TableRefIdCounter, + out_where_clause: &mut Vec, qualified_name: QualifiedName, maybe_alias: Option, + maybe_args: Option>, ) -> Result<()> { let normalized_qualified_name = normalize_ident(qualified_name.name.0.as_str()); // Check if the FROM clause table is referring to a CTE in the current scope. @@ -343,7 +335,16 @@ fn parse_table( ast::As::Elided(id) => id, }) .map(|a| a.0); + let internal_id = table_ref_counter.next(); let tbl_ref = if let Table::Virtual(tbl) = table.as_ref() { + if let Some(args) = maybe_args { + transform_args_into_where_terms( + args, + internal_id, + out_where_clause, + table.as_ref(), + )?; + } Table::Virtual(tbl.clone()) } else if let Table::BTree(table) = table.as_ref() { Table::BTree(table.clone()) @@ -359,7 +360,7 @@ fn parse_table( }, table: tbl_ref, identifier: alias.unwrap_or(normalized_qualified_name), - internal_id: table_ref_counter.next(), + internal_id, join_info: None, col_used_mask: ColumnUsedMask::default(), }); @@ -395,6 +396,72 @@ fn parse_table( crate::bail_parse_error!("no such table: {}", normalized_qualified_name); } +fn transform_args_into_where_terms( + args: Vec, + internal_id: TableInternalId, + out_where_clause: &mut Vec, + table: &Table, +) -> Result<()> { + let mut args_iter = args.into_iter(); + let mut hidden_count = 0; + for (i, col) in table.columns().iter().enumerate() { + if !col.hidden { + continue; + } + hidden_count += 1; + + if let Some(arg_expr) = args_iter.next() { + if contains_column_reference(&arg_expr)? { + crate::bail_parse_error!( + "Column references are not supported as table-valued function arguments yet" + ); + } + let column_expr = Expr::Column { + database: None, + table: internal_id, + column: i, + is_rowid_alias: col.is_rowid_alias, + }; + let expr = match arg_expr { + Expr::Literal(Null) => Expr::IsNull(Box::new(column_expr)), + other => Expr::Binary( + Box::new(column_expr), + ast::Operator::Equals, + Box::new(other), + ), + }; + out_where_clause.push(WhereTerm { + expr, + from_outer_join: None, + consumed: Cell::new(false), + }); + } + } + + if args_iter.next().is_some() { + return Err(crate::LimboError::ParseError(format!( + "Too many arguments for {}: expected at most {}, got {}", + table.get_name(), + hidden_count, + hidden_count + 1 + args_iter.count() + ))); + } + + Ok(()) +} + +fn contains_column_reference(top_level_expr: &Expr) -> Result { + let mut contains = false; + walk_expr(top_level_expr, &mut |expr: &Expr| -> Result { + match expr { + Expr::Id(_) | Expr::Qualified(_, _) | Expr::Column { .. } => contains = true, + _ => {} + }; + Ok(WalkControl::Continue) + })?; + Ok(contains) +} + pub fn parse_from( schema: &Schema, mut from: Option, @@ -485,6 +552,7 @@ pub fn parse_from( schema, select_owned, table_references, + out_where_clause, &mut ctes_as_subqueries, syms, table_ref_counter, @@ -723,6 +791,7 @@ fn parse_join( schema, table, table_references, + out_where_clause, ctes, syms, table_ref_counter, diff --git a/core/translate/pragma.rs b/core/translate/pragma.rs index eba48e517..f5a32192c 100644 --- a/core/translate/pragma.rs +++ b/core/translate/pragma.rs @@ -386,7 +386,7 @@ fn query_pragma( translate_integrity_check(schema, &mut program)?; } PragmaName::UnstableCaptureDataChangesConn => { - let pragma = pragma_for(pragma); + let pragma = pragma_for(&pragma); let second_column = program.alloc_register(); let opts = connection.get_capture_data_changes(); program.emit_string8(opts.mode_name().to_string(), register); diff --git a/core/util.rs b/core/util.rs index 0204357ac..d735678cc 100644 --- a/core/util.rs +++ b/core/util.rs @@ -1074,35 +1074,6 @@ pub fn parse_pragma_bool(expr: &Expr) -> Result { )) } -// for TVF's we need these at planning time so we cannot emit translate_expr -pub fn vtable_args(args: &[ast::Expr]) -> Vec { - let mut vtable_args = Vec::new(); - for arg in args { - match arg { - Expr::Literal(lit) => match lit { - Literal::Numeric(i) => { - if i.contains('.') { - vtable_args.push(turso_ext::Value::from_float(i.parse().unwrap())); - } else { - vtable_args.push(turso_ext::Value::from_integer(i.parse().unwrap())); - } - } - Literal::String(s) => { - vtable_args.push(turso_ext::Value::from_text(s.clone())); - } - Literal::Blob(b) => { - vtable_args.push(turso_ext::Value::from_blob(b.as_bytes().into())); - } - _ => { - vtable_args.push(turso_ext::Value::null()); - } - }, - _ => vtable_args.push(turso_ext::Value::null()), - } - } - vtable_args -} - #[cfg(test)] pub mod tests { use super::*; diff --git a/core/vtab.rs b/core/vtab.rs index b0e6e67c2..4fbbca249 100644 --- a/core/vtab.rs +++ b/core/vtab.rs @@ -1,6 +1,6 @@ use crate::pragma::{PragmaVirtualTable, PragmaVirtualTableCursor}; use crate::schema::Column; -use crate::util::{columns_from_create_table_body, vtable_args}; +use crate::util::columns_from_create_table_body; use crate::{Connection, LimboError, SymbolTable, Value}; use fallible_iterator::FallibleIterator; use std::cell::RefCell; @@ -19,29 +19,33 @@ enum VirtualTableType { #[derive(Clone, Debug)] pub struct VirtualTable { pub(crate) name: String, - pub(crate) args: Option>, pub(crate) columns: Vec, pub(crate) kind: VTabKind, vtab_type: VirtualTableType, } impl VirtualTable { - pub(crate) fn function( - name: &str, - args: Option>, - syms: &SymbolTable, - ) -> crate::Result> { + pub(crate) fn builtin_functions() -> Vec> { + PragmaVirtualTable::functions() + .into_iter() + .map(|(tab, schema)| { + let vtab = VirtualTable { + name: format!("pragma_{}", tab.pragma_name), + columns: Self::resolve_columns(schema) + .expect("built-in function schema resolution should not fail"), + kind: VTabKind::TableValuedFunction, + vtab_type: VirtualTableType::Pragma(tab), + }; + Rc::new(vtab) + }) + .collect() + } + + pub(crate) fn function(name: &str, syms: &SymbolTable) -> crate::Result> { let module = syms.vtab_modules.get(name); let (vtab_type, schema) = if module.is_some() { - let ext_args = match args { - Some(ref args) => vtable_args(args), - None => vec![], - }; - ExtVirtualTable::create(name, module, ext_args, VTabKind::TableValuedFunction) + ExtVirtualTable::create(name, module, Vec::new(), VTabKind::TableValuedFunction) .map(|(vtab, columns)| (VirtualTableType::External(vtab), columns))? - } else if let Some(pragma_name) = name.strip_prefix("pragma_") { - PragmaVirtualTable::create(pragma_name) - .map(|(vtab, columns)| (VirtualTableType::Pragma(vtab), columns))? } else { return Err(LimboError::ParseError(format!( "No such table-valued function: {name}" @@ -50,7 +54,6 @@ impl VirtualTable { let vtab = VirtualTable { name: name.to_owned(), - args, columns: Self::resolve_columns(schema)?, kind: VTabKind::TableValuedFunction, vtab_type, @@ -69,7 +72,6 @@ impl VirtualTable { ExtVirtualTable::create(module_name, module, args, VTabKind::VirtualTable)?; let vtab = VirtualTable { name: tbl_name.unwrap_or(module_name).to_owned(), - args: None, columns: Self::resolve_columns(schema)?, kind: VTabKind::VirtualTable, vtab_type: VirtualTableType::External(table), diff --git a/extensions/completion/src/lib.rs b/extensions/completion/src/lib.rs index da784cf43..9a93b65df 100644 --- a/extensions/completion/src/lib.rs +++ b/extensions/completion/src/lib.rs @@ -15,12 +15,13 @@ register_extension! { vtabs: { CompletionVTabModule } } -macro_rules! try_option { - ($expr:expr, $err:expr) => { - match $expr { - Some(val) => val, - None => return $err, - } +macro_rules! extract_arg_text { + ($args:expr, $idx:expr) => { + $args + .get($idx) + .map(|v| v.to_text().unwrap_or("")) + .unwrap_or("") + .to_string() }; } @@ -170,19 +171,21 @@ impl CompletionCursor { impl VTabCursor for CompletionCursor { type Error = ResultCode; - fn filter(&mut self, args: &[Value], _: Option<(&str, i32)>) -> ResultCode { - if args.is_empty() || args.len() > 2 { - return ResultCode::InvalidArgs; - } + fn filter(&mut self, args: &[Value], idx_info: Option<(&str, i32)>) -> ResultCode { self.reset(); - let prefix = try_option!(args[0].to_text(), ResultCode::InvalidArgs); - let wholeline = args.get(1).map(|v| v.to_text().unwrap_or("")).unwrap_or(""); + if let Some((_, idx_num)) = idx_info { + let mut arg_idx = 0; + // For the semantics of `idx_num`, see the comment in the `best_index` method. + if idx_num & 1 != 0 { + self.prefix = extract_arg_text!(args, arg_idx); + arg_idx += 1; + } + if idx_num & 2 != 0 { + self.line = extract_arg_text!(args, arg_idx); + } + } - self.line = wholeline.to_string(); - self.prefix = prefix.to_string(); - - // Currently best index is not implemented so the correct arg parsing is not done here if !self.line.is_empty() && self.prefix.is_empty() { let mut i = self.line.len(); while let Some(ch) = self.line.chars().next() { diff --git a/testing/cli_tests/extensions.py b/testing/cli_tests/extensions.py index e82792ea2..f0fc950ad 100755 --- a/testing/cli_tests/extensions.py +++ b/testing/cli_tests/extensions.py @@ -316,13 +316,41 @@ def _test_series(limbo: TestTursoShell): "SELECT * FROM generate_series(1, 10);", lambda res: res == "1\n2\n3\n4\n5\n6\n7\n8\n9\n10", ) + limbo.run_test_fn( + "SELECT * FROM generate_series WHERE start = 1 AND stop = 10;", + lambda res: res == "1\n2\n3\n4\n5\n6\n7\n8\n9\n10", + ) + limbo.run_test_fn( + "SELECT * FROM generate_series(1, 10) WHERE value < 5;", + lambda res: res == "1\n2\n3\n4", + ) + limbo.run_test_fn( + "SELECT * FROM generate_series WHERE start = 1 AND stop = 10 AND value < 5;", + lambda res: res == "1\n2\n3\n4", + ) + limbo.run_test_fn( + "SELECT * FROM generate_series WHERE start = 1 AND stop = 10 AND start = 5;", + lambda res: res == "", + ) + limbo.run_test_fn( + "SELECT * FROM generate_series WHERE start = 1 AND stop = 10 AND start > 5;", + lambda res: res == "", + ) + limbo.run_test_fn( + "SELECT * FROM generate_series;", + lambda res: "Invalid Argument" in res or 'first argument to "generate_series()" missing or unusable' in res, + ) limbo.run_test_fn( "SELECT * FROM generate_series(1, 10, 2);", lambda res: res == "1\n3\n5\n7\n9", ) + limbo.run_test_fn( + "SELECT * FROM generate_series WHERE start = 1 AND stop = 10 AND step = 2;", + lambda res: res == "1\n3\n5\n7\n9", + ) limbo.run_test_fn( "SELECT * FROM generate_series(1, 10, 2, 3);", - lambda res: "Invalid Argument" in res or "too many arguments" in res, + lambda res: "too many arguments" in res.lower(), ) limbo.run_test_fn( "SELECT * FROM generate_series(10, 1, -2);", @@ -949,6 +977,22 @@ def _test_hidden_columns(exec_name, ext_path): "SELECT * FROM r NATURAL JOIN l NATURAL JOIN r;", lambda res: "comment0|2|3" == res, ) + limbo.run_test_fn( + "SELECT * FROM (SELECT * FROM l JOIN r USING(key, value)) JOIN r USING(comment, key, value);", + lambda res: "2|3|comment0" == res, + ) + limbo.run_test_fn( + "SELECT * FROM (SELECT * FROM l NATURAL JOIN r) JOIN r USING(comment, key, value);", + lambda res: "2|3|comment0" == res, + ) + limbo.run_test_fn( + "SELECT * FROM l JOIN r USING(key, value) JOIN r USING(comment, key, value);", + lambda res: "" == res, + ) + limbo.run_test_fn( + "SELECT * FROM l NATURAL JOIN r JOIN r USING(comment, key, value);", + lambda res: "" == res, + ) limbo.quit() diff --git a/testing/pragma.test b/testing/pragma.test index ff42424bf..6e6b42d93 100755 --- a/testing/pragma.test +++ b/testing/pragma.test @@ -83,6 +83,15 @@ do_execsql_test pragma-function-table-info { 4|sql|TEXT|0||0 } +do_execsql_test pragma-vtab-table-info { + SELECT * FROM pragma_table_info WHERE arg = 'sqlite_schema' +} {0|type|TEXT|0||0 +1|name|TEXT|0||0 +2|tbl_name|TEXT|0||0 +3|rootpage|INT|0||0 +4|sql|TEXT|0||0 +} + do_execsql_test pragma-table-info-invalid-table { PRAGMA table_info=pekka } {} @@ -91,6 +100,10 @@ do_execsql_test pragma-function-table-info-invalid-table { SELECT * FROM pragma_table_info('pekka') } {} +do_execsql_test pragma-vtab-table-info-invalid-table { + SELECT * FROM pragma_table_info WHERE arg = 'pekka' +} {} + # temporarily skip this test case. The issue is detailed in #1407 #do_execsql_test_on_specific_db ":memory:" pragma-page-count-empty { # PRAGMA page_count @@ -128,17 +141,17 @@ do_execsql_test pragma-legacy-file-format { PRAGMA legacy_file_format } {} -do_execsql_test_error_content pragma-function-legacy-file-format { +do_execsql_test_error pragma-function-legacy-file-format { SELECT * FROM pragma_legacy_file_format() -} {"No such table"} +} {(no such table|Table.*not found)} do_execsql_test_error_content pragma-function-too-many-arguments { SELECT * FROM pragma_table_info('sqlite_schema', 'main', 'arg3') } {"Too many arguments"} -do_execsql_test_error_content pragma-function-update { +do_execsql_test_error pragma-function-update { SELECT * FROM pragma_wal_checkpoint() -} {"No such table"} +} {(no such table|Table.*not found)} do_execsql_test pragma-function-nontext-argument { SELECT * FROM pragma_table_info('sqlite_schema', NULL); @@ -149,10 +162,23 @@ do_execsql_test pragma-function-nontext-argument { 4|sql|TEXT|0||0 } +do_execsql_test pragma-vtab-nontext-argument { + SELECT * FROM pragma_table_info WHERE arg ='sqlite_schema' AND schema IS NULL; +} {0|type|TEXT|0||0 +1|name|TEXT|0||0 +2|tbl_name|TEXT|0||0 +3|rootpage|INT|0||0 +4|sql|TEXT|0||0 +} + do_execsql_test pragma-function-no-arguments { SELECT * FROM pragma_table_info(); } {} +do_execsql_test pragma-vtab-no-arguments { + SELECT * FROM pragma_table_info; +} {} + do_execsql_test_on_specific_db ":memory:" pragma-function-argument-with-space { CREATE TABLE "foo bar"(c0); SELECT * FROM pragma_table_info('foo bar')