Add file read/write ability to ObjectContent

Cargo.toml

@@ -20,16 +20,20 @@ chrono = "0.4.27"
 crc = "3.0.1"
 dashmap = "5.5.3"
 derivative = "2.2.0"
+env_logger = "0.11.2"
 futures-util = "0.3.28"
 hex = "0.4.3"
 hmac = "0.12.1"
+home = "0.5.9"
 http = "0.2.9"
 hyper = { version = "0.14.27", features = ["full"] }
 lazy_static = "1.4.0"
+log = "0.4.20"
 md5 = "0.7.0"
 multimap = "0.10.0"
 os_info = "3.7.0"
 percent-encoding = "2.3.0"
+rand = { version = "0.8.5", features = ["small_rng"] }
 regex = "1.9.4"
 serde = { version = "1.0.188", features = ["derive"] }
 serde_json = "1.0.105"
@@ -39,17 +43,13 @@ tokio-stream = "0.1.14"
 tokio-util = { version = "0.7.8", features = ["io"] }
 urlencoding = "2.1.3"
 xmltree = "0.10.3"
-log = "0.4.20"
-env_logger = "0.11.2"
-home = "0.5.9"
 
 [dependencies.reqwest]
 version = "0.11.20"
 features = ["native-tls", "blocking", "rustls-tls", "stream"]
 
 [dev-dependencies]
 async-std = { version = "1.12.0", features = ["attributes", "tokio1"] }
-rand = { version = "0.8.5", features = ["small_rng"] }
 
 [[example]]
 name = "file-uploader"

examples/file-uploader.rs

@@ -1,5 +1,6 @@
 use log::info;
 use minio::s3::args::{BucketExistsArgs, MakeBucketArgs};
+use minio::s3::builders::ObjectContent;
 use minio::s3::client::ClientBuilder;
 use minio::s3::creds::StaticProvider;
 use minio::s3::http::BaseUrl;
@@ -47,8 +48,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     info!("filename {}", &filename.to_str().unwrap());
 
+    let content = ObjectContent::from(filename);
     client
-        .put_object_from_file(bucket_name, object_name, filename)
+        .put_object_content(bucket_name, object_name, content)
         .send()
         .await?;
 

src/s3/builders/object_content.rs

@@ -13,55 +13,164 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::pin::Pin;
+use std::path::PathBuf;
+use std::{ffi::OsString, path::Path, pin::Pin};
 
 use bytes::{Bytes, BytesMut};
 use futures_util::Stream;
-use tokio::io::AsyncRead;
+use rand::prelude::random;
+use tokio::fs;
+use tokio::io::AsyncWriteExt;
 use tokio_stream::StreamExt;
 
 type IoResult<T> = Result<T, std::io::Error>;
 
