|
| 1 | +use std::{marker::Unpin, sync::atomic::Ordering}; |
| 2 | + |
| 3 | +use futures_util::{ |
| 4 | + io::{AsyncRead, AsyncReadExt}, |
| 5 | + stream::TryStreamExt, |
| 6 | +}; |
| 7 | + |
| 8 | +use super::{options::GridFsUploadOptions, Chunk, FilesCollectionDocument, GridFsBucket}; |
| 9 | +use crate::{ |
| 10 | + bson::{doc, oid::ObjectId, spec::BinarySubtype, Bson, DateTime, Document, RawBinaryRef}, |
| 11 | + bson_util::get_int, |
| 12 | + error::{ErrorKind, Result}, |
| 13 | + index::IndexModel, |
| 14 | + options::{FindOneOptions, ReadPreference, SelectionCriteria}, |
| 15 | + Collection, |
| 16 | +}; |
| 17 | + |
| 18 | +impl GridFsBucket { |
| 19 | + /// Uploads a user file to a GridFS bucket. Bytes are read from `source` and stored in chunks in |
| 20 | + /// the bucket's chunks collection. After all the chunks have been uploaded, a corresponding |
| 21 | + /// [`FilesCollectionDocument`] is stored in the bucket's files collection. |
| 22 | + /// |
| 23 | + /// This method generates an [`ObjectId`] for the `files_id` field of the |
| 24 | + /// [`FilesCollectionDocument`] and returns it. |
| 25 | + pub async fn upload_from_futures_0_3_reader<T>( |
| 26 | + &self, |
| 27 | + filename: impl AsRef<str>, |
| 28 | + source: T, |
| 29 | + options: impl Into<Option<GridFsUploadOptions>>, |
| 30 | + ) -> Result<ObjectId> |
| 31 | + where |
| 32 | + T: AsyncRead + Unpin, |
| 33 | + { |
| 34 | + let id = ObjectId::new(); |
| 35 | + self.upload_from_futures_0_3_reader_with_id(id.into(), filename, source, options) |
| 36 | + .await?; |
| 37 | + Ok(id) |
| 38 | + } |
| 39 | + |
| 40 | + /// Uploads a user file to a GridFS bucket with the given `files_id`. Bytes are read from |
| 41 | + /// `source` and stored in chunks in the bucket's chunks collection. After all the chunks have |
| 42 | + /// been uploaded, a corresponding [`FilesCollectionDocument`] is stored in the bucket's files |
| 43 | + /// collection. |
| 44 | + pub async fn upload_from_futures_0_3_reader_with_id<T>( |
| 45 | + &self, |
| 46 | + files_id: Bson, |
| 47 | + filename: impl AsRef<str>, |
| 48 | + mut source: T, |
| 49 | + options: impl Into<Option<GridFsUploadOptions>>, |
| 50 | + ) -> Result<()> |
| 51 | + where |
| 52 | + T: AsyncRead + Unpin, |
| 53 | + { |
| 54 | + let options = options.into(); |
| 55 | + |
| 56 | + self.create_indexes().await?; |
| 57 | + |
| 58 | + let chunk_size = options |
| 59 | + .as_ref() |
| 60 | + .and_then(|opts| opts.chunk_size_bytes) |
| 61 | + .unwrap_or_else(|| self.chunk_size_bytes()); |
| 62 | + let mut length = 0u64; |
| 63 | + let mut n = 0; |
| 64 | + |
| 65 | + let mut buf = vec![0u8; chunk_size as usize]; |
| 66 | + loop { |
| 67 | + let bytes_read = match source.read(&mut buf).await { |
| 68 | + Ok(0) => break, |
| 69 | + Ok(n) => n, |
| 70 | + Err(error) => { |
| 71 | + self.chunks() |
| 72 | + .delete_many(doc! { "files_id": &files_id }, None) |
| 73 | + .await?; |
| 74 | + return Err(ErrorKind::Io(error.into()).into()); |
| 75 | + } |
| 76 | + }; |
| 77 | + |
| 78 | + let chunk = Chunk { |
| 79 | + id: ObjectId::new(), |
| 80 | + files_id: files_id.clone(), |
| 81 | + n, |
| 82 | + data: RawBinaryRef { |
| 83 | + subtype: BinarySubtype::Generic, |
| 84 | + bytes: &buf[..bytes_read], |
| 85 | + }, |
| 86 | + }; |
| 87 | + self.chunks().insert_one(chunk, None).await?; |
| 88 | + |
| 89 | + length += bytes_read as u64; |
| 90 | + n += 1; |
| 91 | + } |
| 92 | + |
| 93 | + let file = FilesCollectionDocument { |
| 94 | + id: files_id, |
| 95 | + length, |
| 96 | + chunk_size, |
| 97 | + upload_date: DateTime::now(), |
| 98 | + filename: Some(filename.as_ref().to_string()), |
| 99 | + metadata: options.and_then(|opts| opts.metadata), |
| 100 | + }; |
| 101 | + self.files().insert_one(file, None).await?; |
| 102 | + |
| 103 | + Ok(()) |
| 104 | + } |
| 105 | + |
| 106 | + async fn create_indexes(&self) -> Result<()> { |
| 107 | + if !self.inner.created_indexes.load(Ordering::SeqCst) { |
| 108 | + let find_options = FindOneOptions::builder() |
| 109 | + .selection_criteria(SelectionCriteria::ReadPreference(ReadPreference::Primary)) |
| 110 | + .projection(doc! { "_id": 1 }) |
| 111 | + .build(); |
| 112 | + if self |
| 113 | + .files() |
| 114 | + .clone_with_type::<Document>() |
| 115 | + .find_one(None, find_options) |
| 116 | + .await? |
| 117 | + .is_none() |
| 118 | + { |
| 119 | + Self::create_index(self.files(), doc! { "filename": 1, "uploadDate": 1 }).await?; |
| 120 | + Self::create_index(self.chunks(), doc! { "files_id": 1, "n": 1 }).await?; |
| 121 | + } |
| 122 | + self.inner.created_indexes.store(true, Ordering::SeqCst); |
| 123 | + } |
| 124 | + |
| 125 | + Ok(()) |
| 126 | + } |
| 127 | + |
| 128 | + async fn create_index<T>(coll: &Collection<T>, keys: Document) -> Result<()> { |
| 129 | + // From the spec: Drivers MUST check whether the indexes already exist before attempting to |
| 130 | + // create them. |
| 131 | + let mut indexes = coll.list_indexes(None).await?; |
| 132 | + 'outer: while let Some(index_model) = indexes.try_next().await? { |
| 133 | + if index_model.keys.len() != keys.len() { |
| 134 | + continue; |
| 135 | + } |
| 136 | + // Indexes should be considered equivalent regardless of numeric value type. |
| 137 | + // e.g. { "filename": 1, "uploadDate": 1 } is equivalent to |
| 138 | + // { "filename": 1.0, "uploadDate": 1.0 } |
| 139 | + let number_matches = |key: &str, value: &Bson| { |
| 140 | + if let Some(model_value) = index_model.keys.get(key) { |
| 141 | + match get_int(value) { |
| 142 | + Some(num) => get_int(model_value) == Some(num), |
| 143 | + None => model_value == value, |
| 144 | + } |
| 145 | + } else { |
| 146 | + false |
| 147 | + } |
| 148 | + }; |
| 149 | + for (key, value) in keys.iter() { |
| 150 | + if !number_matches(key, value) { |
| 151 | + continue 'outer; |
| 152 | + } |
| 153 | + } |
| 154 | + return Ok(()); |
| 155 | + } |
| 156 | + |
| 157 | + let index_model = IndexModel::builder().keys(keys).build(); |
| 158 | + coll.create_index(index_model, None).await?; |
| 159 | + |
| 160 | + Ok(()) |
| 161 | + } |
| 162 | +} |
0 commit comments