mongor 0.1.10

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

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

/// A Cursor over a `aggregate` query.
///
/// https://www.mongodb.com/docs/manual/reference/method/db.collection.aggregate/
///
/// ---
/// Example Usage:
/// ```
///
/// let collection: Collection<D> = ...;
/// let pipeline: Vec<Document> = vec![
///     doc! {
///         "$match": {
///             ...
///         }
///     },
///     doc! {
///         "$addFields": {
///             ...
///         }
///     },
///     ...
/// ];
///
/// let mut session: ClientSession = ...;
///
/// let mut cursor = AggregateCursor::from(
///     &collection,
///     pipeline,
///     None::<AggregateOptions>,
///     Some(&mut session),
/// ).await?;
///
/// // We used `$addFields`, so we cannot cast the resulting documents as `D`, as a new field
/// // is appended to the documents, thus we  use `next_generic_with_session` to yield the
/// // generic `Document`s.
/// while let Some(document: Document) = cursor.next_generic_with_session(&mut session).await? {
///     ...
/// }
///
/// // I.E. If we didn't modify the documents of type `D` we can use `next_with_session`
/// while let Some(document: D) = cursor.next_with_session(&mut session).await? {
///     ...
/// }
///
/// // Or, the equivalent without a session
/// while let Some(document: Document) = cursor.next_generic().await? {
///     ...
/// }
///
/// while let Some(document: D) = cursor.next().await? {
///     ...
/// }
/// ```
pub struct AggregateCursor<D>
where
    D: Serialize + for<'a> Deserialize<'a> + Send + Sync,
{
    cursor: Option<Cursor<Document>>,
    session_cursor: Option<SessionCursor<Document>>,
    _phantom: PhantomData<D>,
}

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

        Ok(Self {
            cursor,
            session_cursor,
            _phantom: PhantomData,
        })
    }

    /// Yields the next available document converted to `D` in the internal `Cursor`,
    /// or `None` if there are no more to be yielded.
    ///
    /// If a `ClientSession` was used in `AggregateCursor::from`, please use `next_with_session`.
    ///
    /// If the `pipeline` used in `AggregateCursor::from` modified the structure from `D`,
    /// please use `next_generic`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut cursor: AggregateCursor<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)? {
                let generic_document = cursor.deserialize_current().map_err(Error::Mongo)?;

                let document: D = from_document(generic_document).map_err(Error::Bson)?;

                return Ok(Some(document));
            } else {
                return Ok(None);
            }
        }

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

    /// Yields the next available document converted to `D` in the internal `SessionCursor`,
    /// or `None` if there are no more to be yielded.
    ///
    /// If a `ClientSession` was not used in `AggregateCursor::from`, please use `next`.
    ///
    /// If the `pipeline` used in `AggregateCursor::from` modified the structure from `D`,
    /// please use `next_generic_with_session`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut session: ClientSession = ...;
    ///
    /// let mut cursor: AggregateCursor<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)?
            {
                let generic_document =
                    session_cursor.deserialize_current().map_err(Error::Mongo)?;

                let document: D = from_document(generic_document).map_err(Error::Bson)?;

                return Ok(Some(document));
            } else {
                return Ok(None);
            }
        }

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

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

                return Ok(Some(generic_document));
            } else {
                return Ok(None);
            }
        }

        Err(Error::internal(
            "Cannot call `next_generic` on `AggregateCursor`\
         when using a `ClientSession`. Please use `next_generic_with_session`",
        ))
    }

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

                return Ok(Some(generic_document));
            } else {
                return Ok(None);
            }
        }

        Err(Error::internal(
            "Cannot call `next_generic_with_session` on `AggregateCursor`\
         without using a `ClientSession`. Please use `next_generic`",
        ))
    }

    /// Yields all documents in the internal `Cursor` converted to `D`.
    ///
    /// If a `ClientSession` was used in `AggregateCursor::from`, please use `all_with_session`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut cursor: AggregateCursor<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 `AggregateCursor`\
             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` converted to `D`.
    ///
    /// If a `ClientSession` was not used in `AggregateCursor::from`, please use `all`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let mut session: ClientSession = ...;
    ///
    /// let mut cursor: AggregateCursor<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 `AggregateCursor`\
                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)
    }

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

        let mut documents = vec![];

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

        Ok(documents)
    }

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

        let mut documents = vec![];

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

        Ok(documents)
    }
}