mongor 0.1.10

// Authors: Robert Lopez
// License: MIT (See `LICENSE.md`)

use crate::error::Error;
use mongodb::{bson::Document, options::*, ClientSession, Collection, Cursor, SessionCursor};
use serde::{Deserialize, Serialize};

/// Executes a `findOne` query on `Collection`
///
/// https://www.mongodb.com/docs/manual/reference/method/db.collection.findOne/
///
/// ---
/// Example Usage:
/// ```
///
/// let client: Client = ...;
///
/// let collection: Collection<D> = ...;
///
/// let mut session: ClientSession = start_session(&client, None).await?;
///
/// let document_option: Option<D> = find_one(
///     &collection,
///     doc! { ... },
///     None::<FindOneOptions>,
///     Some(&mut session),
/// ).await?;
/// ```
pub async fn find_one<D>(
    collection: &Collection<D>,
    filter: Document,
    options: Option<FindOneOptions>,
    session: Option<&mut ClientSession>,
) -> Result<Option<D>, Error>
where
    D: Serialize + for<'a> Deserialize<'a> + Send + Sync + Unpin,
{
    let document = if let Some(session) = session {
        collection
            .find_one(filter)
            .with_options(options)
            .session(session)
            .await
    } else {
        collection.find_one(filter).with_options(options).await
    }
    .map_err(Error::Mongo)?;

    Ok(document)
}

/// A Cursor over a `find` query
///
/// https://www.mongodb.com/docs/manual/reference/method/db.collection.find/
///
/// ---
/// Example Usage:
/// ```
///
/// let client: Client = ...;
///
/// let collection: Collection<D> = ...;
///
/// let mut session: ClientSession = start_session(&client, None).await?;
///
/// let mut cursor = FindManyCursor::from(
///     &collection,
///     doc! { ... },
///     None::<FindOptions>,
///     Some(&mut session),
/// ).await?;
///
/// // Iteration:
/// while let Some(document: D) = cursor.next_with_session(&mut session).await? {
///     ...
/// }
///
/// // Or to get all at once:
/// let documents: Vec<D> = cursor.all_with_session().await?;
/// ```
pub struct FindManyCursor<D>
where
    D: Serialize + for<'a> Deserialize<'a> + Send + Sync,
{
    cursor: Option<Cursor<D>>,
    session_cursor: Option<SessionCursor<D>>,
}

