Add support for pragma table-valued functions

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;

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::<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()
+    }
+}

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<Vec<ast::Expr>>,
         syms: &SymbolTable,
     ) -> crate::Result<Rc<VirtualTable>> {
-        let ext_args = match args {
+        let module = syms.vtab_modules.get(name);
-            Some(ref args) => vtable_args(args),
+        let (vtab_type, schema) = if let Some(_) = module {
-            None => vec![],
+            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<limbo_ext::Value>,
         syms: &SymbolTable,
     ) -> crate::Result<Rc<VirtualTable>> {
-        let (table, columns) =
+        let module = syms.vtab_modules.get(module_name);
-            ExtVirtualTable::from_args(module_name, args, syms, VTabKind::VirtualTable)?;
+        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,19 +208,16 @@ 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>)> {
+    ) -> crate::Result<(Self, String)> {
-        let module = syms
+        let module = module.ok_or(LimboError::ExtensionError(format!(
-            .vtab_modules
+            "Virtual table module not found: {}",
-            .get(module_name)
+            module_name
-            .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());
+        let vtab = ExtVirtualTable {
-        if let ast::Cmd::Stmt(ast::Stmt::CreateTable { body, .. }) = parser.next()?.ok_or(
+            connection_ptr: RefCell::new(None),
-            LimboError::ParseError("Failed to parse schema from virtual table module".to_string()),
+            implementation: module.implementation.clone(),
-        )? {
+            table_ptr,
-            let columns = columns_from_create_table_body(&body)?;
+        };
-            let vtab = ExtVirtualTable {
+        Ok((vtab, schema))
-                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(),
-        ))
     }
 
     /// Accepts a Weak pointer to the connection that owns the VTable, that the module

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(

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