RUST-796 Provide easy way to reborrow session in between cursor iterations (#341)

Author: patrickfreed

src/coll/mod.rs

@@ -488,7 +488,7 @@ where
         let mut cursor = self
             .find_with_session(filter, Some(options), session)
             .await?;
-        let mut cursor = cursor.with_session(session);
+        let mut cursor = cursor.stream(session);
         cursor.next().await.transpose()
     }
 

src/cursor/mod.rs

@@ -21,12 +21,13 @@ use crate::{
 pub(crate) use common::{CursorInformation, CursorSpecification};
 use common::{GenericCursor, GetMoreProvider, GetMoreProviderResult};
 
-/// A `Cursor` streams the result of a query. When a query is made, a `Cursor` will be returned with
-/// the first batch of results from the server; the documents will be returned as the `Cursor` is
-/// iterated. When the batch is exhausted and if there are more results, the `Cursor` will fetch the
-/// next batch of documents, and so forth until the results are exhausted. Note that because of this
-/// batching, additional network I/O may occur on any given call to `Cursor::next`. Because of this,
-/// a `Cursor` iterates over `Result<Document>` items rather than simply `Document` items.
+/// A [`Cursor`] streams the result of a query. When a query is made, the returned [`Cursor`] will
+/// contain the first batch of results from the server; the individual results will then be returned
+/// as the [`Cursor`] is iterated. When the batch is exhausted and if there are more results, the
+/// [`Cursor`] will fetch the next batch of documents, and so forth until the results are exhausted.
+/// Note that because of this batching, additional network I/O may occur on any given call to
+/// `next`. Because of this, a [`Cursor`] iterates over `Result<T>` items rather than
+/// simply `T` items.
 ///
 /// The batch size of the `Cursor` can be configured using the options to the method that returns
 /// it. For example, setting the `batch_size` field of
@@ -38,45 +39,42 @@ use common::{GenericCursor, GetMoreProvider, GetMoreProviderResult};
 /// results from the server; both of these factors should be taken into account when choosing the
 /// optimal batch size.
 ///
-/// A cursor can be used like any other [`Stream`](https://docs.rs/futures/0.3.4/futures/stream/trait.Stream.html). The simplest way is just to iterate over the
-/// documents it yields:
+/// [`Cursor`] implements [`Stream`](https://docs.rs/futures/latest/futures/stream/trait.Stream.html), which means
+/// it can be iterated over much in the same way that an `Iterator` can be in synchronous Rust. In
+/// order to do so, the [`StreamExt`](https://docs.rs/futures/latest/futures/stream/trait.StreamExt.html) trait must
+/// be imported. Because a [`Cursor`] iterates over a `Result<T>`, it also has access to the
+/// potentially more ergonomic functionality provided by
+/// [`TryStreamExt`](https://docs.rs/futures/latest/futures/stream/trait.TryStreamExt.html), which can be
+/// imported instead of or in addition to
+/// [`StreamExt`](https://docs.rs/futures/latest/futures/stream/trait.StreamExt.html). The methods from
+/// [`TryStreamExt`](https://docs.rs/futures/latest/futures/stream/trait.TryStreamExt.html) are especially useful when
+/// used in conjunction with the `?` operator.
 ///
 /// ```rust
-/// # use futures::stream::StreamExt;
 /// # use mongodb::{bson::Document, Client, error::Result};
 /// #
 /// # async fn do_stuff() -> Result<()> {
 /// # let client = Client::with_uri_str("mongodb://example.com").await?;
 /// # let coll = client.database("foo").collection::<Document>("bar");
-/// # let mut cursor = coll.find(None, None).await?;
 /// #
+/// use futures::stream::{StreamExt, TryStreamExt};
+///
+/// let mut cursor = coll.find(None, None).await?;
+/// // regular Stream uses next() and iterates over Option<Result<T>>
 /// while let Some(doc) = cursor.next().await {
 ///   println!("{}", doc?)
 /// }
-/// #
-/// # Ok(())
-/// # }
-/// ```
-///
-/// Additionally, all the other methods that an [`Stream`](https://docs.rs/futures/0.3.4/futures/stream/trait.Stream.html) has are available on `Cursor` as well.
-/// This includes all of the functionality provided by [`StreamExt`](https://docs.rs/futures/0.3.4/futures/stream/trait.StreamExt.html), which provides similar functionality to the standard library `Iterator` trait.
-/// For instance, if the number of results from a query is known to be small, it might make sense
-/// to collect them into a vector:
+/// // regular Stream uses collect() and collects into a Vec<Result<T>>
+/// let v: Vec<Result<_>> = cursor.collect().await;
 ///
-/// ```rust
-/// # use futures::stream::StreamExt;
-/// # use mongodb::{
-/// #     bson::{doc, Document},
-/// #     error::Result,
-/// #     Client,
-/// # };
-/// #
-/// # async fn do_stuff() -> Result<()> {
-/// # let client = Client::with_uri_str("mongodb://example.com").await?;
-/// # let coll = client.database("foo").collection("bar");
-/// # let cursor = coll.find(Some(doc! { "x": 1 }), None).await?;
+/// let mut cursor = coll.find(None, None).await?;
+/// // TryStream uses try_next() and iterates over Result<Option<T>>
+/// while let Some(doc) = cursor.try_next().await? {
+///   println!("{}", doc)
+/// }
+/// // TryStream uses try_collect() and collects into a Result<Vec<T>>
+/// let v: Vec<_> = cursor.try_collect().await?;
 /// #
-/// let results: Vec<Result<Document>> = cursor.collect().await;
 /// # Ok(())
 /// # }
 /// ```

src/cursor/session.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use futures_core::{future::BoxFuture, Stream};
+use futures_util::StreamExt;
 use serde::de::DeserializeOwned;
 
 use super::common::{CursorInformation, GenericCursor, GetMoreProvider, GetMoreProviderResult};
@@ -19,22 +20,29 @@ use crate::{
     RUNTIME,
 };
 
-/// A `SessionCursor` is a cursor that was created with a `ClientSession` and must be iterated using
-/// one. To iterate, retrieve a `SessionCursorHandle` using `SessionCursor::with_session`:
+/// A [`SessionCursor`] is a cursor that was created with a [`ClientSession`] that must be iterated
+/// using one. To iterate, use [`SessionCursor::next`] or retrieve a [`SessionCursorStream`] using
+/// [`SessionCursor::stream`]:
 ///
 /// ```rust
-/// # use futures::stream::StreamExt;
 /// # use mongodb::{bson::Document, Client, error::Result, ClientSession, SessionCursor};
 /// #
 /// # async fn do_stuff() -> Result<()> {
 /// # let client = Client::with_uri_str("mongodb://example.com").await?;
 /// # let mut session = client.start_session(None).await?;
 /// # let coll = client.database("foo").collection::<Document>("bar");
-/// # let mut cursor = coll.find_with_session(None, None, &mut session).await?;
 /// #
-/// while let Some(doc) = cursor.with_session(&mut session).next().await {
-///     println!("{}", doc?)
+/// // iterate using next()
+/// let mut cursor = coll.find_with_session(None, None, &mut session).await?;
+/// while let Some(doc) = cursor.next(&mut session).await.transpose()? {
+///     println!("{}", doc)
 /// }
+///
+/// // iterate using `Stream`:
+/// use futures::stream::TryStreamExt;
+///
+/// let mut cursor = coll.find_with_session(None, None, &mut session).await?;
+/// let results: Vec<_> = cursor.stream(&mut session).try_collect().await?;
 /// #
 /// # Ok(())
 /// # }
@@ -67,12 +75,52 @@ where
         }
     }
 
