objectstore_client/
multipart.rs

1use std::borrow::Cow;
2use std::collections::BTreeMap;
3
4use base64::Engine as _;
5use bytes::Bytes;
6use futures_util::StreamExt as _;
7use objectstore_types::metadata::Metadata;
8use objectstore_types::multipart::{
9    CompleteErrorDetail, CompleteRequest, CompleteSuccessResponse, InitiateResponse,
10    ListPartsResponse, UploadPartResponse,
11};
12use reqwest::Body;
13use serde::Deserialize;
14use tokio::io::AsyncRead;
15use tokio_util::io::ReaderStream;
16
17use crate::response::ResponseExt as _;
18use crate::{ClientStream, ObjectKey, Session};
19
20pub use objectstore_types::multipart::CompletePart;
21pub use objectstore_types::multipart::ETag;
22pub use objectstore_types::multipart::PartInfo;
23pub use objectstore_types::multipart::PartNumber;
24pub use objectstore_types::multipart::UploadId;
25
26#[derive(Deserialize)]
27#[serde(untagged)]
28enum CompleteResponse {
29    Error { error: CompleteErrorDetail },
30    Success(CompleteSuccessResponse),
31}
32
33impl Session {
34    /// Creates a builder for initiating a multipart upload.
35    ///
36    /// The returned [`InitiateMultipartBuilder`] inherits the session's default compression
37    /// and expiration settings.
38    ///
39    /// IMPORTANT: unlike single-object uploads, the client does not automatically compress the
40    /// contents of [`MultipartUpload::put`]/[`MultipartUpload::put_stream`] based on the
41    /// configured `compression`.
42    /// The caller is responsible to compress the payload in accordance with the configured
43    /// `compression`.
44    /// That's because we require `content_length` on each part to be the length of the compressed
45    /// content, which we wouldn't be able to know beforehand if `objectstore_client` automatically
46    /// compressed payloads on the fly.
47    pub fn initiate_multipart_upload(&self) -> InitiateMultipartBuilder {
48        let metadata = Metadata {
49            expiration_policy: self.scope.usecase().expiration_policy(),
50            compression: self.scope.usecase().compression(),
51            ..Default::default()
52        };
53
54        InitiateMultipartBuilder {
55            session: self.clone(),
56            metadata,
57            key: None,
58        }
59    }
60
61    /// Resumes an existing multipart upload from its key and upload ID.
62    ///
63    /// This reconstructs a [`MultipartUpload`] handle from previously obtained identifiers, and
64    /// doesn't make any network calls.
65    /// Use this to resume an upload after a process restart or to continue an upload initiated elsewhere.
66    pub fn resume_multipart_upload(
67        &self,
68        key: impl Into<ObjectKey>,
69        upload_id: impl Into<String>,
70    ) -> crate::Result<MultipartUpload> {
71        Ok(MultipartUpload {
72            session: self.clone(),
73            key: key.into(),
74            upload_id: UploadId::new(upload_id.into())?,
75        })
76    }
77}
78
79/// A builder for initiating a multipart upload.
80#[derive(Debug)]
81pub struct InitiateMultipartBuilder {
82    session: Session,
83    metadata: Metadata,
84    key: Option<ObjectKey>,
85}
86
87impl InitiateMultipartBuilder {
88    /// Sets an explicit object key.
89    ///
90    /// If a key is specified, the object will be stored under that key. Otherwise, the Objectstore
91    /// server will automatically assign a random key, which is then returned from this request.
92    pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
93        self.key = Some(key.into()).filter(|k| !k.is_empty());
94        self
95    }
96
97    /// Sets the compression algorithm recorded in this object's metadata.
98    ///
99    /// IMPORTANT: unlike single-object uploads, the client does not automatically compress the
100    /// contents of [`MultipartUpload::put`]/[`MultipartUpload::put_stream`] based on the
101    /// configured `compression`.
102    /// The caller is responsible to compress the payload in accordance with the configured
103    /// `compression`.
104    ///
105    /// By default, the compression algorithm set on this Session's Usecase is used.
106    pub fn compression(mut self, compression: impl Into<Option<crate::Compression>>) -> Self {
107        self.metadata.compression = compression.into();
108        self
109    }
110
111    /// Sets the expiration policy of the object to be uploaded.
112    ///
113    /// By default, the expiration policy set on this Session's Usecase is used.
114    pub fn expiration_policy(mut self, expiration_policy: crate::ExpirationPolicy) -> Self {
115        self.metadata.expiration_policy = expiration_policy;
116        self
117    }
118
119    /// Sets the content type of the object to be uploaded.
120    ///
121    /// You can use the utility function [`crate::utils::guess_mime_type`] to attempt to guess a
122    /// `content_type` based on magic bytes.
123    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
124        self.metadata.content_type = content_type.into();
125        self
126    }
127
128    /// Sets the origin of the object, typically the IP address of the original source.
129    ///
130    /// This is an optional but encouraged field that tracks where the payload was
131    /// originally obtained from. For example, the IP address of the Sentry SDK or CLI
132    /// that uploaded the data.
133    ///
134    /// # Example
135    ///
136    /// ```no_run
137    /// # async fn example(session: objectstore_client::Session) {
138    /// session.initiate_multipart_upload()
139    ///     .origin("203.0.113.42")
140    ///     .send()
141    ///     .await
142    ///     .unwrap();
143    /// # }
144    /// ```
145    pub fn origin(mut self, origin: impl Into<String>) -> Self {
146        self.metadata.origin = Some(origin.into());
147        self
148    }
149
150    /// Sets the filename of the object.
151    ///
152    /// When present, the server will include a `Content-Disposition: attachment; filename="<filename>"`
153    /// header in GET responses.
154    pub fn filename(mut self, filename: impl Into<String>) -> Self {
155        self.metadata.filename = Some(filename.into());
156        self
157    }
158
159    /// Sets the custom metadata to the provided map.
160    ///
161    /// It will clear any previously set metadata.
162    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
163        self.metadata.custom = metadata.into();
164        self
165    }
166
167    /// Appends the `key`/`value` to the custom metadata of this object.
168    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
169        self.metadata.custom.insert(key.into(), value.into());
170        self
171    }
172
173    /// Sends the initiate request and returns a [`MultipartUpload`] handle.
174    pub async fn send(self) -> crate::Result<MultipartUpload> {
175        let method = match self.key {
176            Some(_) => reqwest::Method::PUT,
177            None => reqwest::Method::POST,
178        };
179
180        let mut builder =
181            self.session
182                .multipart_request(method, None, self.key.as_deref(), None)?;
183
184        builder = builder.headers(self.metadata.to_headers("")?);
185
186        let response: InitiateResponse = builder
187            .send()
188            .await?
189            .error_for_status_and_drain()
190            .await?
191            .json()
192            .await?;
193
194        Ok(MultipartUpload {
195            session: self.session,
196            key: response.key,
197            upload_id: response.upload_id,
198        })
199    }
200}
201
202/// Represents an ongoing Multipart Upload, tied to a specific [`Session`] and [`UploadId`].
203///
204/// Create a Multipart Upload handle using [`Session::initiate_multipart_upload`] or [`Session::resume_multipart_upload`].
205#[derive(Debug)]
206pub struct MultipartUpload {
207    session: Session,
208    key: String,
209    upload_id: UploadId,
210}
211
212impl MultipartUpload {
213    /// Returns the upload session identifier.
214    pub fn upload_id(&self) -> &UploadId {
215        &self.upload_id
216    }
217
218    /// Returns the key of the object that this upload will create.
219    pub fn key(&self) -> &ObjectKey {
220        &self.key
221    }
222
223    /// Uploads a part using a [`Bytes`]-like payload.
224    ///
225    /// IMPORTANT: unlike single-object uploads, the client does not automatically compress
226    /// contents based on this upload's `Metadata::compression`.
227    /// The caller is responsible to compress the payload in accordance with the `compression`,
228    /// and, optionally, to pass the `content_md5` of the compressed payload.
229    pub async fn put(
230        &self,
231        body: impl Into<Bytes>,
232        part_number: u32,
233        content_md5: Option<&[u8; 16]>,
234    ) -> crate::Result<CompletePart> {
235        let bytes = body.into();
236        let content_length = bytes.len() as u64;
237        self.upload_part(bytes.into(), part_number, content_length, content_md5)
238            .await
239    }
240
241    /// Uploads a part using a streaming payload.
242    ///
243    /// IMPORTANT: unlike single-object uploads, the client does not automatically compress
244    /// contents based on this upload's `Metadata::compression`.
245    /// The caller is responsible to compress the payload in accordance with the `compression`,
246    /// and to pass the `content_length` and, optionally, `content_md5` of the compressed payload.
247    pub async fn put_stream(
248        &self,
249        stream: ClientStream,
250        part_number: u32,
251        content_length: u64,
252        content_md5: Option<&[u8; 16]>,
253    ) -> crate::Result<CompletePart> {
254        self.upload_part(
255            Body::wrap_stream(stream),
256            part_number,
257            content_length,
258            content_md5,
259        )
260        .await
261    }
262
263    /// Uploads a part from an [`AsyncRead`] source.
264    ///
265    /// IMPORTANT: unlike single-object uploads, the client does not automatically compress
266    /// contents based on this upload's `Metadata::compression`.
267    /// The caller is responsible to compress the payload in accordance with the `compression`,
268    /// and to pass the `content_length` and, optionally, `content_md5` of the compressed payload.
269    pub async fn put_read<R>(
270        &self,
271        reader: R,
272        part_number: u32,
273        content_length: u64,
274        content_md5: Option<&[u8; 16]>,
275    ) -> crate::Result<CompletePart>
276    where
277        R: AsyncRead + Send + Sync + 'static,
278    {
279        let stream = ReaderStream::new(reader).boxed();
280        self.put_stream(stream, part_number, content_length, content_md5)
281            .await
282    }
283
284    async fn upload_part(
285        &self,
286        body: Body,
287        part_number: u32,
288        content_length: u64,
289        content_md5: Option<&[u8; 16]>,
290    ) -> crate::Result<CompletePart> {
291        let part_number =
292            PartNumber::new(part_number).ok_or(crate::Error::InvalidPartNumber(part_number))?;
293
294        let mut builder = self
295            .session
296            .multipart_request(
297                reqwest::Method::PUT,
298                Some("parts"),
299                Some(&self.key),
300                Some(vec![
301                    ("upload_id", self.upload_id.to_string()),
302                    ("part_number", part_number.to_string()),
303                ]),
304            )?
305            .header(reqwest::header::CONTENT_LENGTH, content_length)
306            .body(body);
307
308        if let Some(md5) = content_md5 {
309            let encoded = base64::engine::general_purpose::STANDARD.encode(md5);
310            builder = builder.header("content-md5", encoded);
311        }
312
313        let response: UploadPartResponse = builder
314            .send()
315            .await?
316            .error_for_status_and_drain()
317            .await?
318            .json()
319            .await?;
320        Ok(CompletePart {
321            part_number,
322            etag: response.etag,
323        })
324    }
325
326    /// Lists all parts that have been uploaded for this multipart upload.
327    pub async fn list_parts(&self) -> crate::Result<Vec<PartInfo>> {
328        let mut all_parts = Vec::new();
329        let mut marker = None;
330
331        loop {
332            let page = self.list_parts_page(None, marker).await?;
333            all_parts.extend(page.parts);
334
335            if !page.is_truncated {
336                return Ok(all_parts);
337            }
338            marker = page.next_part_number_marker;
339            if marker.is_none() {
340                return Err(crate::Error::MalformedResponse(
341                    "server returned is_truncated=true but no next_part_number_marker. Please report a bug.".into(),
342                ));
343            }
344        }
345    }
346
347    async fn list_parts_page(
348        &self,
349        max_parts: Option<u32>,
350        part_number_marker: Option<PartNumber>,
351    ) -> crate::Result<ListPartsResponse> {
352        let mut params: Vec<(&str, String)> = vec![("upload_id", self.upload_id.to_string())];
353        if let Some(max) = max_parts {
354            params.push(("max_parts", max.to_string()));
355        }
356        if let Some(marker) = part_number_marker {
357            params.push(("part_number_marker", marker.to_string()));
358        }
359
360        let builder = self.session.multipart_request(
361            reqwest::Method::GET,
362            Some("parts"),
363            Some(&self.key),
364            Some(params),
365        )?;
366
367        let response: ListPartsResponse = builder
368            .send()
369            .await?
370            .error_for_status_and_drain()
371            .await?
372            .json()
373            .await?;
374        Ok(response)
375    }
376
377    /// Aborts this multipart upload.
378    pub async fn abort(self) -> crate::Result<()> {
379        let builder = self.session.multipart_request(
380            reqwest::Method::DELETE,
381            None,
382            Some(&self.key),
383            Some(vec![("upload_id", self.upload_id.to_string())]),
384        )?;
385        builder
386            .send()
387            .await?
388            .error_for_status_and_drain()
389            .await?
390            .drain_body()
391            .await;
392        Ok(())
393    }
394
395    /// Completes the multipart upload, assembling all parts into the final object.
396    pub async fn complete(
397        self,
398        parts: impl IntoIterator<Item = CompletePart>,
399    ) -> crate::Result<ObjectKey> {
400        let mut parts: Vec<_> = parts.into_iter().collect();
401        parts.sort_by_key(|p| p.part_number);
402
403        let builder = self
404            .session
405            .multipart_request(
406                reqwest::Method::POST,
407                Some("complete"),
408                Some(&self.key),
409                Some(vec![("upload_id", self.upload_id.to_string())]),
410            )?
411            .json(&CompleteRequest { parts });
412
413        let response = builder.send().await?.error_for_status_and_drain().await?;
414        match response.json::<CompleteResponse>().await? {
415            CompleteResponse::Success(s) => Ok(s.key),
416            CompleteResponse::Error { error } => Err(crate::Error::MultipartComplete {
417                code: error.code,
418                message: error.message,
419            }),
420        }
421    }
422}
objectstore_client/multipart.rs

objectstore_client/
multipart.rs
