Merge 'Support returning column names from prepared statement' from P…

…reston Thorpe

Closes #764
Go driver needs support for fetching the column names after `prepare`

Reviewed-by: Sonny (@sonhmai)

Closes #766

core/lib.rs

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

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

core/vdbe/builder.rs

@@ -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)]
@@ -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,
         }
     }
 }

core/vdbe/mod.rs

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

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