libsql: add batched rows for local conn (#1499)

* libsql: add batched rows for local conn

* continue impl

* Fix local batched rows test

* remove dbg

* fix feature flag

Author: LucioFranco

libsql/src/connection.rs

@@ -61,7 +61,7 @@ impl BatchRows {
         }
     }
 
-    #[cfg(feature = "hrana")]
+    #[cfg(any(feature = "hrana", feature = "core"))]
     pub(crate) fn new(rows: Vec<Option<Rows>>) -> Self {
         Self {
             inner: rows.into(),

libsql/src/local/connection.rs

@@ -1,7 +1,8 @@
 #![allow(dead_code)]
 
+use crate::local::rows::BatchedRows;
 use crate::params::Params;
-use crate::errors;
+use crate::{connection::BatchRows, errors};
 
 use super::{Database, Error, Result, Rows, RowsFuture, Statement, Transaction};
 
@@ -136,21 +137,93 @@ impl Connection {
     ///
     /// Will return `Err` if `sql` cannot be converted to a C-compatible string
     /// or if the underlying SQLite call fails.
-    pub fn execute_batch<S>(&self, sql: S) -> Result<()>
+    pub fn execute_batch<S>(&self, sql: S) -> Result<BatchRows>
     where
         S: Into<String>,
     {
         let sql = sql.into();
         let mut sql = sql.as_str();
 
+        let mut batch_rows = Vec::new();
+
         while !sql.is_empty() {
             let stmt = self.prepare(sql)?;
 
-            if !stmt.inner.raw_stmt.is_null() {
-                stmt.step()?;
-            }
+            let tail = if !stmt.inner.raw_stmt.is_null() {
+                let returned_rows = stmt.step()?;
 
-            let tail = stmt.tail();
+                let tail = stmt.tail();
+
+                // Check if there are rows to be extracted, we must do this upfront due to the lazy
+                // nature of sqlite and our somewhat hacked batch command.
+                if returned_rows {
+                    // Extract columns
+                    let cols = stmt
+                        .columns()
+                        .iter()
+                        .enumerate()
+                        .map(|(i, c)| {
+                            use crate::value::ValueType;
+
+                            let val = stmt.inner.column_type(i as i32);
+                            let t = match val {
+                                libsql_sys::ffi::SQLITE_INTEGER => ValueType::Integer,
+                                libsql_sys::ffi::SQLITE_FLOAT => ValueType::Real,
+                                libsql_sys::ffi::SQLITE_BLOB => ValueType::Blob,
+                                libsql_sys::ffi::SQLITE_TEXT => ValueType::Text,
+                                libsql_sys::ffi::SQLITE_NULL => ValueType::Null,
+                                _ => unreachable!("unknown column type {} at index {}", val, i),
+                            };
+
+                            (c.name.to_string(), t)
+                        })
+                        .collect::<Vec<_>>();
+
+                    let mut rows = Vec::new();
+
+                    // If returned rows we must extract the rows available right away instead of
+                    // using the `Rows` type we have already. This is due to the step api once its
+                    // returned SQLITE_ROWS we must extract them before we call step again.
+                    {
+                        let row = crate::local::Row { stmt: stmt.clone() };
+
+                        let mut values = Vec::with_capacity(cols.len());
+
+                        for i in 0..cols.len() {
+                            let value = row.get_value(i as i32)?;
+
+                            values.push(value);
+                        }
+
+                        rows.push(values);
+                    }
+
+                    // Now we can use the normal rows type to extract any n+1 rows
+                    let rows_sys = Rows::new(stmt);
+
+                    while let Some(row) = rows_sys.next()? {
+                        let mut values = Vec::with_capacity(cols.len());
+
+                        for i in 0..cols.len() {
+                            let value = row.get_value(i as i32)?;
+
+                            values.push(value);
+                        }
+
+                        rows.push(values);
+                    }
+
+                    rows.len();
+
+                    batch_rows.push(Some(crate::Rows::new(BatchedRows::new(cols, rows))));
+                } else {
+                    batch_rows.push(None);
+                }
+
+                tail
+            } else {
+                stmt.tail()
+            };
 
             if tail == 0 || tail >= sql.len() {
                 break;
@@ -159,12 +232,12 @@ impl Connection {
             sql = &sql[tail..];
         }
 
-        Ok(())
+        Ok(BatchRows::new(batch_rows))
     }
 
     fn execute_transactional_batch_inner<S>(&self, sql: S) -> Result<()>
-        where
-            S: Into<String>,
+    where
+        S: Into<String>,
     {
         let sql = sql.into();
         let mut sql = sql.as_str();
@@ -177,13 +250,16 @@ impl Connection {
             } else {
                 &sql[..tail]
             };
-            let prefix_count = stmt_sql
-                .chars()
-                .take_while(|c| c.is_whitespace())
-                .count();
+            let prefix_count = stmt_sql.chars().take_while(|c| c.is_whitespace()).count();
             let stmt_sql = &stmt_sql[prefix_count..];
-            if stmt_sql.starts_with("BEGIN") || stmt_sql.starts_with("COMMIT") || stmt_sql.starts_with("ROLLBACK") || stmt_sql.starts_with("END") {
-                return Err(Error::TransactionalBatchError("Transactions forbidden inside transactional batch".to_string()));
+            if stmt_sql.starts_with("BEGIN")
+                || stmt_sql.starts_with("COMMIT")
+                || stmt_sql.starts_with("ROLLBACK")
+                || stmt_sql.starts_with("END")
+            {
+                return Err(Error::TransactionalBatchError(
+                    "Transactions forbidden inside transactional batch".to_string(),
+                ));
             }
 
             if !stmt.inner.raw_stmt.is_null() {
@@ -299,35 +375,53 @@ impl Connection {
     pub fn enable_load_extension(&self, onoff: bool) -> Result<()> {
         // SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION configration verb accepts 2 additional parameters: an on/off flag and a pointer to an c_int where new state of the parameter will be written (or NULL if reporting back the setting is not needed)
         // See: https://sqlite.org/c3ref/c_dbconfig_defensive.html#sqlitedbconfigenableloadextension
-        let err = unsafe { ffi::sqlite3_db_config(self.raw, ffi::SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION, onoff as i32, std::ptr::null::<c_int>()) };
+        let err = unsafe {
+            ffi::sqlite3_db_config(
+                self.raw,
+                ffi::SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION,
+                onoff as i32,
+                std::ptr::null::<c_int>(),
+            )
+        };
         match err {
             ffi::SQLITE_OK => Ok(()),
-            _ => Err(errors::Error::SqliteFailure(err, errors::error_from_code(err))),
+            _ => Err(errors::Error::SqliteFailure(
+                err,
+                errors::error_from_code(err),
+            )),
         }
     }
 
-    pub fn load_extension(
-        &self,
-        dylib_path: &Path,
-        entry_point: Option<&str>,
-    ) -> Result<()> {
+    pub fn load_extension(&self, dylib_path: &Path, entry_point: Option<&str>) -> Result<()> {
         let mut raw_err_msg: *mut std::ffi::c_char = std::ptr::null_mut();
         let dylib_path = match dylib_path.to_str() {
-            Some(dylib_path) => {
-                std::ffi::CString::new(dylib_path).unwrap()
-            },
-            None => return Err(crate::Error::Misuse(format!(
-                "dylib path is not a valid utf8 string"
-            ))),
+            Some(dylib_path) => std::ffi::CString::new(dylib_path).unwrap(),
+            None => {
+                return Err(crate::Error::Misuse(format!(
+                    "dylib path is not a valid utf8 string"
+                )))
+            }
         };
         let err = match entry_point {
             Some(entry_point) => {
                 let entry_point = std::ffi::CString::new(entry_point).unwrap();
-                unsafe { ffi::sqlite3_load_extension(self.raw, dylib_path.as_ptr(), entry_point.as_ptr(), &mut raw_err_msg) }
-            }
-            None => {
-                unsafe { ffi::sqlite3_load_extension(self.raw, dylib_path.as_ptr(), std::ptr::null(), &mut raw_err_msg) }
+                unsafe {
+                    ffi::sqlite3_load_extension(
+                        self.raw,
+                        dylib_path.as_ptr(),
+                        entry_point.as_ptr(),
+                        &mut raw_err_msg,
+                    )
+                }
             }
+            None => unsafe {
+                ffi::sqlite3_load_extension(
+                    self.raw,
+                    dylib_path.as_ptr(),
+                    std::ptr::null(),
+                    &mut raw_err_msg,
+                )
+            },
         };
         match err {
             ffi::SQLITE_OK => Ok(()),

libsql/src/local/impls.rs

@@ -24,8 +24,7 @@ impl Conn for LibsqlConnection {
     }
 
     async fn execute_batch(&self, sql: &str) -> Result<BatchRows> {
-        self.conn.execute_batch(sql)?;
-        Ok(BatchRows::empty())
+        self.conn.execute_batch(sql)
     }
 
     async fn execute_transactional_batch(&self, sql: &str) -> Result<BatchRows> {

libsql/src/local/rows.rs

@@ -1,12 +1,15 @@
 use crate::local::{Connection, Statement};
 use crate::params::Params;
+use crate::rows::{RowInner, RowsInner};
 use crate::{errors, Error, Result};
 use crate::{Value, ValueRef};
 use libsql_sys::ValueType;
 
-use std::fmt;
 use std::cell::RefCell;
+use std::collections::VecDeque;
 use std::ffi::c_char;
+use std::fmt;
+use std::sync::Arc;
 /// Query result rows.
 #[derive(Debug, Clone)]
 pub struct Rows {
@@ -165,7 +168,9 @@ impl fmt::Debug for Row {
                         ValueRef::Null => dbg_map.value(&(value_type, ())),
                         ValueRef::Integer(i) => dbg_map.value(&(value_type, i)),
                         ValueRef::Real(f) => dbg_map.value(&(value_type, f)),
-                        ValueRef::Text(s) => dbg_map.value(&(value_type, String::from_utf8_lossy(s))),
+                        ValueRef::Text(s) => {
+                            dbg_map.value(&(value_type, String::from_utf8_lossy(s)))
+                        }
                         ValueRef::Blob(b) => dbg_map.value(&(value_type, b.len())),
                     };
                 }
@@ -178,6 +183,94 @@ impl fmt::Debug for Row {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct BatchedRows {
+    /// Colname, decl_type
+    cols: Arc<Vec<(String, crate::value::ValueType)>>,
+    rows: VecDeque<Vec<Value>>,
+}
+
+impl BatchedRows {
+    pub fn new(cols: Vec<(String, crate::value::ValueType)>, rows: Vec<Vec<Value>>) -> Self {
+        Self {
+            cols: Arc::new(cols),
+            rows: rows.into(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl RowsInner for BatchedRows {
+    async fn next(&mut self) -> Result<Option<crate::Row>> {
+        let cols = self.cols.clone();
+        let row = self.rows.pop_front();
+
+        if let Some(row) = row {
+            Ok(Some(crate::Row {
+                inner: Box::new(BatchedRow { cols, row }),
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn column_count(&self) -> i32 {
+        self.cols.len() as i32
+    }
+
+    fn column_name(&self, idx: i32) -> Option<&str> {
+        self.cols.get(idx as usize).map(|s| s.0.as_str())
+    }
+
+    fn column_type(&self, idx: i32) -> Result<crate::value::ValueType> {
+        self.cols
+            .get(idx as usize)
+            .ok_or(Error::InvalidColumnIndex)
+            .map(|(_, vt)| vt.clone())
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct BatchedRow {
+    cols: Arc<Vec<(String, crate::value::ValueType)>>,
+    row: Vec<Value>,
+}
+
+impl RowInner for BatchedRow {
+    fn column_value(&self, idx: i32) -> Result<Value> {
+        self.row
+            .get(idx as usize)
+            .cloned()
+            .ok_or(Error::InvalidColumnIndex)
+    }
+
+    fn column_name(&self, idx: i32) -> Option<&str> {
+        self.cols.get(idx as usize).map(|c| c.0.as_str())
+    }
+
+    fn column_str(&self, idx: i32) -> Result<&str> {
+        self.row
+            .get(idx as usize)
+            .ok_or(Error::InvalidColumnIndex)
+            .and_then(|v| {
+                v.as_text()
+                    .map(String::as_str)
+                    .ok_or(Error::InvalidColumnType)
+            })
+    }
+
+    fn column_count(&self) -> usize {
+        self.cols.len()
+    }
+
+    fn column_type(&self, idx: i32) -> Result<crate::value::ValueType> {
+        self.cols
+            .get(idx as usize)
+            .ok_or(Error::InvalidColumnIndex)
+            .map(|(_, vt)| vt.clone())
+    }
+}
+
 pub trait FromValue {
     fn from_sql(val: libsql_sys::Value) -> Result<Self>
     where

libsql/src/local/statement.rs

@@ -221,11 +221,12 @@ impl Statement {
         }
     }
 
-    pub(crate) fn step(&self) -> Result<()> {
+    /// Returns true if this statement has rows ready to be read.
+    pub(crate) fn step(&self) -> Result<bool> {
         let err = self.inner.step();
         match err {
-            crate::ffi::SQLITE_DONE => Ok(()),
-            crate::ffi::SQLITE_ROW => Err(Error::ExecuteReturnedRows),
+            crate::ffi::SQLITE_DONE => Ok(false),
+            crate::ffi::SQLITE_ROW => Ok(true),
             _ => Err(Error::SqliteFailure(
                 errors::extended_error_code(self.conn.raw),
                 errors::error_from_handle(self.conn.raw),