mongor 0.1.10

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

#[cfg(feature = "grid_fs")]
use crate::{core::find::FindManyCursor, error::Error, util::convert_bson_to_oid};
use futures_util::{AsyncRead, AsyncWrite, AsyncWriteExt};
#[cfg(feature = "grid_fs")]
#[cfg(feature = "grid_fs")]
use mongodb::{
    bson::{doc, oid::ObjectId, Bson, Document},
    gridfs::{FilesCollectionDocument, GridFsBucket, GridFsDownloadStream, GridFsUploadStream},
    options::*,
    Database,
};
// use tokio::io::{copy, AsyncRead, AsyncWrite};
// use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};

#[cfg(feature = "grid_fs")]
/// A structure to control a `GridFsBucket`.
///
/// https://www.mongodb.com/docs/manual/core/gridfs/
///
/// ---
/// Example Usage:
/// ```
///
/// let db: Database = ...;
///
/// let options = GridFsBucketOptions::builder()
///     .bucket_name(Some("bucket_name".to_string()))
///     .build();
///
/// let grid_fs = GridFs::new(&db, Some(options));
///
/// let file: tokio::fs::File = File::open("some path").await?;
/// let file_stream = file.compat();
///
/// grid_fs.upload("filename", file_stream, None).await?;
///
/// let mut writer: futures_util::io::Cursor<Vec<u8>> = Cursor::new(vec![]);
/// grid_fs.download("filename", &mut writer, None).await?;
///
/// grid_fs.rename("filename", "new_filename", None).await?;
///
/// // Most methods take a Option<i32> for a revision number
/// grid_fs.delete_by_filename("filename", Some(-1)).await?;
/// ```
pub struct GridFs {
    pub bucket: GridFsBucket,
}

#[cfg(feature = "grid_fs")]
impl GridFs {
    /// Private helper to build `GridFsDownloadByNameOptions`
    fn _build_download_options(
        &self,
        revision: Option<i32>,
    ) -> Option<GridFsDownloadByNameOptions> {
        if let Some(revision) = revision {
            Some(
                GridFsDownloadByNameOptions::builder()
                    .revision(revision)
                    .build(),
            )
        } else {
            None
        }
    }

    /// Construct a new `GridFs` containing a single bucket.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let db: Database = ...;
    ///
    /// let options = GridFsBucketOptions::builder()
    ///     .bucket_name(Some("bucket_name".to_string()))
    ///     .build();
    ///
    /// let grid_fs = GridFs::new(&db, Some(options));
    /// ```
    pub fn new(db: &Database, options: Option<GridFsBucketOptions>) -> Self {
        let bucket = db.gridfs_bucket(options);

        Self { bucket }
    }

    /// Drops all files and chunks from `self.bucket`
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// grid_fs.drop().await?;
    /// ```
    pub async fn drop(&self) -> Result<(), Error> {
        self.bucket.drop().await.map_err(Error::Mongo)
    }

    /// Returns the size in bytes and the number of files
    /// in `self.bucket`.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let (total_bytes, total_files) = grid_fs.size().await?;
    /// ```
    pub async fn size(&self) -> Result<(usize, usize), Error> {
        let mut total_bytes = 0;
        let mut total_files = 0;

        let mut cursor = self.find_many(doc! {}, None).await?;

        while let Some(FilesCollectionDocument { length, .. }) = cursor.next().await? {
            total_files += 1;
            total_bytes += length as usize;
        }

        Ok((total_bytes, total_files))
    }

    /// Returns true if a file by `filename` exists in the bucket.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// if grid_fs.exists("filename").await? {
    ///     ...
    /// }
    /// ```
    pub async fn exists(&self, filename: &str) -> Result<bool, Error> {
        Ok(self
            .find_many(
                doc! { "filename": filename },
                Some(GridFsFindOptions::builder().limit(1).build()),
            )
            .await?
            .next()
            .await?
            .is_some())
    }

