Add column return support for create and insert statements

tursodatabase · Jan 23, 2025 · 1696534 · 1696534
1 parent 1d9a503
commit 1696534
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 9 deletions.
diff --git a/core/translate/insert.rs b/core/translate/insert.rs
@@ -257,6 +257,10 @@ pub fn translate_insert(
     program.emit_insn(Insn::Goto {
         target_pc: start_offset,
     });
+    program.columns = columns
+        .as_ref()
+        .map(|c| c.iter().map(|r| r.0.clone()).collect())
+        .unwrap_or_default();
 
     Ok(())
 }

diff --git a/core/translate/mod.rs b/core/translate/mod.rs
@@ -394,13 +394,13 @@ fn translate_create_table(
             program.emit_insn(Insn::Goto {
                 target_pc: start_offset,
             });
-
             return Ok(());
         }
         bail_parse_error!("Table {} already exists", tbl_name);
     }
 
     let sql = create_table_body_to_str(&tbl_name, &body);
+    extract_column_names(program, &body)?;
 
     let parse_schema_label = program.allocate_label();
     let init_label = program.allocate_label();
@@ -691,17 +691,21 @@ impl Display for TableFormatter<'_> {
     }
 }
 
+fn extract_column_names(prog: &mut ProgramBuilder, body: &ast::CreateTableBody) -> Result<()> {
+    match body {
+        ast::CreateTableBody::ColumnsAndConstraints { columns, .. } => {
+            prog.columns = columns.iter().map(|c| c.0 .0.clone()).collect::<Vec<_>>();
+            Ok(())
+        }
+        ast::CreateTableBody::AsSelect(_) => {
+            bail_parse_error!("CREATE TABLE AS SELECT not supported yet")
+        }
+    }
+}
+
 fn create_table_body_to_str(tbl_name: &ast::QualifiedName, body: &ast::CreateTableBody) -> String {
     let mut sql = String::new();
     let formatter = TableFormatter { body };
     sql.push_str(format!("CREATE TABLE {} {}", tbl_name.name.0, formatter).as_str());
-    match body {
-        ast::CreateTableBody::ColumnsAndConstraints {
-            columns: _,
-            constraints: _,
-            options: _,
-        } => {}
-        ast::CreateTableBody::AsSelect(_select) => todo!("as select not yet supported"),
-    }
     sql
 }
diff --git a/tests/integration/common.rs b/tests/integration/common.rs
@@ -93,6 +93,25 @@ mod tests {
         assert_eq!(columns.len(), 2);
         assert_eq!(&columns[0], "foo");
         assert_eq!(&columns[1], "bar");
+
+        let stmt = conn.prepare("insert into test (foo, bar) values (1, 2);")?;
+        let columns = stmt.columns();
+        assert_eq!(columns.len(), 2);
+        assert_eq!(&columns[0], "foo");
+        assert_eq!(&columns[1], "bar");
+
+        let stmt = conn.prepare("insert into test (foo) values (1);")?;
+        let columns = stmt.columns();
+        assert_eq!(columns.len(), 1);
+        assert_eq!(&columns[0], "foo");
+
+        let stmt = conn.prepare("create table test_join (foo2, bar2, baz2);")?;
+        let columns = stmt.columns();
+        assert_eq!(columns.len(), 3);
+        assert_eq!(&columns[0], "foo2");
+        assert_eq!(&columns[1], "bar2");
+        assert_eq!(&columns[2], "baz2");
+
         Ok(())
     }
 }