Skip to content

Commit b8446b2

Browse files
authored
RUST-1757 Fix final cursor batch handling (#951) (#954)
1 parent 640a30a commit b8446b2

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
@@ -265,3 +265,34 @@ async fn session_cursor_with_type() {
265265

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

0 commit comments

Comments
 (0)