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

objectstore_client/
put.rs
