objectstore-client 0.2.10

//! Low-level client API for resumable uploads.
//!
//! Enable the `resumable-upload-api` feature to use this API.
//! It maps directly to the server operations and is intended for advanced use cases that need
//! low-level protocol calls.
//! A future release of `objectstore-client` will use this API under the hood for eligible
//! `Session::put` calls, so users can benefit from this feature without the need to interact with
//! this API directly.

use std::borrow::Cow;
use std::collections::BTreeMap;
use std::fmt;

use bytes::Bytes;
use objectstore_types::metadata::Metadata;
use objectstore_types::resumable::{
    CompleteUploadResponse, CreateSessionResponse, HEADER_UPLOAD_LENGTH, HEADER_UPLOAD_OFFSET,
    UploadOffset,
};
use reqwest::{Body, Method, Response, StatusCode};
use serde::Serialize;
use tokio::io::AsyncRead;
use tokio_util::io::ReaderStream;

pub use objectstore_types::resumable::{SessionToken, UploadProgress};

use crate::response::ResponseExt as _;
use crate::{ClientStream, Compression, Error, ExpirationPolicy, ObjectKey, Session};

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
enum UploadType {
    Resumable,
}

#[derive(Serialize)]
struct UploadTypeQuery {
    upload_type: UploadType,
}

#[derive(Serialize)]
struct SessionQuery<'a> {
    session: &'a SessionToken,
}

/// A handle bound to one resumable upload session.
///
/// See the [crate-level documentation](crate#resumable-upload-api) for more information and
/// examples on using this API.
#[derive(Clone, Debug)]
pub struct ResumableUpload {
    session: Session,
    key: ObjectKey,
    token: SessionToken,
    total_length: Option<u64>,
    granularity: Option<u64>,
}

impl Session {
    /// Starts building a resumable upload for an object with `object_length` bytes.
    ///
    /// `object_length` is the total length of the object to be uploaded.
    /// Unlike regular uploads, the client does not automatically compress chunk contents.
    /// The caller must pre-compress the object according to `compression`, then create the upload
    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
    ///
    /// The object key is generated by the server unless one is supplied with
    /// [`CreateResumableUploadBuilder::key`].
    pub fn create_upload(&self, object_length: u64) -> CreateResumableUploadBuilder {
        let metadata = Metadata {
            expiration_policy: self.scope.usecase().expiration_policy(),
            ..Default::default()
        };

        CreateResumableUploadBuilder {
            session: self.clone(),
            total_length: object_length,
            key: None,
            metadata,
        }
    }

    /// Reconstructs a resumable upload handle without contacting the server.
    ///
    /// The `key` and `token` must come from the same previously created upload and scope. The
    /// server rejects a mismatched or unknown token when an operation is sent.
    pub fn resume_upload(&self, key: impl Into<ObjectKey>, token: SessionToken) -> ResumableUpload {
        ResumableUpload {
            session: self.clone(),
            key: key.into(),
            token,
            total_length: None,
            granularity: None,
        }
    }
}

impl ResumableUpload {
    /// Returns the object key bound to this upload.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Returns the session token bound to this upload.
    pub fn token(&self) -> &SessionToken {
        &self.token
    }

    /// Returns this upload's granularity in bytes, if known.
    ///
    /// The granularity is the persistence unit for non-final chunks: chunks shorter than one
    /// unit are rejected, and larger chunks may persist only an aligned prefix.
    pub fn granularity(&self) -> Option<u64> {
        self.granularity
    }

    /// Builds a request for the server's authoritative upload progress.
    pub fn progress(&self) -> UploadProgressBuilder {
        UploadProgressBuilder {
            upload: self.clone(),
        }
    }

    /// Builds a request to write `chunk` starting at `offset`.
    ///
    /// `put` doesn't perform any automatic compression of the payload, so the caller is
    /// responsible for applying compression to the entire payload beforehand and passing chunks
    /// of the already compressed payload to this method.
    ///
    /// The returned progress contains the authoritative server offset, which should be used for
    /// subsequent requests.
    pub fn put(&self, offset: u64, chunk: impl Into<Bytes>) -> PutChunkBuilder {
        let chunk = chunk.into();
        self.put_body(offset, chunk.len() as u64, chunk.into())
    }

    /// Builds a request to write a streaming chunk of `length` bytes at `offset`.
    ///
    /// The stream must yield exactly `length` bytes. If this upload records compression,
    /// those bytes must come from the object after it has been compressed in full.
    pub fn put_stream(&self, offset: u64, length: u64, stream: ClientStream) -> PutChunkBuilder {
        self.put_body(offset, length, Body::wrap_stream(stream))
    }