    /// Find one file document by a filter.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let document_option: Option<FilesCollectionDocument> = grid_fs.find_one(
    ///     doc! { ... },
    ///     None::<Document>,
    /// ).await?;
    /// ```
    pub async fn find_one(
        &self,
        filter: Document,
        sort: Option<Document>,
    ) -> Result<Option<FilesCollectionDocument>, Error> {
        Ok(self
            .find_many(
                filter,
                Some(GridFsFindOptions::builder().limit(1).sort(sort).build()),
            )
            .await?
            .next()
            .await?)
    }

    /// Find one file document by a filename and a specific revision.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let document_option: Option<FilesCollectionDocument> = grid_fs.find_one_by_filename_revision(
    ///     "filename",
    ///     -2,
    /// ).await?;
    /// ```
    async fn find_one_by_filename_revision(
        &self,
        filename: &str,
        revision: i32,
    ) -> Result<Option<FilesCollectionDocument>, Error> {
        let (sort, skip) = if revision >= 0 {
            (1, revision)
        } else {
            (-1, -revision - 1)
        };

        let options = GridFsFindOptions::builder()
            .sort(doc! { "uploadDate": sort })
            .skip(skip as u64)
            .limit(Some(1))
            .build();

        let mut cursor = self
            .find_many(doc! { "filename": filename }, Some(options))
            .await?;

        cursor.next().await
    }

    /// Upload a file to GridFS via a `source` returning its `ObjectId`.
    ///
    /// The `mongodb` crate implements io to the bucket via `futures_util`,
    /// so the source `S` must implement `futures_io::AsyncRead + Unpin`.
    ///
    /// However, if you are using "tokio", you can call `.compat()` to
    /// wrap it into a source that implements `futures_io::AsyncRead`.
    ///
    /// See: `https://docs.rs/tokio-util/latest/tokio_util/compat/index.html`
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("some file path").await?;
    ///
    /// let file_metadata = doc! { ... };
    /// let options = GridFsUploadOptions::builder()
    ///     .metadata(Some(file_metadata))
    ///     .build();
    ///
    /// let oid: ObjectId = grid_fs.upload(
    ///     "file_one",
    ///     file.compat(),
    ///     Some(options),
    /// ).await?;
    /// ```
    pub async fn upload<S>(
        &self,
        filename: &str,
        mut source: S,
        options: Option<GridFsUploadOptions>,
    ) -> Result<ObjectId, Error>
    where
        S: AsyncRead + Unpin,
    {
        let mut upload_stream = self.open_upload_stream(filename, options).await?;
        let upload_id = convert_bson_to_oid(upload_stream.id().clone())?;
        futures_util::io::copy(&mut source, &mut upload_stream)
            .await
            .map_err(|err| Error::IO(err.kind()))?;
        upload_stream
            .close()
            .await
            .map_err(|err| Error::IO(err.kind()))?;

        Ok(upload_id)
    }

    /// Open a `GridFsUploadStream` via a files `filename`. You can
    /// obtain the `ObjectId` from `GridFsUploadStream.id`.
    ///
    /// If multiple files under the same `filename` exist, and you want
    /// to download a specific revision (instead of the most recent) then
    /// provide a `revision` number.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let upload_stream: GridFsUploadOptions = grid_fs
    ///     .open_download_stream(
    ///         "filename",
    ///         None::<GridFsUploadOptions>,
    ///     )
    ///     .await?;
    /// ```
    pub async fn open_upload_stream(
        &self,
        filename: &str,
        options: Option<GridFsUploadOptions>,
    ) -> Result<GridFsUploadStream, Error> {
        self.bucket
            .open_upload_stream(filename)
            .with_options(options)
            .await
            .map_err(Error::Mongo)
    }

