RUST-1478 GridFS upload methods

src/gridfs.rs

@@ -3,6 +3,7 @@
 
 mod download;
 pub mod options;
+mod upload;
 
 use core::task::{Context, Poll};
 use std::{
@@ -241,64 +242,6 @@ impl GridFsBucket {
         todo!()
     }
 
-    /// Uploads a user file to a GridFS bucket. The application supplies a custom file id. Uses the
-    /// `tokio` crate's `AsyncRead` trait for the `source`.
-    pub async fn upload_from_tokio_reader_with_id(
-        &self,
-        id: Bson,
-        filename: String,
-        source: impl tokio::io::AsyncRead,
-        options: impl Into<Option<GridFsUploadOptions>>,
-    ) {
-        todo!()
-    }
-
-    /// Uploads a user file to a GridFS bucket. The application supplies a custom file id. Uses the
-    /// `futures-0.3` crate's `AsyncRead` trait for the `source`.
-    pub async fn upload_from_futures_0_3_reader_with_id(
-        &self,
-        id: Bson,
-        filename: String,
-        source: impl futures_util::AsyncRead,
-        options: impl Into<Option<GridFsUploadOptions>>,
-    ) {
-        todo!()
-    }
-
-    /// Uploads a user file to a GridFS bucket. The driver generates a unique [`Bson::ObjectId`] for
-    /// the file id. Uses the `tokio` crate's `AsyncRead` trait for the `source`.
-    pub async fn upload_from_tokio_reader(
-        &self,
-        filename: String,
-        source: impl tokio::io::AsyncRead,
-        options: impl Into<Option<GridFsUploadOptions>>,
-    ) {
-        self.upload_from_tokio_reader_with_id(
-            Bson::ObjectId(ObjectId::new()),
-            filename,
-            source,
-            options,
-        )
-        .await
-    }
-
-    /// Uploads a user file to a GridFS bucket. The driver generates a unique [`Bson::ObjectId`] for
-    /// the file id. Uses the `futures-0.3` crate's `AsyncRead` trait for the `source`.
-    pub async fn upload_from_futures_0_3_reader(
-        &self,
-        filename: String,
-        source: impl futures_util::AsyncRead,
-        options: impl Into<Option<GridFsUploadOptions>>,
-    ) {
-        self.upload_from_futures_0_3_reader_with_id(
-            Bson::ObjectId(ObjectId::new()),
-            filename,
-            source,
-            options,
-        )
-        .await
-    }
-
     /// Opens a [`GridFsUploadStream`] that the application can write the contents of the file to.
     /// The driver generates a unique [`Bson::ObjectId`] for the file id.
     ///

src/gridfs/upload.rs

@@ -0,0 +1,162 @@
+use std::{marker::Unpin, sync::atomic::Ordering};
+
+use futures_util::{
+    io::{AsyncRead, AsyncReadExt},
+    stream::TryStreamExt,
+};
+
+use super::{options::GridFsUploadOptions, Chunk, FilesCollectionDocument, GridFsBucket};
+use crate::{
+    bson::{doc, oid::ObjectId, spec::BinarySubtype, Bson, DateTime, Document, RawBinaryRef},
+    bson_util::get_int,
+    error::{ErrorKind, Result},
+    index::IndexModel,
+    options::{FindOneOptions, ReadPreference, SelectionCriteria},
+    Collection,
+};
+
+impl GridFsBucket {
+    /// Uploads a user file to a GridFS bucket. Bytes are read from `source` and stored in chunks in
+    /// the bucket's chunks collection. After all the chunks have been uploaded, a corresponding
+    /// [`FilesCollectionDocument`] is stored in the bucket's files collection.
+    ///
+    /// This method generates an [`ObjectId`] for the `files_id` field of the
+    /// [`FilesCollectionDocument`] and returns it.
+    pub async fn upload_from_futures_0_3_reader<T>(
+        &self,
+        filename: impl AsRef<str>,
+        source: T,
+        options: impl Into<Option<GridFsUploadOptions>>,
+    ) -> Result<ObjectId>
+    where
+        T: AsyncRead + Unpin,
+    {
+        let id = ObjectId::new();
+        self.upload_from_futures_0_3_reader_with_id(id.into(), filename, source, options)
+            .await?;
+        Ok(id)
+    }
+
+    /// Uploads a user file to a GridFS bucket with the given `files_id`. Bytes are read from
+    /// `source` and stored in chunks in the bucket's chunks collection. After all the chunks have
+    /// been uploaded, a corresponding [`FilesCollectionDocument`] is stored in the bucket's files
+    /// collection.
+    pub async fn upload_from_futures_0_3_reader_with_id<T>(
+        &self,
+        files_id: Bson,
+        filename: impl AsRef<str>,
+        mut source: T,
+        options: impl Into<Option<GridFsUploadOptions>>,
+    ) -> Result<()>
+    where
+        T: AsyncRead + Unpin,
+    {
+        let options = options.into();
+
+        self.create_indexes().await?;
+
+        let chunk_size = options
+            .as_ref()
+            .and_then(|opts| opts.chunk_size_bytes)
+            .unwrap_or_else(|| self.chunk_size_bytes());
+        let mut length = 0u64;
+        let mut n = 0;
+
+        let mut buf = vec![0u8; chunk_size as usize];
+        loop {
+            let bytes_read = match source.read(&mut buf).await {
+                Ok(0) => break,
+                Ok(n) => n,
+                Err(error) => {
+                    self.chunks()
+                        .delete_many(doc! { "files_id": &files_id }, None)
+                        .await?;
+                    return Err(ErrorKind::Io(error.into()).into());
+                }
+            };
+
+            let chunk = Chunk {
+                id: ObjectId::new(),
+                files_id: files_id.clone(),
+                n,
+                data: RawBinaryRef {
+                    subtype: BinarySubtype::Generic,
+                    bytes: &buf[..bytes_read],
+                },
+            };
+            self.chunks().insert_one(chunk, None).await?;
+
+            length += bytes_read as u64;
+            n += 1;
+        }
+
+        let file = FilesCollectionDocument {
+            id: files_id,
+            length,
+            chunk_size,
+            upload_date: DateTime::now(),
+            filename: Some(filename.as_ref().to_string()),
+            metadata: options.and_then(|opts| opts.metadata),
+        };
+        self.files().insert_one(file, None).await?;
+
+        Ok(())
+    }
+
+    async fn create_indexes(&self) -> Result<()> {
+        if !self.inner.created_indexes.load(Ordering::SeqCst) {
+            let find_options = FindOneOptions::builder()
+                .selection_criteria(SelectionCriteria::ReadPreference(ReadPreference::Primary))
+                .projection(doc! { "_id": 1 })
+                .build();
+            if self
+                .files()
+                .clone_with_type::<Document>()
+                .find_one(None, find_options)
+                .await?
+                .is_none()
+            {
+                Self::create_index(self.files(), doc! { "filename": 1, "uploadDate": 1 }).await?;
+                Self::create_index(self.chunks(), doc! { "files_id": 1, "n": 1 }).await?;
+            }
+            self.inner.created_indexes.store(true, Ordering::SeqCst);
+        }
+
+        Ok(())
+    }
+
+    async fn create_index<T>(coll: &Collection<T>, keys: Document) -> Result<()> {
+        // From the spec: Drivers MUST check whether the indexes already exist before attempting to
+        // create them.
+        let mut indexes = coll.list_indexes(None).await?;
+        'outer: while let Some(index_model) = indexes.try_next().await? {
+            if index_model.keys.len() != keys.len() {
+                continue;
+            }
+            // Indexes should be considered equivalent regardless of numeric value type.
+            // e.g. { "filename": 1, "uploadDate": 1 } is equivalent to
+            // { "filename": 1.0, "uploadDate": 1.0 }
+            let number_matches = |key: &str, value: &Bson| {
+                if let Some(model_value) = index_model.keys.get(key) {
+                    match get_int(value) {
+                        Some(num) => get_int(model_value) == Some(num),
+                        None => model_value == value,
+                    }
+                } else {
+                    false
+                }
+            };
+            for (key, value) in keys.iter() {
+                if !number_matches(key, value) {
+                    continue 'outer;
+                }
+            }
+            return Ok(());
+        }
+
+        let index_model = IndexModel::builder().keys(keys).build();
+        coll.create_index(index_model, None).await?;
+
+        Ok(())
+    }
+}

