[incubator-opendal] 04/05: Implement sink

[xu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xuanwo pushed a commit to branch implement-sink-for-s3
in repository https://gitbox.apache.org/repos/asf/incubator-opendal.git

commit 13e85c9c00a5c35be93397c252fa4c93b6cffc5d
Author: Xuanwo <gi...@xuanwo.io>
AuthorDate: Wed Jun 21 18:46:04 2023 +0800

    Implement sink
    
    Signed-off-by: Xuanwo <gi...@xuanwo.io>
---
 core/src/types/writer.rs     | 23 +++++++++++++++++++++++
 core/tests/behavior/write.rs | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/core/src/types/writer.rs b/core/src/types/writer.rs
index 049116566..98eebdac0 100644
--- a/core/src/types/writer.rs
+++ b/core/src/types/writer.rs
@@ -26,6 +26,7 @@ use bytes::Bytes;
 use futures::future::BoxFuture;
 use futures::AsyncWrite;
 use futures::FutureExt;
+use futures::TryStreamExt;
 
 use crate::raw::oio::Write;
 use crate::raw::*;
@@ -90,6 +91,28 @@ impl Writer {
         }
     }
 
+    /// Sink into writer.
+    ///
+    /// sink will read data from given streamer and write them into writer
+    /// directly without extra in-memory buffer.
+    pub async fn sink<S, T>(&mut self, size: u64, s: S) -> Result<()>
+    where
+        S: futures::Stream<Item = Result<T>> + Send + Sync + Unpin + 'static,
+        T: Into<Bytes>,
+    {
+        if let State::Idle(Some(w)) = &mut self.state {
+            let s = Box::new(oio::into_stream::from_futures_stream(
+                s.map_ok(|v| v.into()),
+            ));
+            w.sink(size, s).await
+        } else {
+            unreachable!(
+                "writer state invalid while pipe_form, expect Idle, actual {}",
+                self.state
+            );
+        }
+    }
+
     /// Abort the writer and clean up all written data.
     ///
     /// ## Notes
diff --git a/core/tests/behavior/write.rs b/core/tests/behavior/write.rs
index cff78532b..f3b935a25 100644
--- a/core/tests/behavior/write.rs
+++ b/core/tests/behavior/write.rs
@@ -21,6 +21,7 @@ use std::time::Duration;
 use anyhow::Result;
 use futures::io::BufReader;
 use futures::io::Cursor;
+use futures::stream;
 use futures::AsyncReadExt;
 use futures::AsyncSeekExt;
 use futures::StreamExt;
@@ -111,6 +112,7 @@ macro_rules! behavior_write_tests {
                 test_delete_stream,
                 test_remove_one_file,
                 test_writer_write,
+                test_writer_sink,
                 test_writer_abort,
                 test_writer_futures_copy,
                 test_fuzz_unsized_writer,
@@ -1121,6 +1123,41 @@ pub async fn test_writer_write(op: Operator) -> Result<()> {
     Ok(())
 }
 
+/// Streaming data into writer
+pub async fn test_writer_sink(op: Operator) -> Result<()> {
+    let path = uuid::Uuid::new_v4().to_string();
+    let size = 5 * 1024 * 1024; // write file with 5 MiB
+    let content_a = gen_fixed_bytes(size);
+    let content_b = gen_fixed_bytes(size);
+    let stream = stream::iter(vec![content_a.clone(), content_b.clone()]).map(Ok);
+
+    let mut w = op
+        .writer_with(&path)
+        .content_length(2 * size as u64)
+        .await?;
+    w.sink(2 * size as u64, stream).await?;
+    w.close().await?;
+
+    let meta = op.stat(&path).await.expect("stat must succeed");
+    assert_eq!(meta.content_length(), (size * 2) as u64);
+
+    let bs = op.read(&path).await?;
+    assert_eq!(bs.len(), size * 2, "read size");
+    assert_eq!(
+        format!("{:x}", Sha256::digest(&bs[..size])),
+        format!("{:x}", Sha256::digest(content_a)),
+        "read content a"
+    );
+    assert_eq!(
+        format!("{:x}", Sha256::digest(&bs[size..])),
+        format!("{:x}", Sha256::digest(content_b)),
+        "read content b"
+    );
+
+    op.delete(&path).await.expect("delete must succeed");
+    Ok(())
+}
+
 /// Copy data from reader to writer
 pub async fn test_writer_futures_copy(op: Operator) -> Result<()> {
     let path = uuid::Uuid::new_v4().to_string();