Skip to main content

s3s/dto/
streaming_blob.rs

1//! Streaming blob
2
3use crate::error::StdError;
4use crate::http::Body;
5use crate::stream::*;
6
7use std::fmt;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10
11use futures::Stream;
12use hyper::body::Bytes;
13
14pub struct StreamingBlob {
15    inner: DynByteStream,
16}
17
18impl StreamingBlob {
19    pub fn new<S>(stream: S) -> Self
20    where
21        S: ByteStream<Item = Result<Bytes, StdError>> + Send + Sync + 'static,
22    {
23        Self { inner: Box::pin(stream) }
24    }
25
26    pub fn wrap<S, E>(stream: S) -> Self
27    where
28        S: Stream<Item = Result<Bytes, E>> + Send + Sync + 'static,
29        E: std::error::Error + Send + Sync + 'static,
30    {
31        Self { inner: wrap(stream) }
32    }
33
34    fn into_inner(self) -> DynByteStream {
35        self.inner
36    }
37}
38
39impl fmt::Debug for StreamingBlob {
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        f.debug_struct("StreamingBlob")
42            .field("remaining_length", &self.remaining_length())
43            .finish_non_exhaustive()
44    }
45}
46
47impl Stream for StreamingBlob {
48    type Item = Result<Bytes, StdError>;
49
50    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
51        Pin::new(&mut self.inner).poll_next(cx)
52    }
53
54    fn size_hint(&self) -> (usize, Option<usize>) {
55        self.inner.size_hint()
56    }
57}
58
59impl ByteStream for StreamingBlob {
60    fn remaining_length(&self) -> RemainingLength {
61        self.inner.remaining_length()
62    }
63}
64
65impl From<StreamingBlob> for DynByteStream {
66    fn from(value: StreamingBlob) -> Self {
67        value.into_inner()
68    }
69}
70
71impl From<DynByteStream> for StreamingBlob {
72    fn from(value: DynByteStream) -> Self {
73        Self { inner: value }
74    }
75}
76
77impl From<StreamingBlob> for Body {
78    fn from(value: StreamingBlob) -> Self {
79        Body::from(value.into_inner())
80    }
81}
82
83impl From<Body> for StreamingBlob {
84    fn from(value: Body) -> Self {
85        Self::new(value)
86    }
87}
88
89pin_project_lite::pin_project! {
90    pub(crate) struct StreamWrapper<S> {
91        #[pin]
92        inner: S
93    }
94}
95
96impl<S, E> Stream for StreamWrapper<S>
97where
98    S: Stream<Item = Result<Bytes, E>> + Send + Sync + 'static,
99    E: std::error::Error + Send + Sync + 'static,
100{
101    type Item = Result<Bytes, StdError>;
102
103    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
104        let this = self.project();
105        this.inner.poll_next(cx).map_err(|e| Box::new(e) as StdError)
106    }
107
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        self.inner.size_hint()
110    }
111}
112
113impl<S> ByteStream for StreamWrapper<S>
114where
115    StreamWrapper<S>: Stream<Item = Result<Bytes, StdError>>,
116{
117    fn remaining_length(&self) -> RemainingLength {
118        RemainingLength::unknown()
119    }
120}
121
122fn wrap<S>(inner: S) -> DynByteStream
123where
124    StreamWrapper<S>: ByteStream<Item = Result<Bytes, StdError>> + Send + Sync + 'static,
125{
126    Box::pin(StreamWrapper { inner })
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use futures::StreamExt;
133    use http_body::Body as HttpBody;
134
135    #[tokio::test]
136    async fn streaming_blob_new_and_poll() {
137        let body = Body::from(Bytes::from_static(b"hello world"));
138        let mut blob = StreamingBlob::new(body);
139        let mut collected = Vec::new();
140        while let Some(chunk) = blob.next().await {
141            collected.push(chunk.unwrap());
142        }
143        assert_eq!(collected, vec![Bytes::from_static(b"hello world")]);
144    }
145
146    #[tokio::test]
147    async fn streaming_blob_wrap() {
148        let data = vec![
149            Ok::<_, std::io::Error>(Bytes::from_static(b"abc")),
150            Ok(Bytes::from_static(b"def")),
151        ];
152        let stream = futures::stream::iter(data);
153        let mut blob = StreamingBlob::wrap(stream);
154        let mut collected = Vec::new();
155        while let Some(chunk) = blob.next().await {
156            collected.push(chunk.unwrap());
157        }
158        assert_eq!(collected, vec![Bytes::from_static(b"abc"), Bytes::from_static(b"def")]);
159    }
160
161    #[test]
162    fn streaming_blob_debug() {
163        let body = Body::from(Bytes::from_static(b"test"));
164        let blob = StreamingBlob::new(body);
165        let debug = format!("{blob:?}");
166        assert!(debug.contains("StreamingBlob"));
167        assert!(debug.contains("remaining_length"));
168    }
169
170    #[test]
171    fn streaming_blob_remaining_length() {
172        let body = Body::from(Bytes::from_static(b"hello"));
173        let blob = StreamingBlob::new(body);
174        let rl = blob.remaining_length();
175        assert_eq!(rl.exact(), Some(5));
176    }
177
178    #[test]
179    fn streaming_blob_from_body_roundtrip() {
180        let body = Body::from(Bytes::from_static(b"data"));
181        let blob = StreamingBlob::from(body);
182        let body_back: Body = Body::from(blob);
183        assert!(!HttpBody::is_end_stream(&body_back));
184    }
185
186    #[test]
187    fn streaming_blob_into_dyn_byte_stream() {
188        let body = Body::from(Bytes::from_static(b"test"));
189        let blob = StreamingBlob::new(body);
190        let _dyn_stream: DynByteStream = blob.into();
191    }
192
193    #[test]
194    fn streaming_blob_from_dyn_byte_stream() {
195        let body = Body::from(Bytes::from_static(b"test"));
196        let dyn_stream: DynByteStream = Box::pin(body);
197        let _blob = StreamingBlob::from(dyn_stream);
198    }
199
200    #[test]
201    fn streaming_blob_size_hint() {
202        let body = Body::from(Bytes::from_static(b"12345"));
203        let blob = StreamingBlob::new(body);
204        let (lower, upper) = blob.size_hint();
205        if let Some(upper) = upper {
206            assert!(lower <= upper);
207        }
208    }
209
210    #[tokio::test]
211    async fn streaming_blob_empty() {
212        let body = Body::empty();
213        let mut blob = StreamingBlob::new(body);
214        let next = blob.next().await;
215        assert!(next.is_none());
216    }
217}