    /// Builds a request to write an [`AsyncRead`] chunk of `length` bytes at `offset`.
    ///
    /// The reader must produce exactly `length` bytes. If this upload records compression,
    /// those bytes must come from the object after it has been compressed in full.
    pub fn put_read<R>(&self, offset: u64, length: u64, reader: R) -> PutChunkBuilder
    where
        R: AsyncRead + Send + Sync + 'static,
    {
        self.put_body(offset, length, Body::wrap_stream(ReaderStream::new(reader)))
    }

    fn put_body(&self, offset: u64, length: u64, body: Body) -> PutChunkBuilder {
        PutChunkBuilder {
            upload: self.clone(),
            offset,
            length,
            body,
        }
    }

    fn validate_chunk_length(&self, offset: u64, chunk_length: u64) -> crate::Result<()> {
        let (Some(granularity), Some(total_length)) = (self.granularity, self.total_length) else {
            // A reconstructed handle cannot tell whether this is the final chunk. The backend
            // retains the authoritative total length and validates the request.
            return Ok(());
        };
        if granularity == 0 || chunk_length == 0 || chunk_length >= granularity {
            return Ok(());
        }
        if !offset
            .checked_add(chunk_length)
            .is_some_and(|end| end < total_length)
        {
            return Ok(());
        }

        Err(Error::ChunkTooSmall {
            chunk_length,
            upload_granularity: granularity,
        })
    }

    /// Builds a request to cancel this upload session, discarding any uploaded bytes.
    pub fn cancel(&self) -> CancelUploadBuilder {
        CancelUploadBuilder {
            upload: self.clone(),
        }
    }

    fn request(&self, method: Method) -> crate::Result<reqwest::RequestBuilder> {
        Ok(self
            .session
            .request(method, &self.key)?
            .query(&SessionQuery {
                session: &self.token,
            }))
    }
}

/// A builder for [`Session::create_upload`].
#[derive(Debug)]
pub struct CreateResumableUploadBuilder {
    session: Session,
    total_length: u64,
    key: Option<ObjectKey>,
    metadata: Metadata,
}

impl CreateResumableUploadBuilder {
    /// Sets an explicit object key instead of asking the server to generate one.
    pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
        self.key = Some(key.into()).filter(|key| !key.is_empty());
        self
    }

    /// Sets the object's content type.
    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
        self.metadata.content_type = content_type.into();
        self
    }

    /// Sets the object's expiration policy.
    pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
        self.metadata.expiration_policy = expiration_policy;
        self
    }

    /// Sets the compression algorithm recorded in this object's metadata.
    ///
    /// Unlike regular uploads, the client does not automatically compress chunk contents.
    /// The caller must pre-compress the object according to `compression`, then create the upload
    /// with the post-compression `object_length`. Offsets similarly refer to the bytes after compression.
    ///
    /// By default, no compression is set.
    pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
        self.metadata.compression = compression.into();
        self
    }

    /// Sets the origin of the object, typically the IP address of the original source.
    pub fn origin(mut self, origin: impl Into<String>) -> Self {
        self.metadata.origin = Some(origin.into());
        self
    }

    /// Sets the filename recorded for the object.
    pub fn filename(mut self, filename: impl Into<String>) -> Self {
        self.metadata.filename = Some(filename.into());
        self
    }

    /// Replaces all custom metadata for the object.
    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
        self.metadata.custom = metadata.into();
        self
    }

    /// Adds or replaces one custom metadata entry.
    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.metadata.custom.insert(key.into(), value.into());
        self
    }

    /// Creates the resumable upload and returns a handle bound to it.
    ///
    /// Returns `None` when the server declines resumable uploads for this object.
    /// Callers should fall back to a regular [`Session::put`] in that case.
    pub async fn send(self) -> crate::Result<Option<ResumableUpload>> {
        let method = if self.key.is_some() {
            Method::PUT
        } else {
            Method::POST
        };
        let request = self
            .session
            .request(method, self.key.as_deref().unwrap_or_default())?
            .query(&UploadTypeQuery {
                upload_type: UploadType::Resumable,
            })
            .headers(self.metadata.to_headers("")?)
            .header(HEADER_UPLOAD_LENGTH, self.total_length.to_string());
        let response = request.send().await?;

        match response.status() {
            StatusCode::OK => {}
            StatusCode::NOT_IMPLEMENTED => {
                response.drain_body().await;
                return Ok(None);
            }
            status => {
                let response = response.error_for_status_and_drain().await?;
                response.drain_body().await;
                return Err(Error::MalformedResponse(format!(
                    "unexpected HTTP status {status} while creating a resumable upload"
                )));
            }
        }

        let response: CreateSessionResponse = response.json().await?;
        let upload = ResumableUpload {
            session: self.session,
            key: response.key,
            token: response.session,
            total_length: Some(self.total_length),
            granularity: Some(response.granularity),
        };
        Ok(Some(upload))
    }
}

/// A builder for [`ResumableUpload::progress`].
#[derive(Debug)]
pub struct UploadProgressBuilder {
    upload: ResumableUpload,
}