-pub struct ObjectContent {
-    r: Pin<Box<dyn Stream<Item = IoResult<Bytes>>>>,
-    extra: Option<Bytes>,
-    size: Option<u64>,
+/// Object content that can be uploaded or downloaded. Can be constructed from a stream of `Bytes`,
+/// a file path, or a `Bytes` object.
+pub struct ObjectContent(ObjectContentInner);
+
+enum ObjectContentInner {
+    Stream(Pin<Box<dyn Stream<Item = IoResult<Bytes>>>>, Option<u64>),
+    FilePath(PathBuf),
+    Bytes(SegmentedBytes),
 }
 
 impl From<Bytes> for ObjectContent {
     fn from(value: Bytes) -> Self {
-        let n = value.len();
-        ObjectContent {
-            r: Box::pin(tokio_stream::iter(vec![Ok(value)])),
-            extra: None,
-            size: Some(n as u64),
-        }
+        ObjectContent(ObjectContentInner::Bytes(SegmentedBytes::from(value)))
     }
 }
 
 impl From<String> for ObjectContent {
     fn from(value: String) -> Self {
-        let n = value.len();
-        ObjectContent {
-            r: Box::pin(tokio_stream::iter(vec![Ok(Bytes::from(value))])),
-            extra: None,
-            size: Some(n as u64),
-        }
+        ObjectContent(ObjectContentInner::Bytes(SegmentedBytes::from(
+            Bytes::from(value),
+        )))
     }
 }
 
 impl From<Vec<u8>> for ObjectContent {
     fn from(value: Vec<u8>) -> Self {
-        let n = value.len();
-        ObjectContent {
-            r: Box::pin(tokio_stream::iter(vec![Ok(Bytes::from(value))])),
-            extra: None,
-            size: Some(n as u64),
-        }
+        ObjectContent(ObjectContentInner::Bytes(SegmentedBytes::from(
+            Bytes::from(value),
+        )))
+    }
+}
+
+impl From<&Path> for ObjectContent {
+    fn from(value: &Path) -> Self {
+        ObjectContent(ObjectContentInner::FilePath(value.to_path_buf()))
+    }
+}
+
+impl Default for ObjectContent {
+    fn default() -> Self {
+        ObjectContent(ObjectContentInner::Bytes(SegmentedBytes::new()))
     }
 }
 
 impl ObjectContent {
+    /// Create a new `ObjectContent` from a stream of `Bytes`.
+    pub fn new_from_stream(
+        r: impl Stream<Item = IoResult<Bytes>> + 'static,
+        size: Option<u64>,
+    ) -> Self {
+        let r = Box::pin(r);
+        ObjectContent(ObjectContentInner::Stream(r, size))
+    }
+
+    pub async fn to_stream(
+        self,
+    ) -> IoResult<(Pin<Box<dyn Stream<Item = IoResult<Bytes>>>>, Option<u64>)> {
+        match self.0 {
+            ObjectContentInner::Stream(r, size) => Ok((r, size)),
+            ObjectContentInner::FilePath(path) => {
+                let file = fs::File::open(&path).await?;
+                let size = file.metadata().await?.len();
+                let r = tokio_util::io::ReaderStream::new(file);
+                Ok((Box::pin(r), Some(size)))
+            }
+            ObjectContentInner::Bytes(sb) => {
+                let k = sb.len();
+                let r = Box::pin(tokio_stream::iter(sb.into_iter().map(Ok)));
+                Ok((r, Some(k as u64)))
+            }
+        }
+    }
+
+    pub(crate) async fn to_content_stream(self) -> IoResult<ContentStream> {
+        let (r, size) = self.to_stream().await?;
+        Ok(ContentStream::new(r, size))
+    }
+
+    /// Load the content into memory and return a `SegmentedBytes` object.
+    pub async fn to_segmented_bytes(self) -> IoResult<SegmentedBytes> {
+        let mut segmented_bytes = SegmentedBytes::new();
+        let (mut r, _) = self.to_stream().await?;
+        while let Some(bytes) = r.next().await {
+            let bytes = bytes?;
+            if bytes.is_empty() {
+                break;
+            }
+            segmented_bytes.append(bytes);
+        }
+        Ok(segmented_bytes)
+    }
+
+    /// Write the content to a file. This function will return the total number
+    /// of bytes written to the file. It first writes the content to a temporary
+    /// file and then renames the temporary file to the final file path. The
+    /// temporary file will be located in the same directory as the final file
+    /// path.
+    ///
+    /// If the file already exists, it will be replaced. If the parent directory
+    /// does not exist, an attempt to create it will be made.
+    pub async fn to_file(self, file_path: &Path) -> IoResult<u64> {
+        if file_path.is_dir() {
+            return Err(std::io::Error::other("path is a directory"));
+        }
+        let parent_dir = file_path.parent().ok_or(std::io::Error::other(format!(
+            "path {:?} does not have a parent directory",
+            file_path
+        )))?;
+        if !parent_dir.is_dir() {
+            fs::create_dir_all(parent_dir).await?;
+        }
+        let file_name = file_path.file_name().ok_or(std::io::Error::other(
+            "could not get filename component of path",
+        ))?;
+        let mut tmp_file_name: OsString = file_name.to_os_string();
+        tmp_file_name.push(format!("_{}", random::<u64>()));
+        let tmp_file_path = parent_dir
+            .to_path_buf()
+            .join(Path::new(tmp_file_name.as_os_str()));
+
+        let mut total = 0;
+        {
+            let mut fp = fs::File::open(&tmp_file_path).await?;
+            let (mut r, _) = self.to_stream().await?;
+            while let Some(bytes) = r.next().await {
+                let bytes = bytes?;
+                if bytes.is_empty() {
+                    break;
+                }
+                total += bytes.len() as u64;
+                fp.write_all(&bytes).await?;
+            }
+            fp.flush().await?;
+        }
+        fs::rename(&tmp_file_path, file_path).await?;
+        Ok(total)
+    }
+}
+
+pub(crate) struct ContentStream {
+    r: Pin<Box<dyn Stream<Item = IoResult<Bytes>>>>,
+    extra: Option<Bytes>,
+    size: Option<u64>,
+}
+
+impl ContentStream {
     pub fn new(r: impl Stream<Item = IoResult<Bytes>> + 'static, size: Option<u64>) -> Self {
         let r = Box::pin(r);
         Self {
@@ -79,24 +188,10 @@ impl ObjectContent {
         }
     }
 
-    pub fn from_reader(r: impl AsyncRead + Send + Sync + 'static, size: Option<u64>) -> Self {
-        let pinned = Box::pin(r);
-        let r = tokio_util::io::ReaderStream::new(pinned);
-        Self {
-            r: Box::pin(r),
-            extra: None,
-            size,
-        }
-    }
-
     pub fn get_size(&self) -> Option<u64> {
         self.size
     }
 
-    pub fn stream(self) -> impl Stream<Item = IoResult<Bytes>> {
-        self.r
-    }
-
     // Read as many bytes as possible up to `n` and return a `SegmentedBytes`
     // object.
     pub async fn read_upto(&mut self, n: usize) -> IoResult<SegmentedBytes> {
@@ -134,20 +229,9 @@ impl ObjectContent {
         }
         Ok(segmented_bytes)
     }
-
-    pub async fn to_segmented_bytes(mut self) -> IoResult<SegmentedBytes> {
-        let mut segmented_bytes = SegmentedBytes::new();
-        while let Some(bytes) = self.r.next().await {
-            let bytes = bytes?;
-            if bytes.is_empty() {
-                break;
-            }
-            segmented_bytes.append(bytes);
-        }
-        Ok(segmented_bytes)
-    }
 }
 
+/// An aggregated collection of `Bytes` objects.
 #[derive(Debug, Clone)]
 pub struct SegmentedBytes {
     segments: Vec<Vec<Bytes>>,
@@ -200,6 +284,14 @@ impl SegmentedBytes {
         }
     }
 
+    pub fn into_iter(self) -> SegmentedBytesIntoIterator {
+        SegmentedBytesIntoIterator {
+            sb: self,
+            current_segment: 0,
+            current_segment_index: 0,
+        }
+    }
+
     // Copy all the content into a single `Bytes` object.
     pub fn to_bytes(&self) -> Bytes {
         let mut buf = BytesMut::with_capacity(self.total_size);
@@ -212,11 +304,42 @@ impl SegmentedBytes {
     }
 }
 
-impl From<Bytes> for SegmentedBytes {
-    fn from(bytes: Bytes) -> Self {
-        let mut sb = SegmentedBytes::new();
-        sb.append(bytes);
-        sb
+pub struct SegmentedBytesIntoIterator {
+    sb: SegmentedBytes,
+    current_segment: usize,
+    current_segment_index: usize,
+}
+
+impl Iterator for SegmentedBytesIntoIterator {
+    type Item = Bytes;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current_segment >= self.sb.segments.len() {
+            return None;
+        }
+        let segment = &self.sb.segments[self.current_segment];
+        if self.current_segment_index >= segment.len() {
+            self.current_segment += 1;
+            self.current_segment_index = 0;
+            return Iterator::next(self);
+        }
+        let bytes = &segment[self.current_segment_index];
+        self.current_segment_index += 1;
+        Some(bytes.clone())
+    }
+}
+
+impl IntoIterator for SegmentedBytes {
+    type Item = Bytes;
+
+    type IntoIter = SegmentedBytesIntoIterator;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SegmentedBytesIntoIterator {
+            sb: self,
+            current_segment: 0,
+            current_segment_index: 0,
+        }
     }
 }
 
@@ -226,7 +349,7 @@ pub struct SegmentedBytesIterator<'a> {
     current_segment_index: usize,
 }
 
-impl Iterator for SegmentedBytesIterator<'_> {
+impl<'a> Iterator for SegmentedBytesIterator<'a> {
     type Item = Bytes;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -237,10 +360,31 @@ impl Iterator for SegmentedBytesIterator<'_> {
         if self.current_segment_index >= segment.len() {
             self.current_segment += 1;
             self.current_segment_index = 0;
-            return self.next();
+            return Iterator::next(self);
         }
         let bytes = &segment[self.current_segment_index];
         self.current_segment_index += 1;
         Some(bytes.clone())
     }
 }
+
+impl<'a> IntoIterator for &'a SegmentedBytes {
+    type Item = Bytes;
+    type IntoIter = SegmentedBytesIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SegmentedBytesIterator {
+            sb: self,
+            current_segment: 0,
+            current_segment_index: 0,
+        }
+    }
+}
+
+impl From<Bytes> for SegmentedBytes {
+    fn from(bytes: Bytes) -> Self {
+        let mut sb = SegmentedBytes::new();
+        sb.append(bytes);
+        sb
+    }
+}