objectstore_client/
get.rs

1use std::{fmt, io};
2
3use async_compression::tokio::bufread::ZstdDecoder;
4use bytes::BytesMut;
5use futures_util::{StreamExt, TryStreamExt};
6use objectstore_types::metadata::{Compression, Metadata};
7use reqwest::StatusCode;
8use tokio_util::io::{ReaderStream, StreamReader};
9
10use crate::response::ResponseExt as _;
11use crate::{ClientStream, ObjectKey, Session};
12
13/// The result from a successful [`get()`](Session::get) call.
14///
15/// This carries the response as a stream, plus the compression algorithm of the data.
16pub struct GetResponse {
17    /// The metadata attached to this object, including the compression algorithm used for the payload.
18    pub metadata: Metadata,
19    /// The response stream.
20    pub stream: ClientStream,
21}
22
23impl GetResponse {
24    /// Loads the object payload fully into memory.
25    pub async fn payload(self) -> crate::Result<bytes::Bytes> {
26        let bytes: BytesMut = self.stream.try_collect().await?;
27        Ok(bytes.freeze())
28    }
29
30    /// Loads the object payload fully into memory and interprets it as UTF-8 text.
31    pub async fn text(self) -> crate::Result<String> {
32        let bytes = self.payload().await?;
33        Ok(String::from_utf8(bytes.to_vec())?)
34    }
35}
36
37impl fmt::Debug for GetResponse {
38    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
39        f.debug_struct("GetResponse")
40            .field("metadata", &self.metadata)
41            .field("stream", &format_args!("[Stream]"))
42            .finish()
43    }
44}
45
46impl Session {
47    /// Retrieves the object with the given `key`.
48    pub fn get(&self, key: &str) -> GetBuilder {
49        GetBuilder {
50            session: self.clone(),
51            key: key.to_owned(),
52            decompress: true,
53            accept_encoding: vec![],
54        }
55    }
56}
57
58/// A [`get`](Session::get) request builder.
59#[derive(Debug)]
60pub struct GetBuilder {
61    pub(crate) session: Session,
62    pub(crate) key: ObjectKey,
63    pub(crate) decompress: bool,
64    pub(crate) accept_encoding: Vec<Compression>,
65}
66
67impl GetBuilder {
68    /// Indicates whether the request should automatically handle decompression of known algorithms,
69    /// or rather return the payload as it is stored, along with the compression algorithm it is stored in.
70    ///
71    /// By default, automatic decompression is enabled.
72    pub fn decompress(mut self, decompress: bool) -> Self {
73        self.decompress = decompress;
74        self
75    }
76
77    /// Specifies compression encodings the caller can handle natively.
78    ///
79    /// When the stored object's compression matches one of these, the payload
80    /// is returned still compressed and `metadata.compression` is preserved.
81    /// An empty list (the default) means the client does not accept any
82    /// compressed encoding, so automatic decompression applies as usual.
83    pub fn accept_encoding(mut self, encodings: impl IntoIterator<Item = Compression>) -> Self {
84        self.accept_encoding = encodings.into_iter().collect();
85        self
86    }
87
88    /// Sends the get request.
89    pub async fn send(self) -> crate::Result<Option<GetResponse>> {
90        let response = self
91            .session
92            .request(reqwest::Method::GET, &self.key)?
93            .send()
94            .await?;
95        if response.status() == StatusCode::NOT_FOUND {
96            response.drain_body().await;
97            return Ok(None);
98        }
99        let response = response.error_for_status_and_drain().await?;
100
101        let mut metadata = match Metadata::from_headers(response.headers(), "") {
102            Ok(metadata) => metadata,
103            Err(error) => {
104                response.drain_body().await;
105                return Err(error.into());
106            }
107        };
108
109        let stream = response.bytes_stream().map_err(io::Error::other).boxed();
110        let stream = maybe_decompress(
111            stream,
112            &mut metadata,
113            self.decompress,
114            &self.accept_encoding,
115        );
116
117        Ok(Some(GetResponse { metadata, stream }))
118    }
119}
120
121/// Wraps a stream in a zstd decompression layer.
122///
123/// Decompresses if the metadata indicates zstd compression, `decompress` is `true`,
124/// and the stored encoding is not listed in `accept_encoding`. When the stored encoding
125/// is in `accept_encoding`, the payload is returned compressed and `metadata.compression`
126/// is preserved. Clears `metadata.compression` when decompression is applied.
127pub(crate) fn maybe_decompress(
128    stream: ClientStream,
129    metadata: &mut Metadata,
130    decompress: bool,
131    accept_encoding: &[Compression],
132) -> ClientStream {
133    let encoding_accepted = metadata
134        .compression
135        .is_some_and(|c| accept_encoding.contains(&c));
136    match (metadata.compression, decompress && !encoding_accepted) {
137        (Some(Compression::Zstd), true) => {
138            metadata.compression = None;
139            let mut decoder = ZstdDecoder::new(StreamReader::new(stream));
140            // Multipart uploads with compression, when each part is compressed individually,
141            // will consist of multiple concatenated zstd frames.
142            // This allows the client to handle automatic decompression for these objects transparently.
143            decoder.multiple_members(true);
144            ReaderStream::new(decoder).boxed()
145        }
146        _ => stream,
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use futures_util::{StreamExt as _, TryStreamExt as _};
153    use objectstore_types::metadata::{Compression, Metadata};
154
155    use super::maybe_decompress;
156    use crate::ClientStream;
157
158    fn compressed_zstd_stream(data: &[u8]) -> ClientStream {
159        let mut encoder = zstd::Encoder::new(vec![], 0).unwrap();
160        std::io::copy(&mut std::io::Cursor::new(data), &mut encoder).unwrap();
161        let compressed = encoder.finish().unwrap();
162        futures_util::stream::once(async move {
163            Ok::<_, std::io::Error>(bytes::Bytes::from(compressed))
164        })
165        .boxed()
166    }
167
168    fn raw_stream(data: &[u8]) -> ClientStream {
169        let bytes = bytes::Bytes::copy_from_slice(data);
170        futures_util::stream::once(async move { Ok::<_, std::io::Error>(bytes) }).boxed()
171    }
172
173    async fn collect(stream: ClientStream) -> Vec<u8> {
174        let chunks: bytes::BytesMut = stream.try_collect().await.unwrap();
175        chunks.to_vec()
176    }
177
178    fn zstd_metadata() -> Metadata {
179        Metadata {
180            compression: Some(Compression::Zstd),
181            ..Default::default()
182        }
183    }
184
185    fn no_compression_metadata() -> Metadata {
186        Metadata::default()
187    }
188
189    #[tokio::test]
190    async fn empty_accept_decompress_true_decompresses() {
191        let payload = b"hello world";
192        let stream = compressed_zstd_stream(payload);
193        let mut metadata = zstd_metadata();
194
195        let out = maybe_decompress(stream, &mut metadata, true, &[]);
196        assert_eq!(collect(out).await, payload);
197        assert_eq!(metadata.compression, None);
198    }
199
200    #[tokio::test]
201    async fn empty_accept_decompress_false_returns_compressed() {
202        let payload = b"hello world";
203        let compressed_bytes = collect(compressed_zstd_stream(payload)).await;
204        let stream = compressed_zstd_stream(payload);
205
206        let mut metadata = zstd_metadata();
207        let out = maybe_decompress(stream, &mut metadata, false, &[]);
208        assert_eq!(collect(out).await, compressed_bytes);
209        assert_eq!(metadata.compression, Some(Compression::Zstd));
210    }
211
212    #[tokio::test]
213    async fn zstd_accept_decompress_true_skips_decompression() {
214        let payload = b"hello world";
215        let compressed_bytes = collect(compressed_zstd_stream(payload)).await;
216        let stream = compressed_zstd_stream(payload);
217
218        let mut metadata = zstd_metadata();
219        let out = maybe_decompress(stream, &mut metadata, true, &[Compression::Zstd]);
220        assert_eq!(collect(out).await, compressed_bytes);
221        assert_eq!(metadata.compression, Some(Compression::Zstd));
222    }
223
224    #[tokio::test]
225    async fn zstd_accept_decompress_false_returns_compressed() {
226        let payload = b"hello world";
227        let compressed_bytes = collect(compressed_zstd_stream(payload)).await;
228        let stream = compressed_zstd_stream(payload);
229
230        let mut metadata = zstd_metadata();
231        let out = maybe_decompress(stream, &mut metadata, false, &[Compression::Zstd]);
232        assert_eq!(collect(out).await, compressed_bytes);
233        assert_eq!(metadata.compression, Some(Compression::Zstd));
234    }
235
236    #[tokio::test]
237    async fn no_compression_returns_raw_regardless_of_accept() {
238        let payload = b"hello world";
239        let stream = raw_stream(payload);
240
241        let mut metadata = no_compression_metadata();
242        let out = maybe_decompress(stream, &mut metadata, true, &[Compression::Zstd]);
243        assert_eq!(collect(out).await, payload);
244        assert_eq!(metadata.compression, None);
245    }
246
247    #[tokio::test]
248    async fn zstd_concatenated_frames_decompress() {
249        let payload1 = b"hello ";
250        let payload2 = b"world";
251        let compressed1 = collect(compressed_zstd_stream(payload1)).await;
252        let compressed2 = collect(compressed_zstd_stream(payload2)).await;
253        let stream = futures_util::stream::iter([
254            Ok::<_, std::io::Error>(bytes::Bytes::from(compressed1)),
255            Ok::<_, std::io::Error>(bytes::Bytes::from(compressed2)),
256        ])
257        .boxed();
258
259        let mut metadata = zstd_metadata();
260        let out = maybe_decompress(stream, &mut metadata, true, &[]);
261        assert_eq!(collect(out).await, b"hello world");
262        assert_eq!(metadata.compression, None);
263    }
264}
objectstore_client/get.rs

objectstore_client/
get.rs