impl<D> FindManyCursor<D>
where
    D: Serialize + for<'a> Deserialize<'a> + Send + Sync,
{
    /// Executes a `find` query on `Collection` and returns a Cursor to iterate over
    ///
    /// https://www.mongodb.com/docs/manual/reference/method/db.collection.find/
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let client: Client = ...;
    ///
    /// let collection: Collection<D> = ...;
    ///
    /// let mut session: ClientSession = start_session(&client, None).await?;
    ///
    /// let mut cursor = FindManyCursor::from(
    ///     &collection,
    ///     doc! { ... },
    ///     None::<FindOptions>,
    ///     Some(&mut session),
    /// ).await?;
    /// ```
    pub async fn from(
        collection: &Collection<D>,
        filter: Document,
        options: Option<FindOptions>,
        session: Option<&mut ClientSession>,
    ) -> Result<Self, Error> {
        let (cursor, session_cursor) = if let Some(session) = session {
            (
                None,
                Some(
                    collection
                        .find(filter)
                        .with_options(options)
                        .session(session)
                        .await
                        .map_err(Error::Mongo)?,
                ),
            )
        } else {
            (
                Some(
                    collection
                        .find(filter)
                        .with_options(options)
                        .await
                        .map_err(Error::Mongo)?,
                ),
                None,
            )
        };

        Ok(Self {
            cursor,
            session_cursor,
        })
    }

    /// Constructs `FindManyCursor<D>` from a `Cursor<D>`.
    ///
    /// This is used to construct the `FindManyCursor` for
    /// `GridFS::find_many`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let cursor: Cursor<D> = ...;
    ///
    /// let find_many_cursor = FindManyCursor::from_cursor(cursor);
    /// ```
    #[cfg(feature = "grid_fs")]

    pub fn from_cursor(cursor: Cursor<D>) -> Self {
        Self {
            cursor: Some(cursor),
            session_cursor: None,
        }
    }

    /// Yields the next available document in the internal `Cursor`, or `None` if there
    /// are no more to be yielded.
    ///
    /// If a `ClientSession` was used in `FindManyCursor::from`, please use `next_with_session`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut cursor: FindManyCursor<D> = ...;
    ///
    /// while let Some(document: D) = cursor.next().await? {
    ///     ...
    /// }
    /// ```
    pub async fn next(&mut self) -> Result<Option<D>, Error> {
        if let Some(cursor) = &mut self.cursor {
            if cursor.advance().await.map_err(Error::Mongo)? {
                return Ok(Some(cursor.deserialize_current().map_err(Error::Mongo)?));
            } else {
                return Ok(None);
            }
        }

        Err(Error::internal(
            "Cannot call `next` on `FindManyCursor`\
         when using a `ClientSession`. Please use `next_with_session`",
        ))
    }

    /// Yields the next available document in the internal `SessionCursor`, or `None`
    /// if there are no more to be yielded.
    ///
    /// If a `ClientSession` was not used in `FindManyCursor::from`, please use `next`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let client: Client = ...;
    ///
    /// let mut session: ClientSession = start_session(&client, None).await?;
    ///
    /// let mut cursor: FindManyCursor<D> = ...;
    ///
    /// while let Some(document: D) = cursor.next_with_session(&mut session).await? {
    ///     ...
    /// }
    /// ```
    pub async fn next_with_session(
        &mut self,
        session: &mut ClientSession,
    ) -> Result<Option<D>, Error> {
        if let Some(session_cursor) = &mut self.session_cursor {
            if session_cursor
                .advance(session)
                .await
                .map_err(Error::Mongo)?
            {
                return Ok(Some(
                    session_cursor.deserialize_current().map_err(Error::Mongo)?,
                ));
            } else {
                return Ok(None);
            }
        }

        Err(Error::internal(
            "Cannot call `next_with_session` on `FindManyCursor`\
         without using a `ClientSession`. Please use `next`",
        ))
    }

    /// Yields all documents in the internal `Cursor`.
    ///
    /// If a `ClientSession` was used in `FindManyCursor::from`, please use `all_with_session`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut cursor: FindManyCursor<D> = ...;
    ///
    /// let documents: Vec<D> = cursor.all().await?;
    /// ```
    pub async fn all(mut self) -> Result<Vec<D>, Error> {
        if self.cursor.is_none() {
            return Err(Error::internal(
                "Cannot call `all` on `FindManyCursor`\
             when using a `ClientSession`. Please use `all_with_session`",
            ));
        }

        let mut documents = vec![];

        while let Some(document) = self.next().await? {
            documents.push(document);
        }

        Ok(documents)
    }

    /// Yields all documents in the internal `SessionCursor`.
    ///
    /// If a `ClientSession` was not used in `FindManyCursor::from`, please use `all`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let client: Client = ...;
    ///
    /// let mut session: ClientSession = start_session(&client, None).await?;
    ///
    /// let mut cursor: FindManyCursor<D> = ...;
    ///
    /// let documents: Vec<D> = cursor.all_with_session(&mut session).await?;
    /// ```
    pub async fn all_with_session(mut self, session: &mut ClientSession) -> Result<Vec<D>, Error> {
        if self.session_cursor.is_none() {
            return Err(Error::internal(
                "Cannot call `all_with_session` on `FindManyCursor`\
                without using a `ClientSession`. Please use `all`",
            ));
        }

        let mut documents = vec![];

        while let Some(document) = self.next_with_session(session).await? {
            documents.push(document);
        }

        Ok(documents)
    }
}