src/test/spec/unified_runner/operation.rs

@@ -38,7 +38,7 @@ use crate::{
     coll::options::Hint,
     collation::Collation,
     error::{ErrorKind, Result},
-    gridfs::options::GridFsDownloadByNameOptions,
+    gridfs::options::{GridFsDownloadByNameOptions, GridFsUploadOptions},
     options::{
         AggregateOptions,
         CountOptions,
@@ -342,6 +342,7 @@ impl<'de> Deserialize<'de> for Operation {
             "download" => deserialize_op::<Download>(definition.arguments),
             "downloadByName" => deserialize_op::<DownloadByName>(definition.arguments),
             "delete" => deserialize_op::<Delete>(definition.arguments),
+            "upload" => deserialize_op::<Upload>(definition.arguments),
             _ => Ok(Box::new(UnimplementedOperation) as Box<dyn TestOperation>),
         }
         .map_err(|e| serde::de::Error::custom(format!("{}", e)))?;
@@ -2720,6 +2721,43 @@ impl TestOperation for Delete {
         .boxed()
     }
 }
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct Upload {
+    source: Document,
+    filename: String,
+    // content_type and disableMD5 are deprecated and no longer supported.
+    // Options included for deserialization.
+    #[serde(rename = "contentType")]
+    _content_type: Option<String>,
+    #[serde(rename = "disableMD5")]
+    _disable_md5: Option<bool>,
+    #[serde(flatten)]
+    options: GridFsUploadOptions,
+}
+
+impl TestOperation for Upload {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let bucket = test_runner.get_bucket(id).await;
+            let hex_string = self.source.get("$$hexBytes").unwrap().as_str().unwrap();
+            let bytes = hex::decode(hex_string).unwrap();
+            let id = bucket
+                .upload_from_futures_0_3_reader(
+                    self.filename.clone(),
+                    &bytes[..],
+                    self.options.clone(),
+                )
+                .await?;
+            Ok(Some(Entity::Bson(id.into())))
+        }
+        .boxed()
+    }
+}
 
 #[derive(Debug, Deserialize)]
 pub(super) struct UnimplementedOperation;

src/test/spec/unified_runner/test_runner.rs

@@ -60,7 +60,6 @@ const SKIPPED_OPERATIONS: &[&str] = &[
     "listCollectionObjects",
     "listDatabaseObjects",
     "mapReduce",
-    "upload",
     "watch",
 ];