From 47e08d34bf8e2c0f57b72ef789a675ce352344ea Mon Sep 17 00:00:00 2001
From: Pekka Enberg <penberg@iki.fi>
Date: Mon, 16 Jun 2025 14:23:26 +0300
Subject: [PATCH] bindings/rust: Fix Rows::next() I/O dispatcher handling

The `next()` function needs to be a loop to make sure we actually return rows.
---
 bindings/rust/src/lib.rs                 | 34 +++++++++++++++---------
 bindings/rust/tests/integration_tests.rs | 27 +++++++++++++++++++
 2 files changed, 49 insertions(+), 12 deletions(-)
 create mode 100644 bindings/rust/tests/integration_tests.rs

diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs
index 9f29fb420..d7925a0e3 100644
--- a/bindings/rust/src/lib.rs
+++ b/bindings/rust/src/lib.rs
@@ -289,19 +289,29 @@ unsafe impl Sync for Rows {}
 
 impl Rows {
     pub async fn next(&mut self) -> Result<Option<Row>> {
-        let mut stmt = self
-            .inner
-            .lock()
-            .map_err(|e| Error::MutexError(e.to_string()))?;
-
-        match stmt.step() {
-            Ok(limbo_core::StepResult::Row) => {
-                let row = stmt.row().unwrap();
-                Ok(Some(Row {
-                    values: row.get_values().map(|v| v.to_owned()).collect(),
-                }))
+        loop {
+            let mut stmt = self
+                .inner
+                .lock()
+                .map_err(|e| Error::MutexError(e.to_string()))?;
+            match stmt.step() {
+                Ok(limbo_core::StepResult::Row) => {
+                    let row = stmt.row().unwrap();
+                    return Ok(Some(Row {
+                        values: row.get_values().map(|v| v.to_owned()).collect(),
+                    }));
+                }
+                Ok(limbo_core::StepResult::Done) => return Ok(None),
+                Ok(limbo_core::StepResult::IO) => {
+                    if let Err(e) = stmt.run_once() {
+                        return Err(e.into());
+                    }
+                    continue;
+                }
+                Ok(limbo_core::StepResult::Busy) => return Ok(None),
+                Ok(limbo_core::StepResult::Interrupt) => return Ok(None),
+                _ => return Ok(None),
             }
-            _ => Ok(None),
         }
     }
 }
diff --git a/bindings/rust/tests/integration_tests.rs b/bindings/rust/tests/integration_tests.rs
new file mode 100644
index 000000000..b9f31cc88
--- /dev/null
+++ b/bindings/rust/tests/integration_tests.rs
@@ -0,0 +1,27 @@
+use limbo::Builder;
+
+#[tokio::test]
+async fn test_rows_next() {
+    let builder = Builder::new_local(":memory:");
+    let db = builder.build().await.unwrap();
+    let conn = db.connect().unwrap();
+    conn.execute("CREATE TABLE test (x INTEGER)", ())
+        .await
+        .unwrap();
+    conn.execute("INSERT INTO test (x) VALUES (1)", ())
+        .await
+        .unwrap();
+    conn.execute("INSERT INTO test (x) VALUES (2)", ())
+        .await
+        .unwrap();
+    let mut res = conn.query("SELECT * FROM test", ()).await.unwrap();
+    assert_eq!(
+        res.next().await.unwrap().unwrap().get_value(0).unwrap(),
+        1.into()
+    );
+    assert_eq!(
+        res.next().await.unwrap().unwrap().get_value(0).unwrap(),
+        2.into()
+    );
+    assert!(res.next().await.unwrap().is_none());
+}