diff --git a/core/lib.rs b/core/lib.rs index f093762fb..d75510f7c 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -295,7 +295,6 @@ impl Connection { pub(crate) fn run_cmd(self: &Rc, cmd: Cmd) -> Result> { let db = self.db.clone(); let syms: &SymbolTable = &db.syms.borrow(); - match cmd { Cmd::Stmt(stmt) => { let program = Rc::new(translate::translate( @@ -466,6 +465,10 @@ impl Statement { Ok(Rows::new(stmt)) } + pub fn columns(&self) -> &[String] { + &self.program.columns + } + pub fn parameters(&self) -> ¶meters::Parameters { &self.program.parameters } @@ -513,6 +516,10 @@ impl Rows { pub fn next_row(&mut self) -> Result> { self.stmt.step() } + + pub fn columns(&self) -> &[String] { + self.stmt.columns() + } } pub(crate) struct SymbolTable { diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 13daa85ed..939a287f0 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -175,7 +175,11 @@ fn emit_program_for_select( // Finalize program epilogue(program, init_label, start_offset)?; - + program.columns = plan + .result_columns + .iter() + .map(|rc| rc.name.clone()) + .collect::>(); Ok(()) } @@ -286,7 +290,11 @@ fn emit_program_for_delete( // Finalize program epilogue(program, init_label, start_offset)?; - + program.columns = plan + .result_columns + .iter() + .map(|rc| rc.name.clone()) + .collect::>(); Ok(()) } diff --git a/core/vdbe/builder.rs b/core/vdbe/builder.rs index 08b35d9e3..0af4d1182 100644 --- a/core/vdbe/builder.rs +++ b/core/vdbe/builder.rs @@ -30,6 +30,7 @@ pub struct ProgramBuilder { // map of instruction index to manual comment (used in EXPLAIN) comments: HashMap, pub parameters: Parameters, + pub columns: Vec, } #[derive(Debug, Clone)] @@ -60,6 +61,7 @@ impl ProgramBuilder { seekrowid_emitted_bitmask: 0, comments: HashMap::new(), parameters: Parameters::new(), + columns: Vec::new(), } } @@ -352,6 +354,7 @@ impl ProgramBuilder { parameters: self.parameters, n_change: Cell::new(0), change_cnt_on, + columns: self.columns, } } } diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 98b328b71..27b39fcab 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -284,6 +284,7 @@ pub struct Program { pub auto_commit: bool, pub n_change: Cell, pub change_cnt_on: bool, + pub columns: Vec, } impl Program { diff --git a/tests/integration/common.rs b/tests/integration/common.rs index 86f4b7b3f..07c840b23 100644 --- a/tests/integration/common.rs +++ b/tests/integration/common.rs @@ -67,3 +67,45 @@ pub(crate) fn compare_string(a: &String, b: &String) { } } } + +#[cfg(test)] +mod tests { + use super::TempDatabase; + + #[test] + fn test_statement_columns() -> anyhow::Result<()> { + let _ = env_logger::try_init(); + let tmp_db = + TempDatabase::new("create table test (foo integer, bar integer, baz integer);"); + let conn = tmp_db.connect_limbo(); + + let stmt = conn.prepare("select * from test;")?; + + let columns = stmt.columns(); + assert_eq!(columns.len(), 3); + assert_eq!(&columns[0], "foo"); + assert_eq!(&columns[1], "bar"); + assert_eq!(&columns[2], "baz"); + + let stmt = conn.prepare("select foo, bar from test;")?; + + let columns = stmt.columns(); + assert_eq!(columns.len(), 2); + assert_eq!(&columns[0], "foo"); + assert_eq!(&columns[1], "bar"); + + let stmt = conn.prepare("delete from test;")?; + let columns = stmt.columns(); + assert_eq!(columns.len(), 0); + + let stmt = conn.prepare("insert into test (foo, bar, baz) values (1, 2, 3);")?; + let columns = stmt.columns(); + assert_eq!(columns.len(), 0); + + let stmt = conn.prepare("delete from test where foo = 1")?; + let columns = stmt.columns(); + assert_eq!(columns.len(), 0); + + Ok(()) + } +}