Support multi-statement queries (#598)

Inspired by the [Go
client](https://github.com/mlafeldt/go-duckdb/blob/2227e78d82a5e48b320f412cae5e572a0620d98e/connection.go#L203-L235),
this enhances `prepare` to handle both single and multi-statement
queries.

Among other things, this enables the use of `PIVOT` (which internally
expands to multiple statements) and also enables multiple statements per
line in the example REPL. 🎉

Fixes #425
Fixes #310

Author: mlafeldt

crates/duckdb/examples/repl.rs

@@ -132,28 +132,14 @@ impl SqlRepl {
     }
 
     fn execute_sql(&self, sql: &str) -> DuckResult<()> {
-        // Check if it's a statement that returns results
-        let sql_upper = sql.trim().to_uppercase();
-        let is_query = sql_upper.starts_with("SELECT")
-            || sql_upper.starts_with("FROM")
-            || sql_upper.starts_with("SHOW")
-            || sql_upper.starts_with("DESCRIBE")
-            || sql_upper.starts_with("EXPLAIN")
-            || sql_upper.starts_with("PRAGMA")
-            || sql_upper.starts_with("WITH");
-
-        if is_query {
-            let mut stmt = self.conn.prepare(sql)?;
-            let rbs: Vec<RecordBatch> = stmt.query_arrow([])?.collect();
-
-            if rbs.is_empty() || rbs[0].num_rows() == 0 {
-                println!("No results returned.");
-            } else {
-                self.print_records(&rbs);
-            }
-        } else {
-            // Execute non-query statements
-            self.conn.execute_batch(sql)?;
+        let mut stmt = self.conn.prepare(sql)?;
+        let rbs: Vec<RecordBatch> = stmt.query_arrow([])?.collect();
+
+        // NOTE: When executing multi-statement queries (e.g., "SELECT 1; SELECT 2;"),
+        // only the result from the final statement will be displayed. This differs from
+        // the DuckDB CLI which shows results from all statements.
+        if !rbs.is_empty() && rbs[0].num_rows() > 0 {
+            self.print_records(&rbs);
         }
 
         Ok(())

crates/duckdb/src/error.rs

@@ -286,3 +286,29 @@ pub fn result_from_duckdb_arrow(code: ffi::duckdb_state, mut out: ffi::duckdb_ar
         error_from_duckdb_code(code, message)
     }
 }
+
+#[cold]
+#[inline]
+pub fn result_from_duckdb_extract(
+    num_statements: ffi::idx_t,
+    mut extracted: ffi::duckdb_extracted_statements,
+) -> Result<()> {
+    if num_statements > 0 {
+        return Ok(());
+    }
+    unsafe {
+        let message = if extracted.is_null() {
+            Some("extracted statements are null".to_string())
+        } else {
+            let c_err = ffi::duckdb_extract_statements_error(extracted);
+            let message = if c_err.is_null() {
+                None
+            } else {
+                Some(CStr::from_ptr(c_err).to_string_lossy().to_string())
+            };
+            ffi::duckdb_destroy_extracted(&mut extracted);
+            message
+        };
+        error_from_duckdb_code(ffi::DuckDBError, message)
+    }
+}

crates/duckdb/src/inner_connection.rs

@@ -8,7 +8,10 @@ use std::{
 
 use super::{ffi, Appender, Config, Connection, Result};
 use crate::{
-    error::{result_from_duckdb_appender, result_from_duckdb_arrow, result_from_duckdb_prepare, Error},
+    error::{
+        result_from_duckdb_appender, result_from_duckdb_arrow, result_from_duckdb_extract, result_from_duckdb_prepare,
+        Error,
+    },
     raw_statement::RawStatement,
     statement::Statement,
 };
@@ -93,11 +96,68 @@ impl InnerConnection {
     }
 
     pub fn prepare<'a>(&mut self, conn: &'a Connection, sql: &str) -> Result<Statement<'a>> {
-        let mut c_stmt: ffi::duckdb_prepared_statement = ptr::null_mut();
         let c_str = CString::new(sql).unwrap();
-        let r = unsafe { ffi::duckdb_prepare(self.con, c_str.as_ptr() as *const c_char, &mut c_stmt) };
-        result_from_duckdb_prepare(r, c_stmt)?;
-        Ok(Statement::new(conn, unsafe { RawStatement::new(c_stmt) }))
+
+        // Extract statements (handles both single and multi-statement queries)
+        let mut extracted = ptr::null_mut();
+        let num_stmts =
+            unsafe { ffi::duckdb_extract_statements(self.con, c_str.as_ptr() as *const c_char, &mut extracted) };
+        result_from_duckdb_extract(num_stmts, extracted)?;
+
+        // Auto-cleanup on drop
+        let _guard = ExtractedStatementsGuard(extracted);
+
+        // Execute all intermediate statements
+        for i in 0..num_stmts - 1 {
+            self.execute_extracted_statement(extracted, i)?;
+        }
+
+        // Prepare and return final statement
+        let final_stmt = self.prepare_extracted_statement(extracted, num_stmts - 1)?;
+        Ok(Statement::new(conn, unsafe { RawStatement::new(final_stmt) }))
+    }
+
+    fn prepare_extracted_statement(
+        &self,
+        extracted: ffi::duckdb_extracted_statements,
+        index: ffi::idx_t,
+    ) -> Result<ffi::duckdb_prepared_statement> {
+        let mut stmt = ptr::null_mut();
+        let res = unsafe { ffi::duckdb_prepare_extracted_statement(self.con, extracted, index, &mut stmt) };
+        result_from_duckdb_prepare(res, stmt)?;
+        Ok(stmt)
+    }
+
+    fn execute_extracted_statement(
+        &self,
+        extracted: ffi::duckdb_extracted_statements,
+        index: ffi::idx_t,
+    ) -> Result<()> {
+        let mut stmt = self.prepare_extracted_statement(extracted, index)?;
+
+        let mut result = unsafe { mem::zeroed() };
+        let rc = unsafe { ffi::duckdb_execute_prepared(stmt, &mut result) };
+
+        let error = if rc != ffi::DuckDBSuccess {
+            unsafe {
+                let c_err = ffi::duckdb_result_error(&mut result as *mut _);
+                let msg = if c_err.is_null() {
+                    None
+                } else {
+                    Some(CStr::from_ptr(c_err).to_string_lossy().to_string())
+                };
+                Some(Error::DuckDBFailure(ffi::Error::new(rc), msg))
+            }
+        } else {
+            None
+        };
+
+        unsafe {
+            ffi::duckdb_destroy_prepare(&mut stmt);
+            ffi::duckdb_destroy_result(&mut result);
+        }
+
+        error.map_or(Ok(()), Err)
     }
 
     pub fn appender<'a>(&mut self, conn: &'a Connection, table: &str, schema: &str) -> Result<Appender<'a>> {
@@ -126,6 +186,14 @@ impl InnerConnection {
     }
 }
 
