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

objectstore_client/
multipart.rs
