Skip to content

Commit 85de932

Browse files
authored
RUST-1757 Fix final cursor batch handling (#951)
1 parent 1e0dd95 commit 85de932

File tree

2 files changed

+39
-6
lines changed

2 files changed

+39
-6
lines changed

src/cursor/common.rs

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -135,12 +135,14 @@ where
135135
let result = self.provider.execute(spec, client, pin).await;
136136
self.handle_get_more_result(result)?;
137137

138-
if self.is_exhausted() {
139-
Ok(AdvanceResult::Exhausted)
140-
} else {
141-
match self.state_mut().buffer.advance() {
142-
true => Ok(AdvanceResult::Advanced),
143-
false => Ok(AdvanceResult::Waiting),
138+
match self.state_mut().buffer.advance() {
139+
true => Ok(AdvanceResult::Advanced),
140+
false => {
141+
if self.is_exhausted() {
142+
Ok(AdvanceResult::Exhausted)
143+
} else {
144+
Ok(AdvanceResult::Waiting)
145+
}
144146
}
145147
}
146148
}

src/test/cursor.rs

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,3 +257,34 @@ async fn session_cursor_with_type() {
257257

258258
let _ = cursor_with_type.next(&mut session).await.unwrap().unwrap();
259259
}
260+
261+
#[cfg_attr(feature = "tokio-runtime", tokio::test)]
262+
#[cfg_attr(feature = "async-std-runtime", async_std::test)]
263+
async fn cursor_final_batch() {
264+
let client = TestClient::new().await;
265+
let coll = client
266+
.create_fresh_collection("test_cursor_final_batch", "test", None)
267+
.await;
268+
coll.insert_many(
269+
vec![
270+
doc! { "foo": 1 },
271+
doc! { "foo": 2 },
272+
doc! { "foo": 3 },
273+
doc! { "foo": 4 },
274+
doc! { "foo": 5 },
275+
],
276+
None,
277+
)
278+
.await
279+
.unwrap();
280+
281+
let mut cursor = coll
282+
.find(None, FindOptions::builder().batch_size(3).build())
283+
.await
284+
.unwrap();
285+
let mut found = 0;
286+
while cursor.advance().await.unwrap() {
287+
found += 1;
288+
}
289+
assert_eq!(found, 5);
290+
}

0 commit comments

Comments
 (0)