diff --git a/core/lib.rs b/core/lib.rs index b70ad7bcb..793c772b6 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -11,6 +11,7 @@ mod io; mod json; pub mod mvcc; mod parameters; +mod pragma; mod pseudo; pub mod result; mod schema; diff --git a/core/pragma.rs b/core/pragma.rs new file mode 100644 index 000000000..b7a457b76 --- /dev/null +++ b/core/pragma.rs @@ -0,0 +1,246 @@ +use crate::{Connection, LimboError, Statement, StepResult, Value}; +use bitflags::bitflags; +use limbo_sqlite3_parser::ast::PragmaName; +use std::rc::{Rc, Weak}; +use std::str::FromStr; + +bitflags! { + // Flag names match those used in SQLite: + // https://github.com/sqlite/sqlite/blob/b3c1884b65400da85636458298bd77cbbfdfb401/tool/mkpragmatab.tcl#L22-L29 + struct PragmaFlags: u8 { + const NeedSchema = 0x01; + const NoColumns = 0x02; + const NoColumns1 = 0x04; + const ReadOnly = 0x08; + const Result0 = 0x10; + const Result1 = 0x20; + const SchemaOpt = 0x40; + const SchemaReq = 0x80; + } +} + +struct Pragma { + flags: PragmaFlags, + columns: &'static [&'static str], +} + +impl Pragma { + const fn new(flags: PragmaFlags, columns: &'static [&'static str]) -> Self { + Self { flags, columns } + } +} + +fn pragma_for(pragma: PragmaName) -> Pragma { + use PragmaName::*; + + match pragma { + CacheSize => Pragma::new( + PragmaFlags::NeedSchema + | PragmaFlags::Result0 + | PragmaFlags::SchemaReq + | PragmaFlags::NoColumns1, + &["cache_size"], + ), + JournalMode => Pragma::new( + PragmaFlags::NeedSchema | PragmaFlags::Result0 | PragmaFlags::SchemaReq, + &["journal_mode"], + ), + LegacyFileFormat => { + unreachable!("pragma_for() called with LegacyFileFormat, which is unsupported") + } + PageCount => Pragma::new( + PragmaFlags::NeedSchema | PragmaFlags::Result0 | PragmaFlags::SchemaReq, + &["page_count"], + ), + PageSize => Pragma::new( + PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1, + &["page_size"], + ), + SchemaVersion => Pragma::new( + PragmaFlags::NoColumns1 | PragmaFlags::Result0, + &["schema_version"], + ), + TableInfo => Pragma::new( + PragmaFlags::NeedSchema | PragmaFlags::Result1 | PragmaFlags::SchemaOpt, + &["cid", "name", "type", "notnull", "dflt_value", "pk"], + ), + UserVersion => Pragma::new( + PragmaFlags::NoColumns1 | PragmaFlags::Result0, + &["user_version"], + ), + WalCheckpoint => Pragma::new(PragmaFlags::NeedSchema, &["busy", "log", "checkpointed"]), + } +} + +#[derive(Debug, Clone)] +pub(crate) struct PragmaVirtualTable { + 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; + } + if pragma + .flags + .intersects(PragmaFlags::SchemaOpt | PragmaFlags::SchemaReq) + { + sql.push_str(", schema HIDDEN"); + max_arg_count += 1; + } + 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)) + } + + fn no_such_pragma(pragma_name: &str) -> LimboError { + LimboError::ParseError(format!( + "No such table-valued function: pragma_{}", + pragma_name + )) + } + + pub(crate) fn open(&self, conn: Weak) -> crate::Result { + Ok(PragmaVirtualTableCursor { + pragma_name: self.pragma_name.clone(), + pos: 0, + conn: conn + .upgrade() + .ok_or_else(|| LimboError::InternalError("Connection was dropped".into()))?, + stmt: None, + arg: None, + visible_column_count: self.visible_column_count, + max_arg_count: self.max_arg_count, + has_pragma_arg: self.has_pragma_arg, + }) + } +} + +pub struct PragmaVirtualTableCursor { + pragma_name: String, + pos: usize, + conn: Rc, + stmt: Option, + arg: Option, + visible_column_count: usize, + max_arg_count: usize, + has_pragma_arg: bool, +} + +impl PragmaVirtualTableCursor { + pub(crate) fn rowid(&self) -> i64 { + self.pos as i64 + } + + pub(crate) fn next(&mut self) -> crate::Result { + let stmt = self + .stmt + .as_mut() + .ok_or_else(|| LimboError::InternalError("Statement is missing".into()))?; + let result = stmt.step()?; + match result { + StepResult::Done => Ok(false), + _ => { + self.pos += 1; + Ok(true) + } + } + } + + pub(crate) fn column(&self, idx: usize) -> crate::Result { + if idx < self.visible_column_count { + let value = self + .stmt + .as_ref() + .ok_or_else(|| LimboError::InternalError("Statement is missing".into()))? + .row() + .ok_or_else(|| LimboError::InternalError("No row available".into()))? + .get_value(idx) + .clone(); + return Ok(value); + } + + let value = match idx - self.visible_column_count { + 0 => self + .arg + .as_ref() + .map_or(Value::Null, |arg| Value::from_text(arg)), + _ => Value::Null, + }; + Ok(value) + } + + pub(crate) fn filter(&mut self, args: Vec) -> crate::Result { + if args.len() > self.max_arg_count { + return Err(LimboError::ParseError(format!( + "Too many arguments for pragma {}: expected at most {}, got {}", + self.pragma_name, + self.max_arg_count, + args.len() + ))); + } + + let to_text = |v: &Value| v.to_text().map(str::to_owned); + let (arg, schema) = match args.as_slice() { + [arg0] if self.has_pragma_arg => (to_text(arg0), None), + [arg0] => (None, to_text(arg0)), + [arg0, arg1] => (to_text(arg0), to_text(arg1)), + _ => (None, None), + }; + + self.arg = arg; + + if let Some(schema) = schema { + // Schema-qualified PRAGMA statements are not supported yet + return Err(LimboError::ParseError(format!( + "Schema argument is not supported yet (got schema: '{schema}')" + ))); + } + + let mut sql = format!("PRAGMA {}", self.pragma_name); + if let Some(arg) = &self.arg { + sql.push_str(&format!("=\"{}\"", arg)); + } + + self.stmt = Some(self.conn.prepare(sql)?); + + self.next() + } +} diff --git a/core/vtab.rs b/core/vtab.rs index 115ce536d..3297b0346 100644 --- a/core/vtab.rs +++ b/core/vtab.rs @@ -1,3 +1,4 @@ +use crate::pragma::{PragmaVirtualTable, PragmaVirtualTableCursor}; use crate::schema::Column; use crate::util::{columns_from_create_table_body, vtable_args}; use crate::{Connection, LimboError, SymbolTable, Value}; @@ -10,6 +11,7 @@ use std::rc::{Rc, Weak}; #[derive(Debug, Clone)] enum VirtualTableType { + Pragma(PragmaVirtualTable), External(ExtVirtualTable), } @@ -28,18 +30,30 @@ impl VirtualTable { args: Option>, syms: &SymbolTable, ) -> crate::Result> { - let ext_args = match args { - Some(ref args) => vtable_args(args), - None => vec![], + let module = syms.vtab_modules.get(name); + let (vtab_type, schema) = if let Some(_) = module { + let ext_args = match args { + Some(ref args) => vtable_args(args), + None => vec![], + }; + ExtVirtualTable::create(name, module, ext_args, 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 + ))); }; - let (vtab, columns) = - ExtVirtualTable::from_args(name, ext_args, syms, VTabKind::TableValuedFunction)?; + let vtab = VirtualTable { name: name.to_owned(), args, - columns, + columns: Self::resolve_columns(schema)?, kind: VTabKind::TableValuedFunction, - vtab_type: VirtualTableType::External(vtab), + vtab_type, }; Ok(Rc::new(vtab)) } @@ -50,20 +64,35 @@ impl VirtualTable { args: Vec, syms: &SymbolTable, ) -> crate::Result> { - let (table, columns) = - ExtVirtualTable::from_args(module_name, args, syms, VTabKind::VirtualTable)?; + let module = syms.vtab_modules.get(module_name); + let (table, schema) = + ExtVirtualTable::create(module_name, module, args, VTabKind::VirtualTable)?; let vtab = VirtualTable { name: tbl_name.unwrap_or(module_name).to_owned(), args: None, - columns, + columns: Self::resolve_columns(schema)?, kind: VTabKind::VirtualTable, vtab_type: VirtualTableType::External(table), }; Ok(Rc::new(vtab)) } + fn resolve_columns(schema: String) -> crate::Result> { + 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()), + )? { + columns_from_create_table_body(&body) + } else { + Err(LimboError::ParseError( + "Failed to parse schema from virtual table module".to_string(), + )) + } + } + pub(crate) fn open(&self, conn: Weak) -> crate::Result { match &self.vtab_type { + VirtualTableType::Pragma(table) => Ok(VirtualTableCursor::Pragma(table.open(conn)?)), VirtualTableType::External(table) => { Ok(VirtualTableCursor::External(table.open(conn)?)) } @@ -72,12 +101,14 @@ impl VirtualTable { pub(crate) fn update(&self, args: &[Value]) -> crate::Result> { match &self.vtab_type { + VirtualTableType::Pragma(_) => Err(LimboError::ReadOnly), VirtualTableType::External(table) => table.update(args), } } pub(crate) fn destroy(&self) -> crate::Result<()> { match &self.vtab_type { + VirtualTableType::Pragma(_) => Ok(()), VirtualTableType::External(table) => table.destroy(), } } @@ -88,30 +119,40 @@ impl VirtualTable { order_by: &[OrderByInfo], ) -> IndexInfo { match &self.vtab_type { + VirtualTableType::Pragma(_) => { + // SQLite tries to estimate cost and row count for pragma_ TVFs, + // but since Limbo doesn't have cost-based planning yet, this + // estimation is not currently implemented. + Default::default() + } VirtualTableType::External(table) => table.best_index(constraints, order_by), } } } pub enum VirtualTableCursor { + Pragma(PragmaVirtualTableCursor), External(ExtVirtualTableCursor), } impl VirtualTableCursor { pub(crate) fn next(&mut self) -> crate::Result { match self { + VirtualTableCursor::Pragma(cursor) => cursor.next(), VirtualTableCursor::External(cursor) => cursor.next(), } } pub(crate) fn rowid(&self) -> i64 { match self { + VirtualTableCursor::Pragma(cursor) => cursor.rowid(), VirtualTableCursor::External(cursor) => cursor.rowid(), } } pub(crate) fn column(&self, column: usize) -> crate::Result { match self { + VirtualTableCursor::Pragma(cursor) => cursor.column(column), VirtualTableCursor::External(cursor) => cursor.column(column), } } @@ -124,6 +165,7 @@ impl VirtualTableCursor { args: Vec, ) -> crate::Result { match self { + VirtualTableCursor::Pragma(cursor) => cursor.filter(args), VirtualTableCursor::External(cursor) => { cursor.filter(idx_num, idx_str, arg_count, args) } @@ -166,19 +208,16 @@ impl ExtVirtualTable { } /// takes ownership of the provided Args - fn from_args( + fn create( module_name: &str, + module: Option<&Rc>, args: Vec, - syms: &SymbolTable, kind: VTabKind, - ) -> crate::Result<(Self, Vec)> { - let module = syms - .vtab_modules - .get(module_name) - .ok_or(LimboError::ExtensionError(format!( - "Virtual table module not found: {}", - module_name - )))?; + ) -> crate::Result<(Self, String)> { + let module = module.ok_or(LimboError::ExtensionError(format!( + "Virtual table module not found: {}", + module_name + )))?; if kind != module.module_kind { let expected = match kind { VTabKind::VirtualTable => "virtual table", @@ -190,21 +229,12 @@ impl ExtVirtualTable { ))); } 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 = ExtVirtualTable { - connection_ptr: RefCell::new(None), - implementation: module.implementation.clone(), - table_ptr, - }; - return Ok((vtab, columns)); - } - Err(LimboError::ParseError( - "Failed to parse schema from virtual table module".to_string(), - )) + let vtab = ExtVirtualTable { + connection_ptr: RefCell::new(None), + implementation: module.implementation.clone(), + table_ptr, + }; + Ok((vtab, schema)) } /// Accepts a Weak pointer to the connection that owns the VTable, that the module diff --git a/testing/cli_tests/extensions.py b/testing/cli_tests/extensions.py index df97b2537..f996942d3 100755 --- a/testing/cli_tests/extensions.py +++ b/testing/cli_tests/extensions.py @@ -315,7 +315,7 @@ def test_series(): ext_path = "./target/debug/liblimbo_series" limbo.run_test_fn( "SELECT * FROM generate_series(1, 10);", - lambda res: "Virtual table module not found: generate_series" in res, + lambda res: "No such table-valued function: generate_series" in res, ) limbo.execute_dot(f".load {ext_path}") limbo.run_test_fn( diff --git a/testing/pragma.test b/testing/pragma.test index bdc3b1102..c4a478bb3 100755 --- a/testing/pragma.test +++ b/testing/pragma.test @@ -7,10 +7,18 @@ do_execsql_test pragma-cache-size { PRAGMA cache_size } {-2000} +do_execsql_test pragma-function-cache-size { + SELECT * FROM pragma_cache_size() +} {-2000} + do_execsql_test pragma-update-journal-mode-wal { PRAGMA journal_mode=WAL } {wal} +do_execsql_test pragma-function-update-journal-mode { + SELECT * FROM pragma_journal_mode() +} {wal} + do_execsql_test pragma-table-info-equal-syntax { PRAGMA table_info=sqlite_schema } {0|type|TEXT|0||0 @@ -29,10 +37,23 @@ do_execsql_test pragma-table-info-call-syntax { 4|sql|TEXT|0||0 } +do_execsql_test pragma-function-table-info { + SELECT * FROM pragma_table_info('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 } {} +do_execsql_test pragma-function-table-info-invalid-table { + SELECT * FROM pragma_table_info('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 @@ -65,3 +86,49 @@ do_execsql_test_on_specific_db ":memory:" pragma-user-version-float-value { PRAGMA user_version = 10.9; PRAGMA user_version; } {10} + +do_execsql_test pragma-legacy-file-format { + PRAGMA legacy_file_format +} {} + +do_execsql_test_error_content pragma-function-legacy-file-format { + SELECT * FROM pragma_legacy_file_format() +} {"No such table"} + +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 { + SELECT * FROM pragma_wal_checkpoint() +} {"No such table"} + +do_execsql_test pragma-function-nontext-argument { + SELECT * FROM pragma_table_info('sqlite_schema', 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_on_specific_db ":memory:" pragma-function-argument-with-space { + CREATE TABLE "foo bar"(c0); + SELECT * FROM pragma_table_info('foo bar') +} {0|c0||0||0} + +# If the argument passed to the first function call were simply concatenated with the underlying PRAGMA statement, +# we would end up with: PRAGMA table_info='sqlite_schema';CREATE TABLE foo(c0);SELECT 'bar'. Depending on how many +# statements are executed at once, at least one of the following would run: +# - PRAGMA table_info='sqlite_schema'; +# - CREATE TABLE foo(c0); +# - SELECT 'bar'; +# No output means that none of them were executed. +do_execsql_test pragma-function-sql-injection { + SELECT * FROM pragma_table_info('sqlite_schema'';CREATE TABLE foo(c0);SELECT ''bar'); + SELECT * FROM pragma_table_info('foo'); +} {}