introduce program execution state in order to run stmt to completion in case of finalize or reset

Author: sivukhin

core/lib.rs

@@ -31,6 +31,7 @@ mod util;
 #[cfg(feature = "uuid")]
 mod uuid;
 mod vdbe;
+pub type ProgramExecutionState = vdbe::ProgramExecutionState;
 pub mod vector;
 mod vtab;
 
@@ -2650,6 +2651,10 @@ impl Statement {
         self.state.interrupt();
     }
 
+    pub fn execution_state(&self) -> ProgramExecutionState {
+        self.state.execution_state
+    }
+
     fn _step(&mut self, waker: Option<&Waker>) -> Result<StepResult> {
         if let Some(busy_timeout) = self.busy_timeout.as_ref() {
             if self.pager.io.now() < busy_timeout.timeout {

core/vdbe/mod.rs

@@ -273,6 +273,16 @@ pub enum TxnCleanup {
     RollbackTxn,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum ProgramExecutionState {
+    Init,
+    Run,
+    Interrupting,
+    Interrupted,
+    Done,
+    Failed,
+}
+
 /// The program state describes the environment in which the program executes.
 pub struct ProgramState {
     pub io_completions: Option<IOCompletions>,
@@ -287,7 +297,7 @@ pub struct ProgramState {
     /// Indicate whether an [Insn::Once] instruction at a given program counter position has already been executed, well, once.
     once: SmallVec<u32, 4>,
     regex_cache: RegexCache,
-    interrupted: bool,
+    pub execution_state: ProgramExecutionState,
     pub parameters: HashMap<NonZero<usize>, Value>,
     commit_state: CommitState,
     #[cfg(feature = "json")]
@@ -356,7 +366,7 @@ impl ProgramState {
             ended_coroutine: Bitfield::new(),
             once: SmallVec::<u32, 4>::new(),
             regex_cache: RegexCache::new(),
-            interrupted: false,
+            execution_state: ProgramExecutionState::Init,
             parameters: HashMap::new(),
             commit_state: CommitState::Ready,
             #[cfg(feature = "json")]
@@ -409,11 +419,7 @@ impl ProgramState {
     }
 
     pub fn interrupt(&mut self) {
-        self.interrupted = true;
-    }
-
-    pub fn is_interrupted(&self) -> bool {
-        self.interrupted
+        self.execution_state = ProgramExecutionState::Interrupting;
     }
 
     pub fn bind_at(&mut self, index: NonZero<usize>, value: Value) {
@@ -450,7 +456,7 @@ impl ProgramState {
         self.deferred_seeks.iter_mut().for_each(|s| *s = None);
         self.ended_coroutine.0 = [0; 4];
         self.regex_cache.like.clear();
-        self.interrupted = false;
+        self.execution_state = ProgramExecutionState::Init;
         self.current_collation = None;
         #[cfg(feature = "json")]
         self.json_cache.clear();
@@ -650,11 +656,25 @@ impl Program {
         query_mode: QueryMode,
         waker: Option<&Waker>,
     ) -> Result<StepResult> {
-        match query_mode {
+        state.execution_state = ProgramExecutionState::Run;
+        let result = match query_mode {
             QueryMode::Normal => self.normal_step(state, mv_store, pager, waker),
             QueryMode::Explain => self.explain_step(state, mv_store, pager),
             QueryMode::ExplainQueryPlan => self.explain_query_plan_step(state, mv_store, pager),
+        };
+        match &result {
+            Ok(StepResult::Done) => {
+                state.execution_state = ProgramExecutionState::Done;
+            }
+            Ok(StepResult::Interrupt) => {
+                state.execution_state = ProgramExecutionState::Interrupted;
+            }
+            Err(_) => {
+                state.execution_state = ProgramExecutionState::Failed;
+            }
+            _ => {}
         }
+        result
     }
 
     fn explain_step(
@@ -673,7 +693,7 @@ impl Program {
             return Err(LimboError::InternalError("Connection closed".to_string()));
         }
 
-        if state.is_interrupted() {
+        if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
             return Ok(StepResult::Interrupt);
         }
 
@@ -823,7 +843,7 @@ impl Program {
                 return Err(LimboError::InternalError("Connection closed".to_string()));
             }
 
-            if state.is_interrupted() {
+            if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
                 return Ok(StepResult::Interrupt);
             }
 
@@ -871,10 +891,11 @@ impl Program {
                 }
                 return Err(LimboError::InternalError("Connection closed".to_string()));
             }
-            if state.is_interrupted() {
+            if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
                 self.abort(mv_store, &pager, None, state);
                 return Ok(StepResult::Interrupt);
             }
+
             if let Some(io) = &state.io_completions {
                 if !io.finished() {
                     io.set_waker(waker);

sqlite3/src/lib.rs

@@ -4,7 +4,7 @@
 use std::ffi::{self, CStr, CString};
 use std::num::{NonZero, NonZeroUsize};
 use tracing::trace;
-use turso_core::{CheckpointMode, LimboError, Value};
+use turso_core::{CheckpointMode, LimboError, ProgramExecutionState, Value};
 
 use std::sync::{Arc, Mutex};
 
@@ -325,13 +325,32 @@ pub unsafe extern "C" fn sqlite3_prepare_v2(
     SQLITE_OK
 }
 
+unsafe fn stmt_run_to_completion(stmt: *mut sqlite3_stmt) -> ffi::c_int {
+    let stmt_ref = &mut *stmt;
+    while stmt_ref.stmt.execution_state() == ProgramExecutionState::Run {
+        let result = sqlite3_step(stmt);
+        if result != SQLITE_DONE && result != SQLITE_ROW {
+            return result;
+        }
+    }
+    SQLITE_OK
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn sqlite3_finalize(stmt: *mut sqlite3_stmt) -> ffi::c_int {
     if stmt.is_null() {
         return SQLITE_MISUSE;
     }
     let stmt_ref = &mut *stmt;
 
+    // first, finalize any execution if it was unfinished
+    // (for example, many drivers can consume just one row and finalize statement after that, while there still can be work to do)
+    // (this is necessary because queries like INSERT INTO t VALUES (1), (2), (3) RETURNING id return values within a transaction)
+    let result = stmt_run_to_completion(stmt);
+    if result != SQLITE_OK {
+        return result;
+    }
+
     if !stmt_ref.db.is_null() {
         let db = &mut *stmt_ref.db;
         let mut db_inner = db.inner.lock().unwrap();
@@ -669,6 +688,13 @@ fn split_sql_statements(sql: &str) -> Vec<&str> {
 #[no_mangle]
 pub unsafe extern "C" fn sqlite3_reset(stmt: *mut sqlite3_stmt) -> ffi::c_int {
     let stmt = &mut *stmt;
+    // first, finalize any execution if it was unfinished
+    // (for example, many drivers can consume just one row and finalize statement after that, while there still can be work to do)
+    // (this is necessary because queries like INSERT INTO t VALUES (1), (2), (3) RETURNING id return values within a transaction)
+    let result = stmt_run_to_completion(stmt);
+    if result != SQLITE_OK {
+        return result;
+    }
     stmt.stmt.reset();
     stmt.clear_text_cache();
     SQLITE_OK
@@ -1420,10 +1446,10 @@ unsafe extern "C" fn sqlite_get_table_cb(
     for i in 0..n_column {
         let value = *argv.add(i as usize);
         let value_cstring = if !value.is_null() {
-            let len = libc::strlen(value);
-            let mut buf = Vec::with_capacity(len + 1);
-            libc::strncpy(buf.as_mut_ptr() as *mut ffi::c_char, value, len);
-            buf.set_len(len + 1);
+            let value_cstr = CStr::from_ptr(value).to_bytes();
+            let len = value_cstr.len();
+            let mut buf = vec![0u8; len + 1];
+            buf[0..len].copy_from_slice(value_cstr);
             CString::from_vec_with_nul(buf).unwrap()
         } else {
             CString::new("NULL").unwrap()

sqlite3/tests/sqlite3_tests.c

@@ -20,6 +20,7 @@ void test_sqlite3_column_type();
 void test_sqlite3_column_decltype();
 void test_sqlite3_next_stmt();
 void test_sqlite3_table_column_metadata();
+void test_sqlite3_insert_returning();
 
 int allocated = 0;
 
@@ -39,6 +40,7 @@ int main(void)
     test_sqlite3_column_decltype();
     test_sqlite3_next_stmt();
     test_sqlite3_table_column_metadata();
+    test_sqlite3_insert_returning();
     return 0;
 }
 
@@ -746,3 +748,47 @@ void test_sqlite3_table_column_metadata()
 	printf("sqlite3_table_column_metadata test passed\n");
 	sqlite3_close(db);
 }
+
+void test_sqlite3_insert_returning()
+{
+    sqlite3 *db;
+    sqlite3_stmt *stmt;
+    char *err_msg = NULL;
+    int rc;
+
+    rc = sqlite3_open(":memory:", &db);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_exec(db,
+                      "CREATE TABLE t(x)",
+                      NULL, NULL, &err_msg);
+    assert(rc == SQLITE_OK);
+
+    if (err_msg)
+    {
+        sqlite3_free(err_msg);
+        err_msg = NULL;
+    }
+    rc = sqlite3_prepare_v2(db, "INSERT INTO t (x) VALUES (1), (2), (3) RETURNING x;", -1, &stmt, NULL);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_step(stmt);
+    assert(rc == SQLITE_ROW);
+
+    rc = sqlite3_finalize(stmt);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_prepare_v2(db, "SELECT COUNT(*) FROM t;", -1, &stmt, NULL);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_step(stmt);
+    assert(rc == SQLITE_ROW);
+
+    sqlite_int64 fetched = sqlite3_column_int64(stmt, 0);
+    assert(fetched == 3);
+
+    sqlite3_finalize(stmt);
+    sqlite3_close(db);
+
+    printf("test_sqlite3_insert_retuning test passed\n");
+}