-    /// Retrieves a `SessionCursorHandle` to iterate this cursor. The session provided must be the
+    /// Retrieves a [`SessionCursorStream`] to iterate this cursor. The session provided must be the
     /// same session used to create the cursor.
-    pub fn with_session<'session>(
+    ///
+    /// Note that the borrow checker will not allow the session to be reused in between iterations
+    /// of this stream. In order to do that, either use [`SessionCursor::next`] instead or drop
+    /// the stream before using the session.
+    ///
+    /// ```
+    /// # use bson::{doc, Document};
+    /// # use mongodb::{Client, error::Result};
+    /// # fn main() {
+    /// # async {
+    /// # let client = Client::with_uri_str("foo").await?;
+    /// # let coll = client.database("foo").collection::<Document>("bar");
+    /// # let other_coll = coll.clone();
+    /// # let mut session = client.start_session(None).await?;
+    /// #
+    /// use futures::stream::TryStreamExt;
+    ///
+    /// // iterate over the results
+    /// let mut cursor = coll.find_with_session(doc! { "x": 1 }, None, &mut session).await?;
+    /// while let Some(doc) = cursor.stream(&mut session).try_next().await? {
+    ///     println!("{}", doc);
+    /// }
+    ///
+    /// // collect the results
+    /// let mut cursor1 = coll.find_with_session(doc! { "x": 1 }, None, &mut session).await?;
+    /// let v: Vec<Document> = cursor1.stream(&mut session).try_collect().await?;
+    ///
+    /// // use session between iterations
+    /// let mut cursor2 = coll.find_with_session(doc! { "x": 1 }, None, &mut session).await?;
+    /// loop {
+    ///     let doc = match cursor2.stream(&mut session).try_next().await? {
+    ///         Some(d) => d,
+    ///         None => break,
+    ///     };
+    ///     other_coll.insert_one_with_session(doc, None, &mut session).await?;
+    /// }
+    /// # Ok::<(), mongodb::error::Error>(())
+    /// # };
+    /// # }
+    /// ```
+    pub fn stream<'session>(
         &mut self,
         session: &'session mut ClientSession,
