aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2025-12-18 17:14:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add support for pragma table-valued functions

Browse Source
This commit is contained in:
Piotr Rzysko
2025-05-30 07:47:46 +02:00
parent 4d35e36b77
commit d1d8ead475
5 changed files with 380 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
core/lib.rs
View File

@@ -11,6 +11,7 @@ mod io;
mod json;
pub mod mvcc;
mod parameters;
mod pragma;
mod pseudo;
pub mod result;
mod schema;

246
core/pragma.rs Normal file
View File

@@ -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::<Vec<_>>()
.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<Connection>) -> crate::Result<PragmaVirtualTableCursor> {
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<Connection>,
stmt: Option<Statement>,
arg: Option<String>,
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<bool> {
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<Value> {
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<Value>) -> crate::Result<bool> {
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()
}
}

78
core/vtab.rs
View File

@@ -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<Vec<ast::Expr>>,
syms: &SymbolTable,
) -> crate::Result<Rc<VirtualTable>> {
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![],
};
let (vtab, columns) =
ExtVirtualTable::from_args(name, ext_args, syms, VTabKind::TableValuedFunction)?;
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 = 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<limbo_ext::Value>,
syms: &SymbolTable,
) -> crate::Result<Rc<VirtualTable>> {
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<Vec<Column>> {
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<Connection>) -> crate::Result<VirtualTableCursor> {
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<Option<i64>> {
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<bool> {
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<Value> {
match self {
VirtualTableCursor::Pragma(cursor) => cursor.column(column),
VirtualTableCursor::External(cursor) => cursor.column(column),
}
}
@@ -124,6 +165,7 @@ impl VirtualTableCursor {
args: Vec<Value>,
) -> crate::Result<bool> {
match self {
VirtualTableCursor::Pragma(cursor) => cursor.filter(args),
VirtualTableCursor::External(cursor) => {
cursor.filter(idx_num, idx_str, arg_count, args)
}
@@ -166,16 +208,13 @@ impl ExtVirtualTable {
}
/// takes ownership of the provided Args
fn from_args(
fn create(
module_name: &str,
module: Option<&Rc<crate::ext::VTabImpl>>,
args: Vec<limbo_ext::Value>,
syms: &SymbolTable,
kind: VTabKind,
) -> crate::Result<(Self, Vec<Column>)> {
let module = syms
.vtab_modules
.get(module_name)
.ok_or(LimboError::ExtensionError(format!(
) -> crate::Result<(Self, String)> {
let module = module.ok_or(LimboError::ExtensionError(format!(
"Virtual table module not found: {}",
module_name
)))?;
@@ -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(),
))
Ok((vtab, schema))
}
/// Accepts a Weak pointer to the connection that owns the VTable, that the module

2
testing/cli_tests/extensions.py
View File

@@ -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(

67
testing/pragma.test
View File

@@ -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');
} {}