diff --git a/core/lib.rs b/core/lib.rs index 1ac907af7..839562f5b 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -634,6 +634,7 @@ pub struct VirtualTable { args: Option>, pub implementation: Rc, columns: Vec, + kind: VTabKind, } impl VirtualTable { @@ -675,6 +676,7 @@ impl VirtualTable { implementation: module.implementation.clone(), columns, args: exprs, + kind, }); return Ok(vtab); } diff --git a/core/translate/schema.rs b/core/translate/schema.rs index 449d1e0e8..c4d570b5c 100644 --- a/core/translate/schema.rs +++ b/core/translate/schema.rs @@ -3,6 +3,7 @@ use std::fmt::Display; use crate::ast; use crate::schema::Schema; use crate::storage::pager::CreateBTreeFlags; +use crate::schema::Table; use crate::translate::ProgramBuilder; use crate::translate::ProgramBuilderOpts; use crate::translate::QueryMode; @@ -543,7 +544,7 @@ pub fn translate_drop_table( approx_num_insns: 30, approx_num_labels: 1, }); - let table = schema.get_btree_table(tbl_name.name.0.as_str()); + let table = schema.get_table(tbl_name.name.0.as_str()); if table.is_none() { if if_exists { let init_label = program.emit_init(); @@ -558,6 +559,7 @@ pub fn translate_drop_table( } bail_parse_error!("No such table: {}", tbl_name.name.0.as_str()); } + let table = table.unwrap(); // safe since we just checked for None let init_label = program.emit_init(); @@ -663,11 +665,31 @@ pub fn translate_drop_table( } // 3. Destroy the table structure - program.emit_insn(Insn::Destroy { - root: table.root_page, - former_root_reg: 0, // no autovacuum (https://www.sqlite.org/opcode.html#Destroy) - is_temp: 0, - }); + match table.as_ref() { + Table::BTree(table) => { + program.emit_insn(Insn::Destroy { + root: table.root_page, + former_root_reg: 0, // no autovacuum (https://www.sqlite.org/opcode.html#Destroy) + is_temp: 0, + }); + } + Table::Virtual(vtab) => { + // From what I see, TableValuedFunction is not stored in the schema as a table. + // But this line here below is a safeguard in case this behavior changes in the future + // And mirrors what SQLite does. + if matches!(vtab.kind, limbo_ext::VTabKind::TableValuedFunction) { + return Err(crate::LimboError::ParseError(format!( + "table {} may not be dropped", + vtab.name + ))); + } + program.emit_insn(Insn::VDestroy { + table_name: vtab.name.clone(), + db: 0, // TODO change this for multiple databases + }); + } + Table::Pseudo(..) => unimplemented!(), + }; let r6 = program.alloc_register(); let r7 = program.alloc_register(); diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 0df9afcc4..40559035d 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -1134,6 +1134,30 @@ pub fn op_vnext( Ok(InsnFunctionStepResult::Step) } +pub fn op_vdestroy( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Rc, + mv_store: Option<&Rc>, +) -> Result { + let Insn::VDestroy { db, table_name } = insn else { + unreachable!("unexpected Insn {:?}", insn) + }; + let Some(conn) = program.connection.upgrade() else { + return Err(crate::LimboError::ExtensionError( + "Failed to upgrade Connection".to_string(), + )); + }; + + { + conn.syms.borrow_mut().vtabs.remove(table_name); + } + + state.pc += 1; + Ok(InsnFunctionStepResult::Step) +} + pub fn op_open_pseudo( program: &Program, state: &mut ProgramState, diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index d4a766d1d..a62ca310b 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -449,6 +449,15 @@ pub fn insn_to_str( 0, "".to_string(), ), + Insn::VDestroy { db, table_name } => ( + "VDestroy", + *db as i32, + 0, + 0, + OwnedValue::build_text(table_name), + 0, + "".to_string(), + ), Insn::OpenPseudo { cursor_id, content_reg, diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index e12293a71..608030195 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -311,6 +311,14 @@ pub enum Insn { pc_if_next: BranchOffset, }, + /// P4 is the name of a virtual table in database P1. Call the xDestroy method of that table. + VDestroy { + /// Name of a virtual table being destroyed + table_name: String, + /// The database within which this virtual table needs to be destroyed (P1). + db: usize, + }, + /// Open a cursor for a pseudo-table that contains a single row. OpenPseudo { cursor_id: CursorID, @@ -856,6 +864,8 @@ impl Insn { Insn::VColumn { .. } => execute::op_vcolumn, Insn::VUpdate { .. } => execute::op_vupdate, Insn::VNext { .. } => execute::op_vnext, + Insn::VDestroy { .. } => execute::op_vdestroy, + Insn::OpenPseudo { .. } => execute::op_open_pseudo, Insn::RewindAsync { .. } => execute::op_rewind_async, Insn::RewindAwait { .. } => execute::op_rewind_await, @@ -920,6 +930,7 @@ impl Insn { Insn::Copy { .. } => execute::op_copy, Insn::CreateBtree { .. } => execute::op_create_btree, Insn::Destroy { .. } => execute::op_destroy, + Insn::DropTable { .. } => execute::op_drop_table, Insn::Close { .. } => execute::op_close, Insn::IsNull { .. } => execute::op_is_null, diff --git a/testing/cli_tests/extensions.py b/testing/cli_tests/extensions.py index 4d289f311..21ef1a7c2 100755 --- a/testing/cli_tests/extensions.py +++ b/testing/cli_tests/extensions.py @@ -526,6 +526,32 @@ def test_vfs(): limbo.quit() +def test_drop_virtual_table(): + ext_path = "target/debug/liblimbo_ext_tests" + limbo = TestLimboShell() + limbo.execute_dot(f".load {ext_path}") + limbo.debug_print( + "create virtual table t using kv_store;", + ) + limbo.run_test_fn(".schema", lambda res: "CREATE VIRTUAL TABLE t" in res) + limbo.run_test_fn( + "insert into t values ('hello', 'world');", + null, + "can insert into kv_store vtable", + ) + limbo.run_test_fn( + "DROP TABLE t;", + null, + "can drop kv_store vtable", + ) + limbo.run_test_fn( + "DROP TABLE t;", + lambda res: "× Parse error: No such table: t" == res, + "should error when drop kv_store vtable", + ) + limbo.quit() + + def test_sqlite_vfs_compat(): sqlite = TestLimboShell( init_commands="", @@ -573,6 +599,7 @@ if __name__ == "__main__": test_vfs() test_sqlite_vfs_compat() test_kv() + test_drop_virtual_table() except Exception as e: print(f"Test FAILED: {e}") cleanup()