1use bytes::Bytes;
10use futures::{Stream, StreamExt, TryStreamExt};
11use std::pin::Pin;
12
13use super::error::FetchError;
14
15pub type ByteStream = Pin<Box<dyn Stream<Item = Result<Bytes, FetchError>> + Send>>;
17
18#[must_use]
30pub fn channel_stream(rx: tokio::sync::mpsc::Receiver<Result<Vec<u8>, String>>) -> ByteStream {
31 futures::stream::unfold(rx, |mut rx| async move {
32 let chunk = rx.recv().await?;
33 let item = match chunk {
34 Ok(bytes) => Ok(Bytes::from(bytes)),
35 Err(message) => Err(FetchError::Body(message)),
36 };
37 Some((item, rx))
38 })
39 .boxed()
40}
41
42pub(crate) enum RequestPayload {
48 Empty,
49 Bytes(Bytes),
50 Stream(ByteStream),
51 Invalid,
53}
54
55pub struct Body(Inner);
57
58enum Inner {
59 Empty,
60 Bytes(Bytes),
61 Stream(ByteStream),
62 Response(Box<reqwest::Response>),
65}
66
67impl Body {
68 #[must_use]
69 pub fn empty() -> Self {
70 Self(Inner::Empty)
71 }
72
73 #[must_use]
74 pub fn from_bytes(bytes: impl Into<Bytes>) -> Self {
75 Self(Inner::Bytes(bytes.into()))
76 }
77
78 #[must_use]
79 pub fn from_stream(stream: ByteStream) -> Self {
80 Self(Inner::Stream(stream))
81 }
82
83 pub(crate) fn from_response(response: reqwest::Response) -> Self {
84 Self(Inner::Response(Box::new(response)))
85 }
86
87 pub(crate) fn into_request_payload(self) -> RequestPayload {
90 match self.0 {
91 Inner::Empty => RequestPayload::Empty,
92 Inner::Bytes(b) => RequestPayload::Bytes(b),
93 Inner::Stream(s) => RequestPayload::Stream(s),
94 Inner::Response(_) => RequestPayload::Invalid,
98 }
99 }
100
101 #[must_use]
103 pub fn is_empty(&self) -> bool {
104 matches!(self.0, Inner::Empty)
105 }
106
107 pub async fn collect(self) -> Result<Bytes, FetchError> {
114 match self.0 {
115 Inner::Empty => Ok(Bytes::new()),
116 Inner::Bytes(b) => Ok(b),
117 Inner::Response(resp) => resp
118 .bytes()
119 .await
120 .map_err(|e| FetchError::Body(format!("read response body: {e}"))),
121 Inner::Stream(mut s) => {
122 let mut buf = Vec::new();
123 while let Some(chunk) = s.next().await {
124 buf.extend_from_slice(&chunk?);
125 }
126 Ok(Bytes::from(buf))
127 },
128 }
129 }
130
131 #[must_use]
134 pub fn into_stream(self) -> ByteStream {
135 match self.0 {
136 Inner::Empty => futures::stream::empty().boxed(),
137 Inner::Bytes(b) => futures::stream::once(async move { Ok(b) }).boxed(),
138 Inner::Stream(s) => s,
139 Inner::Response(resp) => resp
140 .bytes_stream()
141 .map_err(|e| FetchError::Body(format!("read response body: {e}")))
142 .boxed(),
143 }
144 }
145}
146
147impl std::fmt::Debug for Body {
148 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149 match &self.0 {
150 Inner::Empty => f.write_str("Body::Empty"),
151 Inner::Bytes(b) => write!(f, "Body::Bytes({} bytes)", b.len()),
152 Inner::Stream(_) => f.write_str("Body::Stream"),
153 Inner::Response(_) => f.write_str("Body::Response"),
154 }
155 }
156}