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::{Body, Method, Response, StatusCode};
21use serde::Serialize;
22use tokio::io::AsyncRead;
23use tokio_util::io::ReaderStream;
24
25pub use objectstore_types::resumable::{SessionToken, UploadProgress};
26
27use crate::response::ResponseExt as _;
28use crate::{ClientStream, Compression, Error, ExpirationPolicy, ObjectKey, Session};
29
30#[derive(Serialize)]
31#[serde(rename_all = "snake_case")]
32enum UploadType {
33    Resumable,
34}
35
36#[derive(Serialize)]
37struct UploadTypeQuery {
38    upload_type: UploadType,
39}
40
41#[derive(Serialize)]
42struct SessionQuery<'a> {
43    session: &'a SessionToken,
44}
45
46/// A handle bound to one resumable upload session.
47///
48/// See the [crate-level documentation](crate#resumable-upload-api) for more information and
49/// examples on using this API.
50#[derive(Clone, Debug)]
51pub struct ResumableUpload {
52    session: Session,
53    key: ObjectKey,
54    token: SessionToken,
55    total_length: Option<u64>,
56    granularity: Option<u64>,
57}
58
59impl Session {
60    /// Starts building a resumable upload for an object with `object_length` bytes.
61    ///
62    /// `object_length` is the total length of the object to be uploaded.
63    /// Unlike regular uploads, the client does not automatically compress chunk contents.
64    /// The caller must pre-compress the object according to `compression`, then create the upload
65    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
66    ///
67    /// The object key is generated by the server unless one is supplied with
68    /// [`CreateResumableUploadBuilder::key`].
69    pub fn create_upload(&self, object_length: u64) -> CreateResumableUploadBuilder {
70        let metadata = Metadata {
71            expiration_policy: self.scope.usecase().expiration_policy(),
72            ..Default::default()
73        };
74
75        CreateResumableUploadBuilder {
76            session: self.clone(),
77            total_length: object_length,
78            key: None,
79            metadata,
80        }
81    }
82
83    /// Reconstructs a resumable upload handle without contacting the server.
84    ///
85    /// The `key` and `token` must come from the same previously created upload and scope. The
86    /// server rejects a mismatched or unknown token when an operation is sent.
87    pub fn resume_upload(&self, key: impl Into<ObjectKey>, token: SessionToken) -> ResumableUpload {
88        ResumableUpload {
89            session: self.clone(),
90            key: key.into(),
91            token,
92            total_length: None,
93            granularity: None,
94        }
95    }
96}
97
98impl ResumableUpload {
99    /// Returns the object key bound to this upload.
100    pub fn key(&self) -> &str {
101        &self.key
102    }
103
104    /// Returns the session token bound to this upload.
105    pub fn token(&self) -> &SessionToken {
106        &self.token
107    }
108
109    /// Returns this upload's granularity in bytes, if known.
110    ///
111    /// The granularity is the persistence unit for non-final chunks: chunks shorter than one
112    /// unit are rejected, and larger chunks may persist only an aligned prefix.
113    pub fn granularity(&self) -> Option<u64> {
114        self.granularity
115    }
116
117    /// Builds a request for the server's authoritative upload progress.
118    pub fn progress(&self) -> UploadProgressBuilder {
119        UploadProgressBuilder {
120            upload: self.clone(),
121        }
122    }
123
124    /// Builds a request to write `chunk` starting at `offset`.
125    ///
126    /// `put` doesn't perform any automatic compression of the payload, so the caller is
127    /// responsible for applying compression to the entire payload beforehand and passing chunks
128    /// of the already compressed payload to this method.
129    ///
130    /// The returned progress contains the authoritative server offset, which should be used for
131    /// subsequent requests.
132    pub fn put(&self, offset: u64, chunk: impl Into<Bytes>) -> PutChunkBuilder {
133        let chunk = chunk.into();
134        self.put_body(offset, chunk.len() as u64, chunk.into())
135    }
136
137    /// Builds a request to write a streaming chunk of `length` bytes at `offset`.
138    ///
139    /// The stream must yield exactly `length` bytes. If this upload records compression,
140    /// those bytes must come from the object after it has been compressed in full.
141    pub fn put_stream(&self, offset: u64, length: u64, stream: ClientStream) -> PutChunkBuilder {
142        self.put_body(offset, length, Body::wrap_stream(stream))
143    }
144
145    /// Builds a request to write an [`AsyncRead`] chunk of `length` bytes at `offset`.
146    ///
147    /// The reader must produce exactly `length` bytes. If this upload records compression,
148    /// those bytes must come from the object after it has been compressed in full.
149    pub fn put_read<R>(&self, offset: u64, length: u64, reader: R) -> PutChunkBuilder
150    where
151        R: AsyncRead + Send + Sync + 'static,
152    {
153        self.put_body(offset, length, Body::wrap_stream(ReaderStream::new(reader)))
154    }
155
156    fn put_body(&self, offset: u64, length: u64, body: Body) -> PutChunkBuilder {
157        PutChunkBuilder {
158            upload: self.clone(),
159            offset,
160            length,
161            body,
162        }
163    }
164
165    fn validate_chunk_length(&self, offset: u64, chunk_length: u64) -> crate::Result<()> {
166        let (Some(granularity), Some(total_length)) = (self.granularity, self.total_length) else {
167            // A reconstructed handle cannot tell whether this is the final chunk. The backend
168            // retains the authoritative total length and validates the request.
169            return Ok(());
170        };
171        if granularity == 0 || chunk_length == 0 || chunk_length >= granularity {
172            return Ok(());
173        }
174        if !offset
175            .checked_add(chunk_length)
176            .is_some_and(|end| end < total_length)
177        {
178            return Ok(());
179        }
180
181        Err(Error::ChunkTooSmall {
182            chunk_length,
183            upload_granularity: granularity,
184        })
185    }
186
187    /// Builds a request to cancel this upload session, discarding any uploaded bytes.
188    pub fn cancel(&self) -> CancelUploadBuilder {
189        CancelUploadBuilder {
190            upload: self.clone(),
191        }
192    }
193
194    fn request(&self, method: Method) -> crate::Result<reqwest::RequestBuilder> {
195        Ok(self
196            .session
197            .request(method, &self.key)?
198            .query(&SessionQuery {
199                session: &self.token,
200            }))
201    }
202}
203
204/// A builder for [`Session::create_upload`].
205#[derive(Debug)]
206pub struct CreateResumableUploadBuilder {
207    session: Session,
208    total_length: u64,
209    key: Option<ObjectKey>,
210    metadata: Metadata,
211}
212
213impl CreateResumableUploadBuilder {
214    /// Sets an explicit object key instead of asking the server to generate one.
215    pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
216        self.key = Some(key.into()).filter(|key| !key.is_empty());
217        self
218    }
219
220    /// Sets the object's content type.
221    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
222        self.metadata.content_type = content_type.into();
223        self
224    }
225
226    /// Sets the object's expiration policy.
227    pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
228        self.metadata.expiration_policy = expiration_policy;
229        self
230    }
231
232    /// Sets the compression algorithm recorded in this object's metadata.
233    ///
234    /// Unlike regular uploads, the client does not automatically compress chunk contents.
235    /// The caller must pre-compress the object according to `compression`, then create the upload
236    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
237    ///
238    /// By default, no compression is set.
239    pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
240        self.metadata.compression = compression.into();
241        self
242    }
243
244    /// Sets the origin of the object, typically the IP address of the original source.
245    pub fn origin(mut self, origin: impl Into<String>) -> Self {
246        self.metadata.origin = Some(origin.into());
247        self
248    }
249
250    /// Sets the filename recorded for the object.
251    pub fn filename(mut self, filename: impl Into<String>) -> Self {
252        self.metadata.filename = Some(filename.into());
253        self
254    }
255
256    /// Replaces all custom metadata for the object.
257    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
258        self.metadata.custom = metadata.into();
259        self
260    }
261
262    /// Adds or replaces one custom metadata entry.
263    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
264        self.metadata.custom.insert(key.into(), value.into());
265        self
266    }
267
268    /// Creates the resumable upload and returns a handle bound to it.
269    ///
270    /// Returns `None` when the server declines resumable uploads for this object.
271    /// Callers should fall back to a regular [`Session::put`] in that case.
272    pub async fn send(self) -> crate::Result<Option<ResumableUpload>> {
273        let method = if self.key.is_some() {
274            Method::PUT
275        } else {
276            Method::POST
277        };
278        let request = self
279            .session
280            .request(method, self.key.as_deref().unwrap_or_default())?
281            .query(&UploadTypeQuery {
282                upload_type: UploadType::Resumable,
283            })
284            .headers(self.metadata.to_headers("")?)
285            .header(HEADER_UPLOAD_LENGTH, self.total_length.to_string());
286        let response = request.send().await?;
287
288        match response.status() {
289            StatusCode::OK => {}
290            StatusCode::NOT_IMPLEMENTED => {
291                response.drain_body().await;
292                return Ok(None);
293            }
294            status => {
295                let response = response.error_for_status_and_drain().await?;
296                response.drain_body().await;
297                return Err(Error::MalformedResponse(format!(
298                    "unexpected HTTP status {status} while creating a resumable upload"
299                )));
300            }
301        }
302
303        let response: CreateSessionResponse = response.json().await?;
304        let upload = ResumableUpload {
305            session: self.session,
306            key: response.key,
307            token: response.session,
308            total_length: Some(self.total_length),
309            granularity: Some(response.granularity),
310        };
311        Ok(Some(upload))
312    }
313}
314
315/// A builder for [`ResumableUpload::progress`].
316#[derive(Debug)]
317pub struct UploadProgressBuilder {
318    upload: ResumableUpload,
319}
320
321impl UploadProgressBuilder {
322    /// Queries the server's authoritative upload progress.
323    ///
324    /// # Errors
325    ///
326    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
327    /// could not be found. The upload must be restarted with a new session in that case.
328    pub async fn send(self) -> crate::Result<UploadProgress> {
329        let response = self
330            .upload
331            .request(Method::PUT)?
332            .header(HEADER_UPLOAD_OFFSET, "*")
333            .send()
334            .await?;
335        parse_progress_response(response).await
336    }
337}
338
339/// A builder for [`ResumableUpload::put`], [`ResumableUpload::put_stream`], or
340/// [`ResumableUpload::put_read`].
341pub struct PutChunkBuilder {
342    upload: ResumableUpload,
343    offset: u64,
344    length: u64,
345    body: Body,
346}
347
348impl fmt::Debug for PutChunkBuilder {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        f.debug_struct("PutChunkBuilder")
351            .field("upload", &self.upload)
352            .field("offset", &self.offset)
353            .field("length", &self.length)
354            .finish()
355    }
356}
357
358impl PutChunkBuilder {
359    /// Writes this chunk and returns the server's authoritative progress.
360    ///
361    /// # Errors
362    ///
363    /// Returns [`Error::UploadOffsetMismatch`] with the server's authoritative offset when
364    /// `offset` differs from the server's current offset.
365    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
366    /// could not be found. The upload must be restarted with a new session in that case.
367    /// Returns [`Error::ChunkTooSmall`] before sending the request when this is known to be a
368    /// non-final chunk shorter than the upload granularity.
369    ///
370    /// ```rust,ignore
371    /// let offset = match upload.put(offset, chunk).send().await {
372    ///     Ok(UploadProgress::Complete) => return Ok(()),
373    ///     Ok(UploadProgress::Incomplete { offset }) => offset,
374    ///     Err(Error::UploadOffsetMismatch { offset }) => offset,
375    ///     Err(error @ Error::ResumableUploadUnavailable) => todo!("retry the whole upload"),
376    ///     Err(error) => todo!("handle error"),
377    /// };
378    /// ```
379    pub async fn send(self) -> crate::Result<UploadProgress> {
380        self.upload
381            .validate_chunk_length(self.offset, self.length)?;
382        let response = self
383            .upload
384            .request(Method::PUT)?
385            .header(HEADER_UPLOAD_OFFSET, self.offset.to_string())
386            .header(reqwest::header::CONTENT_LENGTH, self.length)
387            .body(self.body)
388            .send()
389            .await?;
390        parse_progress_response(response).await
391    }
392}
393
394/// A builder for [`ResumableUpload::cancel`].
395#[derive(Debug)]
396pub struct CancelUploadBuilder {
397    upload: ResumableUpload,
398}
399
400impl CancelUploadBuilder {
401    /// Cancels this upload and discards any bytes already uploaded.
402    ///
403    /// # Errors
404    ///
405    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was already
406    /// canceled, or could not be found.
407    pub async fn send(self) -> crate::Result<()> {
408        let response = self.upload.request(Method::DELETE)?.send().await?;
409        match response.status() {
410            StatusCode::NO_CONTENT => {
411                response.drain_body().await;
412                Ok(())
413            }
414            StatusCode::NOT_FOUND | StatusCode::GONE => {
415                response.drain_body().await;
416                Err(Error::ResumableUploadUnavailable)
417            }
418            status => {
419                let response = response.error_for_status_and_drain().await?;
420                response.drain_body().await;
421                Err(Error::MalformedResponse(format!(
422                    "unexpected HTTP status {status} while canceling a resumable upload"
423                )))
424            }
425        }
426    }
427}
428
429async fn parse_progress_response(response: Response) -> crate::Result<UploadProgress> {
430    match response.status() {
431        status @ (StatusCode::NO_CONTENT | StatusCode::CONFLICT) => {
432            let offset = parse_offset(&response);
433            response.drain_body().await;
434            let offset = offset.ok_or_else(|| {
435                crate::Error::MalformedResponse(
436                    "resumable upload response has no valid Upload-Offset header".into(),
437                )
438            })?;
439            if status == StatusCode::CONFLICT {
440                Err(Error::UploadOffsetMismatch { offset })
441            } else {
442                Ok(UploadProgress::Incomplete { offset })
443            }
444        }
445        StatusCode::CREATED => {
446            let _: CompleteUploadResponse = response.json().await?;
447            Ok(UploadProgress::Complete)
448        }
449        StatusCode::NOT_FOUND | StatusCode::GONE => {
450            response.drain_body().await;
451            Err(Error::ResumableUploadUnavailable)
452        }
453        status => {
454            let response = response.error_for_status_and_drain().await?;
455            response.drain_body().await;
456            Err(Error::MalformedResponse(format!(
457                "unexpected HTTP status {status} while continuing a resumable upload"
458            )))
459        }
460    }
461}
462
463fn parse_offset(response: &Response) -> Option<u64> {
464    let value = response
465        .headers()
466        .get(HEADER_UPLOAD_OFFSET)?
467        .to_str()
468        .ok()?;
469    match value.parse().ok()? {
470        UploadOffset::At(offset) => Some(offset),
471        UploadOffset::Unknown => None,
472    }
473}