Skip to main content

objectstore_client/
resumable.rs

1//! Low-level client API for resumable uploads.
2//!
3//! Enable the `resumable-upload-api` feature to use this API.
4//! It maps directly to the server operations and is intended for advanced use cases that need
5//! low-level protocol calls.
6//! A future release of `objectstore-client` will use this API under the hood for eligible
7//! `Session::put` calls, so users can benefit from this feature without the need to interact with
8//! this API directly.
9
10use std::borrow::Cow;
11use std::collections::BTreeMap;
12use std::fmt;
13
14use bytes::Bytes;
15use objectstore_types::metadata::Metadata;
16use objectstore_types::resumable::{
17    CompleteUploadResponse, CreateSessionResponse, HEADER_UPLOAD_LENGTH, HEADER_UPLOAD_OFFSET,
18    UploadOffset,
19};
20use reqwest::{Method, Response, StatusCode};
21use serde::Serialize;
22
23pub use objectstore_types::resumable::{SessionToken, UploadProgress};
24
25use crate::response::ResponseExt as _;
26use crate::{Compression, Error, ExpirationPolicy, ObjectKey, Session};
27
28#[derive(Serialize)]
29#[serde(rename_all = "snake_case")]
30enum UploadType {
31    Resumable,
32}
33
34#[derive(Serialize)]
35struct UploadTypeQuery {
36    upload_type: UploadType,
37}
38
39#[derive(Serialize)]
40struct SessionQuery<'a> {
41    session: &'a SessionToken,
42}
43
44/// A handle bound to one resumable upload session.
45///
46/// See the [crate-level documentation](crate#resumable-upload-api) for more information and
47/// examples on using this API.
48#[derive(Clone, Debug)]
49pub struct ResumableUpload {
50    session: Session,
51    key: ObjectKey,
52    token: SessionToken,
53}
54
55impl Session {
56    /// Starts building a resumable upload for an object with `object_length` bytes.
57    ///
58    /// `object_length` is the total length of the object to be uploaded.
59    /// Unlike regular uploads, the client does not automatically compress chunk contents.
60    /// The caller must pre-compress the object according to `compression`, then create the upload
61    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
62    ///
63    /// The object key is generated by the server unless one is supplied with
64    /// [`CreateResumableUploadBuilder::key`].
65    pub fn create_upload(&self, object_length: u64) -> CreateResumableUploadBuilder {
66        let metadata = Metadata {
67            expiration_policy: self.scope.usecase().expiration_policy(),
68            ..Default::default()
69        };
70
71        CreateResumableUploadBuilder {
72            session: self.clone(),
73            total_length: object_length,
74            key: None,
75            metadata,
76        }
77    }
78
79    /// Reconstructs a resumable upload handle without contacting the server.
80    ///
81    /// The `key` and `token` must come from the same previously created upload and scope. The
82    /// server rejects a mismatched or unknown token when an operation is sent.
83    pub fn resume_upload(&self, key: impl Into<ObjectKey>, token: SessionToken) -> ResumableUpload {
84        ResumableUpload {
85            session: self.clone(),
86            key: key.into(),
87            token,
88        }
89    }
90}
91
92impl ResumableUpload {
93    /// Returns the object key bound to this upload.
94    pub fn key(&self) -> &str {
95        &self.key
96    }
97
98    /// Returns the session token bound to this upload.
99    pub fn token(&self) -> &SessionToken {
100        &self.token
101    }
102
103    /// Builds a request for the server's authoritative upload progress.
104    pub fn progress(&self) -> UploadProgressBuilder {
105        UploadProgressBuilder {
106            upload: self.clone(),
107        }
108    }
109
110    /// Builds a request to write `chunk` starting at `offset`.
111    ///
112    /// `put` doesn't perform any automatic compression of the payload, so the caller is
113    /// responsible for applying compression to the entire payload beforehand and passing chunks
114    /// of the already compressed payload to this method.
115    ///
116    /// The returned progress contains the authoritative server offset, which should be used for
117    /// subsequent requests.
118    pub fn put(&self, offset: u64, chunk: impl Into<Bytes>) -> PutChunkBuilder {
119        PutChunkBuilder {
120            upload: self.clone(),
121            offset,
122            chunk: chunk.into(),
123        }
124    }
125
126    /// Builds a request to cancel this upload session, discarding any uploaded bytes.
127    pub fn cancel(&self) -> CancelUploadBuilder {
128        CancelUploadBuilder {
129            upload: self.clone(),
130        }
131    }
132
133    fn request(&self, method: Method) -> crate::Result<reqwest::RequestBuilder> {
134        Ok(self
135            .session
136            .request(method, &self.key)?
137            .query(&SessionQuery {
138                session: &self.token,
139            }))
140    }
141}
142
143/// A builder for [`Session::create_upload`].
144#[derive(Debug)]
145pub struct CreateResumableUploadBuilder {
146    session: Session,
147    total_length: u64,
148    key: Option<ObjectKey>,
149    metadata: Metadata,
150}
151
152impl CreateResumableUploadBuilder {
153    /// Sets an explicit object key instead of asking the server to generate one.
154    pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
155        self.key = Some(key.into()).filter(|key| !key.is_empty());
156        self
157    }
158
159    /// Sets the object's content type.
160    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
161        self.metadata.content_type = content_type.into();
162        self
163    }
164
165    /// Sets the object's expiration policy.
166    pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
167        self.metadata.expiration_policy = expiration_policy;
168        self
169    }
170
171    /// Sets the compression algorithm recorded in this object's metadata.
172    ///
173    /// Unlike regular uploads, the client does not automatically compress chunk contents.
174    /// The caller must pre-compress the object according to `compression`, then create the upload
175    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
176    ///
177    /// By default, no compression is set.
178    pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
179        self.metadata.compression = compression.into();
180        self
181    }
182
183    /// Sets the origin of the object, typically the IP address of the original source.
184    pub fn origin(mut self, origin: impl Into<String>) -> Self {
185        self.metadata.origin = Some(origin.into());
186        self
187    }
188
189    /// Sets the filename recorded for the object.
190    pub fn filename(mut self, filename: impl Into<String>) -> Self {
191        self.metadata.filename = Some(filename.into());
192        self
193    }
194
195    /// Replaces all custom metadata for the object.
196    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
197        self.metadata.custom = metadata.into();
198        self
199    }
200
201    /// Adds or replaces one custom metadata entry.
202    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
203        self.metadata.custom.insert(key.into(), value.into());
204        self
205    }
206
207    /// Creates the resumable upload and returns a handle bound to it.
208    ///
209    /// Returns `None` when the server declines resumable uploads for this object.
210    /// Callers should fall back to a regular [`Session::put`] in that case.
211    pub async fn send(self) -> crate::Result<Option<ResumableUpload>> {
212        let method = if self.key.is_some() {
213            Method::PUT
214        } else {
215            Method::POST
216        };
217        let request = self
218            .session
219            .request(method, self.key.as_deref().unwrap_or_default())?
220            .query(&UploadTypeQuery {
221                upload_type: UploadType::Resumable,
222            })
223            .headers(self.metadata.to_headers("")?)
224            .header(HEADER_UPLOAD_LENGTH, self.total_length.to_string());
225        let response = request.send().await?;
226
227        match response.status() {
228            StatusCode::OK => {}
229            StatusCode::NOT_IMPLEMENTED => {
230                response.drain_body().await;
231                return Ok(None);
232            }
233            status => {
234                let response = response.error_for_status_and_drain().await?;
235                response.drain_body().await;
236                return Err(Error::MalformedResponse(format!(
237                    "unexpected HTTP status {status} while creating a resumable upload"
238                )));
239            }
240        }
241
242        let response: CreateSessionResponse = response.json().await?;
243        Ok(Some(
244            self.session.resume_upload(response.key, response.session),
245        ))
246    }
247}
248
249/// A builder for [`ResumableUpload::progress`].
250#[derive(Debug)]
251pub struct UploadProgressBuilder {
252    upload: ResumableUpload,
253}
254
255impl UploadProgressBuilder {
256    /// Queries the server's authoritative upload progress.
257    ///
258    /// # Errors
259    ///
260    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
261    /// could not be found. The upload must be restarted with a new session in that case.
262    pub async fn send(self) -> crate::Result<UploadProgress> {
263        let response = self
264            .upload
265            .request(Method::PUT)?
266            .header(HEADER_UPLOAD_OFFSET, "*")
267            .send()
268            .await?;
269        parse_progress_response(response).await
270    }
271}
272
273/// A builder for [`ResumableUpload::put`].
274pub struct PutChunkBuilder {
275    upload: ResumableUpload,
276    offset: u64,
277    chunk: Bytes,
278}
279
280impl fmt::Debug for PutChunkBuilder {
281    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
282        f.debug_struct("PutChunkBuilder")
283            .field("upload", &self.upload)
284            .field("offset", &self.offset)
285            .field("content_length", &self.chunk.len())
286            .finish()
287    }
288}
289
290impl PutChunkBuilder {
291    /// Writes this chunk and returns the server's authoritative progress.
292    ///
293    /// If `offset` differs from the server's current offset, the response is returned as
294    /// [`UploadProgress::Incomplete`] with the server's authoritative offset.
295    ///
296    /// # Errors
297    ///
298    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
299    /// could not be found. The upload must be restarted with a new session in that case.
300    ///
301    /// ```rust,ignore
302    /// let offset = match upload.put(offset, chunk).send().await {
303    ///     Ok(UploadProgress::Complete) => return Ok(()),
304    ///     Ok(UploadProgress::Incomplete { offset }) => offset,
305    ///     Err(error @ Error::ResumableUploadUnavailable) => todo!("retry the whole upload"),
306    ///     Err(error) => todo!("handle error"),
307    /// };
308    /// ```
309    pub async fn send(self) -> crate::Result<UploadProgress> {
310        let content_length = self.chunk.len();
311        let response = self
312            .upload
313            .request(Method::PUT)?
314            .header(HEADER_UPLOAD_OFFSET, self.offset.to_string())
315            .header(reqwest::header::CONTENT_LENGTH, content_length)
316            .body(self.chunk)
317            .send()
318            .await?;
319        parse_progress_response(response).await
320    }
321}
322
323/// A builder for [`ResumableUpload::cancel`].
324#[derive(Debug)]
325pub struct CancelUploadBuilder {
326    upload: ResumableUpload,
327}
328
329impl CancelUploadBuilder {
330    /// Cancels this upload and discards any bytes already uploaded.
331    ///
332    /// # Errors
333    ///
334    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was already
335    /// canceled, or could not be found.
336    pub async fn send(self) -> crate::Result<()> {
337        let response = self.upload.request(Method::DELETE)?.send().await?;
338        match response.status() {
339            StatusCode::NO_CONTENT => {
340                response.drain_body().await;
341                Ok(())
342            }
343            StatusCode::NOT_FOUND | StatusCode::GONE => {
344                response.drain_body().await;
345                Err(Error::ResumableUploadUnavailable)
346            }
347            status => {
348                let response = response.error_for_status_and_drain().await?;
349                response.drain_body().await;
350                Err(Error::MalformedResponse(format!(
351                    "unexpected HTTP status {status} while canceling a resumable upload"
352                )))
353            }
354        }
355    }
356}
357
358async fn parse_progress_response(response: Response) -> crate::Result<UploadProgress> {
359    match response.status() {
360        StatusCode::NO_CONTENT | StatusCode::CONFLICT => {
361            let offset = parse_offset(&response);
362            response.drain_body().await;
363            let offset = offset.ok_or_else(|| {
364                crate::Error::MalformedResponse(
365                    "resumable upload response has no valid Upload-Offset header".into(),
366                )
367            })?;
368            Ok(UploadProgress::Incomplete { offset })
369        }
370        StatusCode::CREATED => {
371            let _: CompleteUploadResponse = response.json().await?;
372            Ok(UploadProgress::Complete)
373        }
374        StatusCode::NOT_FOUND | StatusCode::GONE => {
375            response.drain_body().await;
376            Err(Error::ResumableUploadUnavailable)
377        }
378        status => {
379            let response = response.error_for_status_and_drain().await?;
380            response.drain_body().await;
381            Err(Error::MalformedResponse(format!(
382                "unexpected HTTP status {status} while continuing a resumable upload"
383            )))
384        }
385    }
386}
387
388fn parse_offset(response: &Response) -> Option<u64> {
389    let value = response
390        .headers()
391        .get(HEADER_UPLOAD_OFFSET)?
392        .to_str()
393        .ok()?;
394    match value.parse().ok()? {
395        UploadOffset::At(offset) => Some(offset),
396        UploadOffset::Unknown => None,
397    }
398}