+struct ExtractedStatementsGuard(ffi::duckdb_extracted_statements);
+
+impl Drop for ExtractedStatementsGuard {
+    fn drop(&mut self) {
+        unsafe { ffi::duckdb_destroy_extracted(&mut self.0) }
+    }
+}
+
 impl Drop for InnerConnection {
     #[allow(unused_must_use)]
     #[inline]

crates/duckdb/src/lib.rs

@@ -1455,4 +1455,56 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_prepare_multi_statement() -> Result<()> {
+        let db = checked_memory_handle();
+
+        {
+            let mut stmt =
+                db.prepare("CREATE TABLE test(x INTEGER); INSERT INTO test VALUES (42); SELECT x FROM test;")?;
+            let result: i32 = stmt.query_row([], |row| row.get(0))?;
+            assert_eq!(result, 42);
+        }
+
+        {
+            let mut stmt = db.prepare(
+                "CREATE TEMP TABLE temp_data(id INTEGER, value TEXT);
+                INSERT INTO temp_data VALUES (1, 'first'), (2, 'second');
+                SELECT COUNT(*) FROM temp_data;",
+            )?;
+            let count: i32 = stmt.query_row([], |row| row.get(0))?;
+            assert_eq!(count, 2);
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_pivot_query() -> Result<()> {
+        let db = checked_memory_handle();
+
+        db.execute_batch(
+            "CREATE TABLE cities(city VARCHAR, year INTEGER, population INTEGER);
+             INSERT INTO cities VALUES
+               ('Amsterdam', 2000, 1005),
+               ('Amsterdam', 2010, 1065),
+               ('Amsterdam', 2020, 1158),
+               ('Berlin', 2000, 3382),
+               ('Berlin', 2010, 3460),
+               ('Berlin', 2020, 3576);",
+        )?;
+
+        // PIVOT queries internally expand to multiple statements
+        let mut stmt = db.prepare("PIVOT cities ON year USING sum(population);")?;
+        let mut rows = stmt.query([])?;
+
+        let mut row_count = 0;
+        while let Some(_row) = rows.next()? {
+            row_count += 1;
+        }
+        assert_eq!(row_count, 2);
+
+        Ok(())
+    }
 }