    /// Download a file from GridFS via a async writer `destination`.
    ///
    /// If multiple files under the same `filename` exist, and you want
    /// to download a specific revision (instead of the most recent) then
    /// provide a `revision` number.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// The `mongodb` crate implements io to the bucket via `futures_util`,
    /// so the destination `D` must implement `futures_io::AsyncWrite + Unpin`.
    ///
    /// However, if you are using "tokio", you can call `.compat_write()` to
    /// wrap it into a source that implements `futures_io::AsyncWrite`.
    ///
    /// See: `https://docs.rs/tokio-util/latest/tokio_util/compat/index.html`
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    ///
    /// grid_fs.upload(
    ///     "filename",
    ///     file.compat(),
    ///     None::<GridFsUploadOptions>,
    /// ).await?;
    ///
    /// // If writer implements tokio::AsyncWrite,
    /// // you would need to call `.compat_write()`
    /// let mut writer = futures_util::io::Cursor::new(vec![]);
    /// grid_fs.download(
    ///     "filename",
    ///     &mut writer,
    ///     None::<i32>,
    /// ).await?;
    ///
    /// let uploaded_data: Vec<u8> = writer.into_inner();
    /// ```
    pub async fn download<D>(
        &self,
        filename: &str,
        mut destination: &mut D,
        revision: Option<i32>,
    ) -> Result<usize, Error>
    where
        D: AsyncWrite + Unpin,
    {
        let mut download_stream = self.open_download_stream(filename, revision).await?;

        futures_util::io::copy(&mut download_stream, &mut destination)
            .await
            .map(|bytes| bytes as usize)
            .map_err(|err| Error::IO(err.kind()))
    }

    /// Open a `GridFsDownloadStream` via a files `filename`.
    ///
    /// If multiple files under the same `filename` exist, and you want
    /// to download a specific revision (instead of the most recent) then
    /// provide a `revision` number.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    /// let file_stream = file.compat();
    ///
    /// grid_fs.upload("filename", file_stream, None).await?;
    ///
    /// let download_stream: GridFsDownloadStream = grid_fs
    ///     .open_download_stream(
    ///         "filename",
    ///         None::<i32>,
    ///     )
    ///     .await?;
    /// ```
    pub async fn open_download_stream(
        &self,
        filename: &str,
        revision: Option<i32>,
    ) -> Result<GridFsDownloadStream, Error> {
        let options = self._build_download_options(revision);

        self.bucket
            .open_download_stream_by_name(filename)
            .with_options(options)
            .await
            .map_err(Error::Mongo)
    }

    /// Find many files by a filter `Document`, returning
    /// a cursor to iterate over.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let mut cursor: FindManyCursor<FilesCollectionDocument> = grid_fs.find_many(
    ///     doc! { ... },
    ///     None::<GridFsFindOptions>,
    /// ).await?;
    ///
    /// // Iteration
    /// while let Some(file_document: FilesCollectionDocument) = cursor.next().await? {
    ///     ...
    /// }
    ///
    /// // Or to get all at once:
    /// let documents: Vec<FilesCollectionDocument> = cursor.all().await?;
    /// ```
    pub async fn find_many(
        &self,
        filter: Document,
        options: Option<GridFsFindOptions>,
    ) -> Result<FindManyCursor<FilesCollectionDocument>, Error> {
        let cursor = self
            .bucket
            .find(filter)
            .with_options(options)
            .await
            .map_err(Error::Mongo)?;

        Ok(FindManyCursor::from_cursor(cursor))
    }

    /// Delete many files by a filter `Document`.
    ///
    /// All file revisions will be deleted for those
    /// that match the filter, as the `filename` field
    /// is taken from the document and deleted by
    /// `GridFs::delete_by_filename` with no revision
    /// number.
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// grid_fs.delete_many(
    ///     doc! { ... },
    ///     None::<GridFsFindOptions>,
    /// ).await?;
    /// ```
    pub async fn delete_many(
        &self,
        filter: Document,
        options: Option<GridFsFindOptions>,
    ) -> Result<usize, Error> {
        let mut cursor = self.find_many(filter, options).await?;
        let mut deletion_count = 0;

        while let Some(FilesCollectionDocument { filename, .. }) = cursor.next().await? {
            if let Some(ref filename) = filename {
                deletion_count += self.delete_by_filename(filename, None).await?;
            }
        }

        Ok(deletion_count)
    }

