objectstore_client/
put.rs

1use std::fmt;
2use std::io::{self, Cursor};
3use std::path::PathBuf;
4use std::{borrow::Cow, collections::BTreeMap};
5
6use async_compression::tokio::bufread::{GzipEncoder, ZstdEncoder};
7use bytes::Bytes;
8use futures_util::StreamExt;
9use objectstore_types::metadata::Metadata;
10use reqwest::Body;
11use serde::Deserialize;
12use tokio::fs::File;
13use tokio::io::{AsyncBufRead, AsyncRead, BufReader};
14use tokio_util::io::{ReaderStream, StreamReader};
15
16pub use objectstore_types::metadata::{Compression, ExpirationPolicy};
17
18use crate::response::ResponseExt as _;
19use crate::{ClientStream, ObjectKey, Session};
20
21/// The response returned from the service after uploading an object.
22#[derive(Debug, Deserialize)]
23pub struct PutResponse {
24    /// The key of the object, as stored.
25    pub key: ObjectKey,
26}
27
28pub(crate) enum PutBody {
29    Buffer(Bytes),
30    Stream(ClientStream),
31    File(File),
32    Path(PathBuf),
33}
34
35impl fmt::Debug for PutBody {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        f.debug_tuple("PutBody").finish_non_exhaustive()
38    }
39}
40
41/// Declares how a payload relates to the compression recorded in its metadata.
42///
43/// Both modes record the same [`Compression`] on the object; they differ only in who performs
44/// the compression.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub(crate) enum CompressionMode {
47    /// The client compresses the payload with this algorithm before uploading it.
48    Compress(Compression),
49    /// The payload is already compressed with this algorithm and is uploaded verbatim.
50    Precompressed(Compression),
51}
52
53impl CompressionMode {
54    /// Returns the compression algorithm applied to the payload.
55    pub fn compression(self) -> Compression {
56        match self {
57            Self::Compress(compression) | Self::Precompressed(compression) => compression,
58        }
59    }
60}
61
62impl Session {
63    fn put_body(&self, body: PutBody) -> PutBuilder {
64        let metadata = Metadata {
65            expiration_policy: self.scope.usecase().expiration_policy(),
66            ..Default::default()
67        };
68
69        PutBuilder {
70            session: self.clone(),
71            metadata,
72            compression: self
73                .scope
74                .usecase()
75                .compression()
76                .map(CompressionMode::Compress),
77            key: None,
78            body,
79        }
80    }
81
82    /// Creates or replaces an object using a [`Bytes`]-like payload.
83    pub fn put(&self, body: impl Into<Bytes>) -> PutBuilder {
84        self.put_body(PutBody::Buffer(body.into()))
85    }
86
87    /// Creates or replaces an object using a streaming payload.
88    pub fn put_stream(&self, body: ClientStream) -> PutBuilder {
89        self.put_body(PutBody::Stream(body))
90    }
91
92    /// Creates or replaces an object using an [`AsyncRead`] payload.
93    pub fn put_read<R>(&self, body: R) -> PutBuilder
94    where
95        R: AsyncRead + Send + Sync + 'static,
96    {
97        let stream = ReaderStream::new(body).boxed();
98        self.put_body(PutBody::Stream(stream))
99    }
100
101    /// Creates or replaces an object using the contents of an opened file.
102    ///
103    /// The file descriptor is held open from the moment this method is called until the
104    /// upload completes. When enqueueing many files via [`Session::many`], prefer
105    /// [`put_path`](Session::put_path) instead: it defers opening the file until just before
106    /// upload, keeping file descriptor usage within the active concurrency window and avoiding
107    /// OS file descriptor limit (e.g., macOS's default `ulimit -n`) exhaustion.
108    pub fn put_file(&self, file: File) -> PutBuilder {
109        self.put_body(PutBody::File(file))
110    }
111
112    /// Creates or replaces an object using the contents of the file at `path`.
113    ///
114    /// Unlike [`put_file`](Session::put_file), this method defers opening the file until the
115    /// request is actually sent. When enqueueing many file uploads via [`Session::many`], this
116    /// ensures that file descriptors are opened only within the active concurrency window,
117    /// preventing the process from exhausting the OS file descriptor limit (e.g., macOS's
118    /// default `ulimit -n`).
119    ///
120    /// Prefer `put_path` over [`put_file`](Session::put_file) whenever you are lining up a
121    /// large number of files for upload.
122    pub fn put_path(&self, path: impl Into<PathBuf>) -> PutBuilder {
123        self.put_body(PutBody::Path(path.into()))
124    }
125}
126
127/// A [`put`](Session::put) request builder.
128#[derive(Debug)]
129pub struct PutBuilder {
130    pub(crate) session: Session,
131    pub(crate) metadata: Metadata,
132    pub(crate) compression: Option<CompressionMode>,
133    pub(crate) key: Option<ObjectKey>,
134    pub(crate) body: PutBody,
135}
136
137impl PutBuilder {
138    /// Sets an explicit object key.
139    ///
140    /// If a key is specified, the object will be stored under that key. Otherwise, the Objectstore
141    /// server will automatically assign a random key, which is then returned from this request.
142    pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
143        self.key = Some(key.into()).filter(|k| !k.is_empty());
144        self
145    }
146
147    /// Sets an explicit compression algorithm to be used for this payload.
148    ///
149    /// The client compresses the payload while uploading it and records the algorithm in the
150    /// object's metadata. [`None`] should be used if no compression should be performed by the
151    /// client, either because the payload is uncompressible (such as a media format), or if the
152    /// compression should not be recorded for this object.
153    ///
154    /// If the payload is already compressed and the algorithm should still be recorded, use
155    /// [`precompressed`](Self::precompressed) instead.
156    ///
157    /// By default, the compression algorithm set on this Session's Usecase is used (see
158    /// [`with_compression`](crate::Usecase::with_compression)).
159    ///
160    /// # Example
161    ///
162    /// ```no_run
163    /// # async fn example(session: objectstore_client::Session, media: Vec<u8>) {
164    /// session.put(media)
165    ///     .compress(None) // uncompressible payload
166    ///     .send()
167    ///     .await
168    ///     .unwrap();
169    /// # }
170    /// ```
171    pub fn compress(mut self, compression: impl Into<Option<Compression>>) -> Self {
172        self.compression = compression.into().map(CompressionMode::Compress);
173        self
174    }
175
176    /// Declares that the payload is already compressed with the given algorithm.
177    ///
178    /// The payload is uploaded verbatim, and the algorithm is recorded in the object's metadata
179    /// so that downloads decompress it transparently. Use this to hand pre-compressed data to
180    /// the client without paying for another compression pass.
181    ///
182    /// This overrides the compression algorithm set on this Session's Usecase. To have the
183    /// client perform the compression instead, use [`compress`](Self::compress).
184    ///
185    /// # Example
186    ///
187    /// ```no_run
188    /// # use objectstore_client::Compression;
189    /// # async fn example(session: objectstore_client::Session, zstd_data: Vec<u8>) {
190    /// session.put(zstd_data)
191    ///     .precompressed(Compression::Zstd)
192    ///     .send()
193    ///     .await
194    ///     .unwrap();
195    /// # }
196    /// ```
197    pub fn precompressed(mut self, compression: Compression) -> Self {
198        self.compression = Some(CompressionMode::Precompressed(compression));
199        self
200    }
201
202    /// Sets the expiration policy of the object to be uploaded.
203    ///
204    /// By default, the expiration policy set on this Session's Usecase is used.
205    pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
206        self.metadata.expiration_policy = expiration_policy;
207        self
208    }
209
210    /// Sets the content type of the object to be uploaded.
211    ///
212    /// You can use the utility function [`crate::utils::guess_mime_type`] to attempt to guess a
213    /// `content_type` based on magic bytes.
214    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
215        self.metadata.content_type = content_type.into();
216        self
217    }
218
219    /// Sets the origin of the object, typically the IP address of the original source.
220    ///
221    /// This is an optional but encouraged field that tracks where the payload was
222    /// originally obtained from. For example, the IP address of the Sentry SDK or CLI
223    /// that uploaded the data.
224    ///
225    /// # Example
226    ///
227    /// ```no_run
228    /// # async fn example(session: objectstore_client::Session) {
229    /// session.put("data")
230    ///     .origin("203.0.113.42")
231    ///     .send()
232    ///     .await
233    ///     .unwrap();
234    /// # }
235    /// ```
236    pub fn origin(mut self, origin: impl Into<String>) -> Self {
237        self.metadata.origin = Some(origin.into());
238        self
239    }
240
241    /// Sets the filename of the object.
242    ///
243    /// When present, the server will include a `Content-Disposition: attachment; filename="<filename>"`
244    /// header in GET responses, prompting browsers and download tools to save the file under
245    /// this name.
246    pub fn filename(mut self, filename: impl Into<String>) -> Self {
247        self.metadata.filename = Some(filename.into());
248        self
249    }
250
251    /// This sets the custom metadata to the provided map.
252    ///
253    /// It will clear any previously set metadata.
254    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
255        self.metadata.custom = metadata.into();
256        self
257    }
258
259    /// Appends they `key`/`value` to the custom metadata of this object.
260    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
261        self.metadata.custom.insert(key.into(), value.into());
262        self
263    }
264}
265
266/// Turns the body into a request body, compressing it if the mode asks for it.
267///
268/// Payloads declared as [`CompressionMode::Precompressed`] are forwarded verbatim.
269pub(crate) async fn encode_body(body: PutBody, mode: Option<CompressionMode>) -> io::Result<Body> {
270    let compression = match mode {
271        Some(CompressionMode::Compress(compression)) => Some(compression),
272        // The payload already carries the encoding, so nothing is left to do here.
273        Some(CompressionMode::Precompressed(_)) | None => None,
274    };
275
276    Ok(match compression {
277        None => match body {
278            PutBody::Buffer(bytes) => bytes.into(),
279            PutBody::Stream(stream) => Body::wrap_stream(stream),
280            PutBody::File(file) => Body::wrap_stream(ReaderStream::new(file)),
281            PutBody::Path(path) => Body::wrap_stream(ReaderStream::new(File::open(path).await?)),
282        },
283        Some(compression) => match body {
284            PutBody::Buffer(bytes) => compress(Cursor::new(bytes), compression),
285            PutBody::Stream(stream) => compress(StreamReader::new(stream), compression),
286            PutBody::File(file) => compress(BufReader::new(file), compression),
287            PutBody::Path(path) => compress(BufReader::new(File::open(path).await?), compression),
288        },
289    })
290}
291
292fn compress(reader: impl AsyncBufRead + Send + 'static, compression: Compression) -> Body {
293    match compression {
294        Compression::Zstd => Body::wrap_stream(ReaderStream::new(ZstdEncoder::new(reader))),
295        Compression::Gzip => Body::wrap_stream(ReaderStream::new(GzipEncoder::new(reader))),
296    }
297}
298
299// TODO: instead of a separate `send` method, it would be nice to just implement `IntoFuture`.
300// However, `IntoFuture` needs to define the resulting future as an associated type,
301// and "impl trait in associated type position" is not yet stable :-(
302impl PutBuilder {
303    /// Sends the built put request to the upstream service.
304    pub async fn send(mut self) -> crate::Result<PutResponse> {
305        let method = match self.key {
306            Some(_) => reqwest::Method::PUT,
307            None => reqwest::Method::POST,
308        };
309
310        let mut builder = self
311            .session
312            .request(method, self.key.as_deref().unwrap_or_default())?;
313
314        self.metadata.compression = self.compression.map(CompressionMode::compression);
315        let body = encode_body(self.body, self.compression).await?;
316
317        builder = builder.headers(self.metadata.to_headers("")?);
318
319        let response = builder.body(body).send().await?;
320        Ok(response.error_for_status_and_drain().await?.json().await?)
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use std::io::Read as _;
327
328    use futures_util::stream;
329    use http_body_util::BodyExt as _;
330
331    use super::*;
332
333    fn zstd_compress(data: &[u8]) -> Vec<u8> {
334        zstd::encode_all(Cursor::new(data), 0).expect("zstd encoding to succeed")
335    }
336
337    fn stream_body(chunks: Vec<&'static [u8]>) -> PutBody {
338        let chunks = chunks.into_iter().map(|c| Ok(Bytes::from_static(c)));
339        PutBody::Stream(stream::iter(chunks).boxed())
340    }
341
342    async fn collect(body: Body) -> Vec<u8> {
343        body.collect()
344            .await
345            .expect("body to be readable")
346            .to_bytes()
347            .to_vec()
348    }
349
350    #[tokio::test]
351    async fn compress_buffer_compresses() {
352        let body = PutBody::Buffer(Bytes::from_static(b"hello world"));
353        let mode = Some(CompressionMode::Compress(Compression::Zstd));
354
355        let encoded = collect(encode_body(body, mode).await.unwrap()).await;
356        assert_eq!(encoded, zstd_compress(b"hello world"));
357    }
358
359    #[tokio::test]
360    async fn compress_stream_compresses() {
361        let body = stream_body(vec![b"hello ", b"world"]);
362        let mode = Some(CompressionMode::Compress(Compression::Zstd));
363
364        let encoded = collect(encode_body(body, mode).await.unwrap()).await;
365        assert_eq!(
366            zstd::decode_all(Cursor::new(encoded)).unwrap(),
367            b"hello world"
368        );
369    }
370
371    #[tokio::test]
372    async fn gzip_body_kinds() {
373        let dir = tempfile::tempdir().unwrap();
374        let path = dir.path().join("payload");
375        std::fs::write(&path, b"hello world").unwrap();
376
377        let bodies = [
378            PutBody::Buffer(Bytes::from_static(b"hello world")),
379            stream_body(vec![b"hello ", b"world"]),
380            PutBody::File(File::open(&path).await.unwrap()),
381            PutBody::Path(path.clone()),
382        ];
383
384        for body in bodies {
385            let mode = Some(CompressionMode::Compress(Compression::Gzip));
386            let encoded = collect(encode_body(body, mode).await.unwrap()).await;
387
388            let mut decoded = Vec::new();
389            flate2::read::GzDecoder::new(&encoded[..])
390                .read_to_end(&mut decoded)
391                .unwrap();
392            assert_eq!(decoded, b"hello world");
393        }
394    }
395
396    #[tokio::test]
397    async fn precompressed_buffer_is_forwarded_verbatim() {
398        let compressed = zstd_compress(b"hello world");
399        let body = PutBody::Buffer(Bytes::from(compressed.clone()));
400        let mode = Some(CompressionMode::Precompressed(Compression::Zstd));
401
402        let encoded = collect(encode_body(body, mode).await.unwrap()).await;
403        assert_eq!(encoded, compressed);
404    }
405
406    #[tokio::test]
407    async fn precompressed_stream_is_forwarded_verbatim() {
408        let body = stream_body(vec![b"\x28\xb5\x2f\xfd", b"trailing"]);
409        let mode = Some(CompressionMode::Precompressed(Compression::Zstd));
410
411        let encoded = collect(encode_body(body, mode).await.unwrap()).await;
412        assert_eq!(encoded, b"\x28\xb5\x2f\xfdtrailing");
413    }
414
415    #[tokio::test]
416    async fn without_compression_is_forwarded_verbatim() {
417        let body = PutBody::Buffer(Bytes::from_static(b"hello world"));
418
419        let encoded = collect(encode_body(body, None).await.unwrap()).await;
420        assert_eq!(encoded, b"hello world");
421    }
422}
objectstore_client/put.rs

objectstore_client/
put.rs