-    ) -> SessionCursorHandle<'_, 'session, T> {
+    ) -> SessionCursorStream<'_, 'session, T> {
         let get_more_provider = ExplicitSessionGetMoreProvider::new(session);
 
         // Pass the buffer into this cursor handle for iteration.
@@ -81,7 +129,7 @@ where
             info: self.info.clone(),
             initial_buffer: std::mem::take(&mut self.buffer),
         };
-        SessionCursorHandle {
+        SessionCursorStream {
             generic_cursor: ExplicitSessionCursor::new(
                 self.client.clone(),
                 spec,
@@ -90,6 +138,33 @@ where
             session_cursor: self,
         }
     }
+
+    /// Retrieve the next result from the cursor.
+    /// The session provided must be the same session used to create the cursor.
+    ///
+    /// Use this method when the session needs to be used again between iterations or when the added
+    /// functionality of `Stream` is not needed.
+    ///
+    /// ```
+    /// # use bson::{doc, Document};
+    /// # use mongodb::Client;
+    /// # fn main() {
+    /// # async {
+    /// # let client = Client::with_uri_str("foo").await?;
+    /// # let coll = client.database("foo").collection::<Document>("bar");
+    /// # let other_coll = coll.clone();
+    /// # let mut session = client.start_session(None).await?;
+    /// let mut cursor = coll.find_with_session(doc! { "x": 1 }, None, &mut session).await?;
+    /// while let Some(doc) = cursor.next(&mut session).await.transpose()? {
+    ///     other_coll.insert_one_with_session(doc, None, &mut session).await?;
+    /// }
+    /// # Ok::<(), mongodb::error::Error>(())
+    /// # };
+    /// # }
+    /// ```
+    pub async fn next(&mut self, session: &mut ClientSession) -> Option<Result<T>> {
+        self.stream(session).next().await
+    }
 }
 
 impl<T> Drop for SessionCursor<T>
@@ -115,19 +190,20 @@ where
 /// This is to be used with cursors associated with explicit sessions borrowed from the user.
 type ExplicitSessionCursor<'session> = GenericCursor<ExplicitSessionGetMoreProvider<'session>>;
 
-/// A handle that borrows a `ClientSession` temporarily for executing getMores or iterating through
-/// the current buffer of a `SessionCursor`.
+/// A type that implements [`Stream`](https://docs.rs/futures/latest/futures/stream/index.html) which can be used to
+/// stream the results of a [`SessionCursor`]. Returned from [`SessionCursor::stream`].
 ///
-/// This updates the buffer of the parent `SessionCursor` when dropped.
-pub struct SessionCursorHandle<'cursor, 'session, T = Document>
+/// This updates the buffer of the parent [`SessionCursor`] when dropped. [`SessionCursor::next`] or
+/// any further streams created from [`SessionCursor::stream`] will pick up where this one left off.
+pub struct SessionCursorStream<'cursor, 'session, T = Document>
 where
     T: DeserializeOwned + Unpin,
 {
     session_cursor: &'cursor mut SessionCursor<T>,
     generic_cursor: ExplicitSessionCursor<'session>,
 }
 
-impl<'cursor, 'session, T> Stream for SessionCursorHandle<'cursor, 'session, T>
+impl<'cursor, 'session, T> Stream for SessionCursorStream<'cursor, 'session, T>
 where
     T: DeserializeOwned + Unpin,
 {
@@ -144,7 +220,7 @@ where
     }
 }
 
-impl<'cursor, 'session, T> Drop for SessionCursorHandle<'cursor, 'session, T>
+impl<'cursor, 'session, T> Drop for SessionCursorStream<'cursor, 'session, T>
 where
     T: DeserializeOwned + Unpin,
 {

src/db/mod.rs

@@ -278,7 +278,7 @@ impl Database {
             .await
             .map(|spec| SessionCursor::new(self.client().clone(), spec))?;
 
-        self.list_collection_names_common(cursor.with_session(session))
+        self.list_collection_names_common(cursor.stream(session))
             .await
     }
 

src/lib.rs

@@ -136,15 +136,15 @@ define_if_single_runtime_enabled! {
     pub use crate::{
         client::Client,
         coll::Collection,
-        cursor::{Cursor, session::{SessionCursor, SessionCursorHandle}},
+        cursor::{Cursor, session::{SessionCursor, SessionCursorStream}},
         db::Database,
     };
 
     #[cfg(feature = "sync")]
     pub(crate) use crate::{
         client::Client,
         coll::Collection,
-        cursor::{Cursor, session::{SessionCursor, SessionCursorHandle}},
+        cursor::{Cursor, session::{SessionCursor, SessionCursorStream}},
         db::Database,
     };