impl UploadProgressBuilder {
    /// Queries the server's authoritative upload progress.
    ///
    /// # Errors
    ///
    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
    /// could not be found. The upload must be restarted with a new session in that case.
    pub async fn send(self) -> crate::Result<UploadProgress> {
        let response = self
            .upload
            .request(Method::PUT)?
            .header(HEADER_UPLOAD_OFFSET, "*")
            .send()
            .await?;
        parse_progress_response(response).await
    }
}

/// A builder for [`ResumableUpload::put`], [`ResumableUpload::put_stream`], or
/// [`ResumableUpload::put_read`].
pub struct PutChunkBuilder {
    upload: ResumableUpload,
    offset: u64,
    length: u64,
    body: Body,
}

impl fmt::Debug for PutChunkBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PutChunkBuilder")
            .field("upload", &self.upload)
            .field("offset", &self.offset)
            .field("length", &self.length)
            .finish()
    }
}

impl PutChunkBuilder {
    /// Writes this chunk and returns the server's authoritative progress.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UploadOffsetMismatch`] with the server's authoritative offset when
    /// `offset` differs from the server's current offset.
    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was canceled, or
    /// could not be found. The upload must be restarted with a new session in that case.
    /// Returns [`Error::ChunkTooSmall`] before sending the request when this is known to be a
    /// non-final chunk shorter than the upload granularity.
    ///
    /// ```rust,ignore
    /// let offset = match upload.put(offset, chunk).send().await {
    ///     Ok(UploadProgress::Complete) => return Ok(()),
    ///     Ok(UploadProgress::Incomplete { offset }) => offset,
    ///     Err(Error::UploadOffsetMismatch { offset }) => offset,
    ///     Err(error @ Error::ResumableUploadUnavailable) => todo!("retry the whole upload"),
    ///     Err(error) => todo!("handle error"),
    /// };
    /// ```
    pub async fn send(self) -> crate::Result<UploadProgress> {
        self.upload
            .validate_chunk_length(self.offset, self.length)?;
        let response = self
            .upload
            .request(Method::PUT)?
            .header(HEADER_UPLOAD_OFFSET, self.offset.to_string())
            .header(reqwest::header::CONTENT_LENGTH, self.length)
            .body(self.body)
            .send()
            .await?;
        parse_progress_response(response).await
    }
}

/// A builder for [`ResumableUpload::cancel`].
#[derive(Debug)]
pub struct CancelUploadBuilder {
    upload: ResumableUpload,
}

impl CancelUploadBuilder {
    /// Cancels this upload and discards any bytes already uploaded.
    ///
    /// # Errors
    ///
    /// Returns [`Error::ResumableUploadUnavailable`] when the session expired, was already
    /// canceled, or could not be found.
    pub async fn send(self) -> crate::Result<()> {
        let response = self.upload.request(Method::DELETE)?.send().await?;
        match response.status() {
            StatusCode::NO_CONTENT => {
                response.drain_body().await;
                Ok(())
            }
            StatusCode::NOT_FOUND | StatusCode::GONE => {
                response.drain_body().await;
                Err(Error::ResumableUploadUnavailable)
            }
            status => {
                let response = response.error_for_status_and_drain().await?;
                response.drain_body().await;
                Err(Error::MalformedResponse(format!(
                    "unexpected HTTP status {status} while canceling a resumable upload"
                )))
            }
        }
    }
}

async fn parse_progress_response(response: Response) -> crate::Result<UploadProgress> {
    match response.status() {
        status @ (StatusCode::NO_CONTENT | StatusCode::CONFLICT) => {
            let offset = parse_offset(&response);
            response.drain_body().await;
            let offset = offset.ok_or_else(|| {
                crate::Error::MalformedResponse(
                    "resumable upload response has no valid Upload-Offset header".into(),
                )
            })?;
            if status == StatusCode::CONFLICT {
                Err(Error::UploadOffsetMismatch { offset })
            } else {
                Ok(UploadProgress::Incomplete { offset })
            }
        }
        StatusCode::CREATED => {
            let _: CompleteUploadResponse = response.json().await?;
            Ok(UploadProgress::Complete)
        }
        StatusCode::NOT_FOUND | StatusCode::GONE => {
            response.drain_body().await;
            Err(Error::ResumableUploadUnavailable)
        }
        status => {
            let response = response.error_for_status_and_drain().await?;
            response.drain_body().await;
            Err(Error::MalformedResponse(format!(
                "unexpected HTTP status {status} while continuing a resumable upload"
            )))
        }
    }
}

fn parse_offset(response: &Response) -> Option<u64> {
    let value = response
        .headers()
        .get(HEADER_UPLOAD_OFFSET)?
        .to_str()
        .ok()?;
    match value.parse().ok()? {
        UploadOffset::At(offset) => Some(offset),
        UploadOffset::Unknown => None,
    }
}