fix: modify step and add reset for statement (#511)

Signed-off-by: bkioshn <[email protected]>

Author: bkioshn

hermes/bin/src/runtime_extensions/hermes/sqlite/statement/core.rs

@@ -5,11 +5,12 @@ use libsqlite3_sys::{
     sqlite3_bind_blob, sqlite3_bind_double, sqlite3_bind_int, sqlite3_bind_int64,
     sqlite3_bind_null, sqlite3_bind_text, sqlite3_column_blob, sqlite3_column_bytes,
     sqlite3_column_double, sqlite3_column_int64, sqlite3_column_text, sqlite3_column_type,
-    sqlite3_finalize, sqlite3_step, sqlite3_stmt, SQLITE_BLOB, SQLITE_DONE, SQLITE_FLOAT,
-    SQLITE_INTEGER, SQLITE_NULL, SQLITE_OK, SQLITE_ROW, SQLITE_TEXT, SQLITE_TRANSIENT,
+    sqlite3_finalize, sqlite3_reset, sqlite3_step, sqlite3_stmt, SQLITE_BLOB, SQLITE_DONE,
+    SQLITE_FLOAT, SQLITE_INTEGER, SQLITE_NULL, SQLITE_OK, SQLITE_ROW, SQLITE_TEXT,
+    SQLITE_TRANSIENT,
 };
 
-use crate::runtime_extensions::bindings::hermes::sqlite::api::{Errno, Value};
+use crate::runtime_extensions::bindings::hermes::sqlite::api::{Errno, StepResult, Value};
 
 /// Stores application data into parameters of the original SQL.
 pub(crate) fn bind(
@@ -60,13 +61,13 @@ pub(crate) fn bind(
 }
 
 /// Advances a statement to the next result row or to completion.
-pub(crate) fn step(stmt_ptr: *mut sqlite3_stmt) -> Result<(), Errno> {
+pub(crate) fn step(stmt_ptr: *mut sqlite3_stmt) -> Result<StepResult, Errno> {
     let rc = unsafe { sqlite3_step(stmt_ptr) };
 
-    if rc != SQLITE_DONE && rc != SQLITE_ROW {
-        Err(Errno::Sqlite(rc))
-    } else {
-        Ok(())
+    match rc {
+        SQLITE_DONE => Ok(StepResult::Done),
+        SQLITE_ROW => Ok(StepResult::Row),
+        _ => Err(Errno::Sqlite(rc)),
     }
 }
 
@@ -114,6 +115,17 @@ pub(crate) fn column(
     Ok(value)
 }
 
+/// Reset the prepared statement.
+pub(crate) fn reset(stmt_ptr: *mut sqlite3_stmt) -> Result<(), Errno> {
+    let rc = unsafe { sqlite3_reset(stmt_ptr) };
+
+    if rc == SQLITE_OK {
+        Ok(())
+    } else {
+        Err(Errno::Sqlite(rc))
+    }
+}
+
 /// Destroys a prepared statement object. If the most recent evaluation of the
 /// statement encountered no errors or if the statement is never been evaluated,
 /// then the function results without errors. If the most recent evaluation of
@@ -144,8 +156,7 @@ mod tests {
     const TMP_DIR: &str = "tmp-dir";
 
     fn init() -> Result<*mut sqlite3, Errno> {
-        let app_name = ApplicationName(String::from(TMP_DIR));
-
+        let app_name = ApplicationName(TMP_DIR.to_string());
         open(false, true, app_name)
     }
 
@@ -155,138 +166,124 @@ mod tests {
         value: Value,
     ) -> Result<(), Errno> {
         let sql = format!("CREATE TABLE Dummy(Id INTEGER PRIMARY KEY, Value {db_value_type});");
-
-        execute(db_ptr, sql.as_str())?;
+        execute(db_ptr, &sql)?;
 
         let sql = "INSERT INTO Dummy(Value) VALUES(?);";
-
         let stmt_ptr = prepare(db_ptr, sql)?;
-
         bind(stmt_ptr, 1, value)?;
         step(stmt_ptr)?;
         finalize(stmt_ptr)?;
-
         Ok(())
     }
 
     fn get_value(db_ptr: *mut sqlite3) -> Result<Value, Errno> {
         let sql = "SELECT Value FROM Dummy WHERE Id = 1;";
 
         let stmt_ptr = prepare(db_ptr, sql)?;
-        step(stmt_ptr)?;
-        let col_result = column(stmt_ptr, 0);
-        finalize(stmt_ptr)?;
-
-        col_result
+        match step(stmt_ptr)? {
+            StepResult::Row => {
+                let val = column(stmt_ptr, 0)?;
+                finalize(stmt_ptr)?;
+                Ok(val)
+            },
+            StepResult::Done => {
+                finalize(stmt_ptr)?;
+                Err(Errno::Sqlite(SQLITE_DONE))
+            },
+        }
     }
 
-    #[test]
-    fn test_value_double() -> Result<(), Errno> {
+    fn test_single_value(
+        db_type: &str,
+        value: &Value,
+    ) -> Result<(), Errno> {
         let db_ptr = init()?;
-
-        let value = Value::Double(std::f64::consts::PI);
-        init_value(db_ptr, "REAL", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(
-            matches!((value, value_result), (Value::Double(x), Ok(Value::Double(y))) if x.eq(&y))
-        );
-
+        init_value(db_ptr, db_type, value.clone())?;
+        let value_result = get_value(db_ptr)?;
+        assert_eq!(format!("{value:?}",), format!("{value_result:?}",));
         close(db_ptr)
     }
 
     #[test]
-    fn test_value_bool() -> Result<(), Errno> {
-        let db_ptr = init()?;
-
-        let value = Value::Int32(1);
-        init_value(db_ptr, "BOOLEAN", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(matches!((value, value_result), (Value::Int32(x), Ok(Value::Int32(y))) if x == y));
-
-        close(db_ptr)
+    fn test_double() -> Result<(), Errno> {
+        test_single_value("REAL", &Value::Double(std::f64::consts::PI))
     }
-
     #[test]
-    fn test_value_int32() -> Result<(), Errno> {
-        let db_ptr = init()?;
-
-        let value = Value::Int32(i32::MAX);
-        init_value(db_ptr, "MEDIUMINT", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(matches!((value, value_result), (Value::Int32(x), Ok(Value::Int32(y))) if x == y));
-
-        close(db_ptr)
+    fn test_bool() -> Result<(), Errno> {
+        test_single_value("BOOLEAN", &Value::Int32(1))
     }
-
     #[test]
-    fn test_value_int32_nullable() -> Result<(), Errno> {
-        let db_ptr = init()?;
-
-        let value = Value::Null;
-        init_value(db_ptr, "MEDIUMINT", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(matches!(
-            (value, value_result),
-            (Value::Null, Ok(Value::Null))
-        ));
-
-        close(db_ptr)
+    fn test_int32() -> Result<(), Errno> {
+        test_single_value("MEDIUMINT", &Value::Int32(i32::MAX))
     }
-
     #[test]
-    fn test_value_int64() -> Result<(), Errno> {
-        let db_ptr = init()?;
-
-        let value = Value::Int64(i64::MAX);
-        init_value(db_ptr, "BIGINT", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(matches!((value, value_result), (Value::Int64(x), Ok(Value::Int64(y))) if x == y));
-
-        close(db_ptr)
+    fn test_int32_nullable() -> Result<(), Errno> {
+        test_single_value("MEDIUMINT", &Value::Null)
     }
-
     #[test]
-    fn test_value_text() -> Result<(), Errno> {
-        let db_ptr = init()?;
-
-        let value = Value::Text(String::from("Hello, World!"));
-        init_value(db_ptr, "TEXT", value.clone())?;
-        let value_result = get_value(db_ptr);
-
-        assert!(matches!((value, value_result), (Value::Text(x), Ok(Value::Text(y))) if x == y));
-
-        close(db_ptr)
+    fn test_int64() -> Result<(), Errno> {
+        test_single_value("BIGINT", &Value::Int64(i64::MAX))
+    }
+    #[test]
+    fn test_text() -> Result<(), Errno> {
+        test_single_value("TEXT", &Value::Text("Hello, World!".to_string()))
+    }
+    #[test]
+    fn test_blob() -> Result<(), Errno> {
+        test_single_value("BLOB", &Value::Blob(vec![1, 2, 3, 4, 5]))
     }
 
     #[test]
-    fn test_value_blob() -> Result<(), Errno> {
+    fn test_finalize_simple() -> Result<(), Errno> {
         let db_ptr = init()?;
 
-        let value = Value::Blob(vec![1, 2, 3, 4, 5]);
-        init_value(db_ptr, "BLOB", value.clone())?;
-        let value_result = get_value(db_ptr);
+        let sql = "SELECT 1;";
+
+        let stmt_ptr = core::prepare(db_ptr, sql)?;
+
+        let result = finalize(stmt_ptr);
 
-        assert!(matches!((value, value_result), (Value::Blob(x), Ok(Value::Blob(y))) if x == y));
+        assert!(result.is_ok());
 
         close(db_ptr)
     }
 
     #[test]
-    fn test_finalize_simple() -> Result<(), Errno> {
+    fn test_loop_over_rows() -> Result<(), Errno> {
         let db_ptr = init()?;
 
-        let sql = "SELECT 1;";
+        let sql = "CREATE TABLE IF NOT EXISTS Dummy(Id INTEGER PRIMARY KEY, Value INTEGER);";
+        execute(db_ptr, sql)?;
 
-        let stmt_ptr = core::prepare(db_ptr, sql)?;
+        // Insert multiple rows
+        let sql = "INSERT INTO Dummy(Value) VALUES(?);";
+        let stmt_ptr = prepare(db_ptr, sql)?;
+        for i in 1..=5 {
+            bind(stmt_ptr, 1, Value::Int32(i))?;
+            step(stmt_ptr)?;
+            reset(stmt_ptr)?;
+        }
+        finalize(stmt_ptr)?;
 
-        let result = finalize(stmt_ptr);
+        let sql = "SELECT Value FROM Dummy ORDER BY Id ASC;";
+        let stmt_ptr = prepare(db_ptr, sql)?;
+        let mut collected = Vec::new();
+
+        loop {
+            let result = step(stmt_ptr)?;
+            match result {
+                StepResult::Row => collected.push(column(stmt_ptr, 0)?),
+                StepResult::Done => break,
+            }
+        }
+        finalize(stmt_ptr)?;
 
-        assert!(result.is_ok());
+        for (i, val) in collected.iter().enumerate() {
+            match val {
+                Value::Int32(v) => assert_eq!(*v, i32::try_from(i).unwrap() + 1),
+                _ => panic!("Unexpected value type: {val:?}"),
+            }
+        }
 
         close(db_ptr)
     }

hermes/bin/src/runtime_extensions/hermes/sqlite/statement/host.rs

@@ -3,7 +3,9 @@
 use super::{super::state::get_statement_state, core};
 use crate::{
     runtime_context::HermesRuntimeContext,
-    runtime_extensions::bindings::hermes::sqlite::api::{Errno, HostStatement, Statement, Value},
+    runtime_extensions::bindings::hermes::sqlite::api::{
+        Errno, HostStatement, Statement, StepResult, Value,
+    },
 };
 
 impl HostStatement for HermesRuntimeContext {
@@ -29,10 +31,14 @@ impl HostStatement for HermesRuntimeContext {
     ///
     /// After a prepared statement has been prepared, this function must be called one or
     /// more times to evaluate the statement.
+    ///
+    /// ## Returns
+    ///
+    /// A `step-result` indicating the status of the step.
     fn step(
         &mut self,
         resource: wasmtime::component::Resource<Statement>,
-    ) -> wasmtime::Result<Result<(), Errno>> {
+    ) -> wasmtime::Result<Result<StepResult, Errno>> {
         let mut app_state = get_statement_state().get_app_state(self.app_name())?;
         let stmt_ptr = app_state.get_object(&resource)?;
         Ok(core::step(*stmt_ptr as *mut _))
@@ -62,6 +68,21 @@ impl HostStatement for HermesRuntimeContext {
         Ok(core::column(*stmt_ptr as *mut _, index))
     }
 
+    /// Reset a prepared statement object back to its initial state, ready to be
+    /// re-executed.
+    ///
+    /// This function clears all previous bindings, resets the statement to the beginning,
+    /// and prepares it for another execution. This must be called before reusing a
+    /// statement with new parameter bindings.
+    fn reset(
+        &mut self,
+        resource: wasmtime::component::Resource<Statement>,
+    ) -> wasmtime::Result<Result<(), Errno>> {
+        let mut app_state = get_statement_state().get_app_state(self.app_name())?;
+        let stmt_ptr = app_state.get_object(&resource)?;
+        Ok(core::reset(*stmt_ptr as *mut _))
+    }
+
     /// Destroys a prepared statement object. If the most recent evaluation of the
     /// statement encountered no errors or if the statement is never been evaluated,
     /// then the function results without errors. If the most recent evaluation of

wasm/wasi/wit/deps/hermes-sqlite/api.wit

@@ -61,6 +61,14 @@ interface api {
         text(string)
     }
 
+    /// The result of advancing a prepared SQLite statement by one step.
+    enum step-result {
+        /// Indicates that the statement has finished executing.
+        done,
+        /// Indicates that there is a new row of result.
+        row,
+    }
+
     /// The database connection object.
     resource sqlite {
         /// Closes a database connection, destructor for `sqlite3`.
@@ -115,7 +123,11 @@ interface api {
         /// Advances a statement to the next result row or to completion.
         ///
         /// After a prepared statement has been prepared, this function must be called one or more times to evaluate the statement.
-        step: func() -> result<_, errno>;
+        /// 
+        /// ## Returns
+        /// 
+        /// A `step-result` indicating the status of the step.
+        step: func() -> result<step-result, errno>;
 
         /// Returns information about a single column of the current result row of a query.
         ///
@@ -130,6 +142,13 @@ interface api {
         /// The value of a result column in a specific data format.
         column: func(index: u32) -> result<value, errno>;
 
+        /// Reset a prepared statement object back to its initial state, ready to be re-executed.
+        ///
+        /// This function clears all previous bindings, resets the statement to the beginning,
+        /// and prepares it for another execution. This must be called before reusing a statement
+        /// with new parameter bindings.
+        reset: func() -> result<_, errno>;
+
         /// Destroys a prepared statement object. If the most recent evaluation of the statement encountered no errors or if the statement is never been evaluated,
         /// then the function results without errors. If the most recent evaluation of statement failed, then the function results the appropriate error code.
         ///
@@ -155,4 +174,4 @@ interface api {
 /// World just for the Hermes 'sqlite' API.
 world sqlite-api {
     import api;
-}
+}