    /// Rename a, or all, files by `filename`, returns the amount
    /// of files renamed.
    ///
    /// If multiple files under the same `filename` exist, and you want
    /// to rename a specific revision (instead of all revisions) then
    /// provide a `revision` number.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    ///
    /// grid_fs.upload(
    ///     "filename",
    ///     file.compat(),
    ///     None::<GridFsUploadOptions>,
    /// ).await?;
    ///
    /// let renamed_files: usize = grid_fs.rename_by_filename(
    ///     "filename",
    ///     "new_filename"
    ///     None::<i32>,
    /// ).await?;
    /// ```
    pub async fn rename_by_filename(
        &self,
        filename: &str,
        new_filename: &str,
        revision: Option<i32>,
    ) -> Result<usize, Error> {
        let mut rename_counter = 0;

        if let Some(revision) = revision {
            if let Some(FilesCollectionDocument { id, .. }) = self
                .find_one_by_filename_revision(filename, revision)
                .await?
            {
                self.rename_by_oid(convert_bson_to_oid(id)?, new_filename)
                    .await?;
                rename_counter += 1;
            }

            return Ok(rename_counter);
        }

        let mut cursor = self.find_many(doc! { "filename": filename }, None).await?;

        while let Some(FilesCollectionDocument { id, .. }) = cursor.next().await? {
            self.rename_by_oid(convert_bson_to_oid(id)?, new_filename)
                .await?;
            rename_counter += 1;
        }

        Ok(rename_counter)
    }

    /// Rename a file by its `_id` field
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    ///
    /// grid_fs.upload(
    ///     "filename",
    ///     file.compat(),
    ///     None::<GridFsUploadOptions>,
    /// ).await?;
    ///
    /// grid_fs.rename_by_oid(oid, "new_filename").await?;
    /// ```
    pub async fn rename_by_oid(&self, oid: ObjectId, new_filename: &str) -> Result<(), Error> {
        self.bucket
            .rename(Bson::ObjectId(oid), new_filename)
            .await
            .map_err(Error::Mongo)
    }

    /// Delete a, or all, files by `filename`, returns the amount
    /// of files deleted.
    ///
    /// If multiple files under the same `filename` exist, and you want
    /// to delete a specific revision (instead of all revisions) then
    /// provide a `revision` number.
    ///
    /// Revision numbers are defined as follows:
    /// ```
    /// 0 = the original stored file
    /// 1 = the first revision
    /// 2 = the second revision
    /// etc...
    /// -2 = the second most recent revision
    /// -1 = the most recent revision
    /// ```
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    ///
    /// grid_fs.upload(
    ///     "filename",
    ///     file.compat(),
    ///     None::<GridFsUploadOptions>,
    /// ).await?;
    ///
    /// let deleted_files: usize = grid_fs.delete_by_filename(
    ///     "filename",
    ///     "new_filename"
    ///     None::<i32>,
    /// ).await?;
    /// ```
    pub async fn delete_by_filename(
        &self,
        filename: &str,
        revision: Option<i32>,
    ) -> Result<usize, Error> {
        let mut deletion_count = 0;

        if let Some(revision) = revision {
            if let Some(FilesCollectionDocument { id, .. }) = self
                .find_one_by_filename_revision(filename, revision)
                .await?
            {
                self.delete_by_oid(convert_bson_to_oid(id)?).await?;
                deletion_count += 1;
            }

            return Ok(deletion_count);
        }

        let mut cursor = self.find_many(doc! { "filename": filename }, None).await?;

        while let Some(FilesCollectionDocument { id, .. }) = cursor.next().await? {
            self.delete_by_oid(convert_bson_to_oid(id)?).await?;

            deletion_count += 1;
        }

        Ok(deletion_count)
    }

    /// Delete a file by its `_id` field
    ///
    /// ---
    /// Example Usage:
    /// ```
    ///
    /// let grid_fs: GridFs = ...;
    ///
    /// let file = tokio::fs::File::open("path").await?;
    ///
    /// grid_fs.upload(
    ///     "filename",
    ///     file.compat(),
    ///     None::<GridFsUploadOptions>,
    /// ).await?;
    ///
    /// grid_fs.delete_by_oid(oid).await?;
    /// ```
    pub async fn delete_by_oid(&self, oid: ObjectId) -> Result<(), Error> {
        self.bucket
            .delete(Bson::ObjectId(oid))
            .await
            .map_err(Error::Mongo)
    }
}