Skip to content

Commit

Permalink
Support returning column names from prepared statement
Browse files Browse the repository at this point in the history
  • Loading branch information
PThorpe92 committed Jan 23, 2025
1 parent a26fc16 commit 545990f
Show file tree
Hide file tree
Showing 5 changed files with 64 additions and 3 deletions.
9 changes: 8 additions & 1 deletion core/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,6 @@ impl Connection {
pub(crate) fn run_cmd(self: &Rc<Connection>, cmd: Cmd) -> Result<Option<Rows>> {
let db = self.db.clone();
let syms: &SymbolTable = &db.syms.borrow();

match cmd {
Cmd::Stmt(stmt) => {
let program = Rc::new(translate::translate(
Expand Down Expand Up @@ -466,6 +465,10 @@ impl Statement {
Ok(Rows::new(stmt))
}

pub fn columns(&self) -> &[String] {
&self.program.columns
}

pub fn parameters(&self) -> &parameters::Parameters {
&self.program.parameters
}
Expand Down Expand Up @@ -513,6 +516,10 @@ impl Rows {
pub fn next_row(&mut self) -> Result<StepResult<'_>> {
self.stmt.step()
}

pub fn columns(&self) -> &[String] {
self.stmt.columns()
}
}

pub(crate) struct SymbolTable {
Expand Down
12 changes: 10 additions & 2 deletions core/translate/emitter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<Vec<_>>();
Ok(())
}

Expand Down Expand Up @@ -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::<Vec<_>>();
Ok(())
}

Expand Down
3 changes: 3 additions & 0 deletions core/vdbe/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ pub struct ProgramBuilder {
// map of instruction index to manual comment (used in EXPLAIN)
comments: HashMap<InsnReference, &'static str>,
pub parameters: Parameters,
pub columns: Vec<String>,
}

#[derive(Debug, Clone)]
Expand Down Expand Up @@ -60,6 +61,7 @@ impl ProgramBuilder {
seekrowid_emitted_bitmask: 0,
comments: HashMap::new(),
parameters: Parameters::new(),
columns: Vec::new(),
}
}

Expand Down Expand Up @@ -352,6 +354,7 @@ impl ProgramBuilder {
parameters: self.parameters,
n_change: Cell::new(0),
change_cnt_on,
columns: self.columns,
}
}
}
1 change: 1 addition & 0 deletions core/vdbe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,7 @@ pub struct Program {
pub auto_commit: bool,
pub n_change: Cell<i64>,
pub change_cnt_on: bool,
pub columns: Vec<String>,
}

impl Program {
Expand Down
42 changes: 42 additions & 0 deletions tests/integration/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
}
}

0 comments on commit 545990f

Please sign in to comment.