objectstore-client 0.2.6

//! 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::{Method, Response, StatusCode};
use serde::Serialize;

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

use crate::response::ResponseExt as _;
use crate::{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,
}

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,
        }
    }
}

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
    }

    /// 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 {
        PutChunkBuilder {
            upload: self.clone(),
            offset,
            chunk: chunk.into(),
        }
    }

    /// 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?;
        Ok(Some(
            self.session.resume_upload(response.key, response.session),
        ))
    }
}

/// 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`].
pub struct PutChunkBuilder {
    upload: ResumableUpload,
    offset: u64,
    chunk: Bytes,
}

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("content_length", &self.chunk.len())
            .finish()
    }
}

impl PutChunkBuilder {
    /// Writes this chunk and returns the server's authoritative progress.
    ///
    /// If `offset` differs from the server's current offset, the response is returned as
    /// [`UploadProgress::Incomplete`] with the server's authoritative offset.
    ///
    /// # 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.
    ///
    /// ```rust,ignore
    /// let offset = match upload.put(offset, chunk).send().await {
    ///     Ok(UploadProgress::Complete) => return Ok(()),
    ///     Ok(UploadProgress::Incomplete { offset }) => offset,
    ///     Err(error @ Error::ResumableUploadUnavailable) => todo!("retry the whole upload"),
    ///     Err(error) => todo!("handle error"),
    /// };
    /// ```
    pub async fn send(self) -> crate::Result<UploadProgress> {
        let content_length = self.chunk.len();
        let response = self
            .upload
            .request(Method::PUT)?
            .header(HEADER_UPLOAD_OFFSET, self.offset.to_string())
            .header(reqwest::header::CONTENT_LENGTH, content_length)
            .body(self.chunk)
            .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() {
        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(),
                )
            })?;
            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,
    }
}