objectstore_client/
put.rs

1use std::fmt;
2use std::io::Cursor;
3use std::{borrow::Cow, collections::BTreeMap};
4
5use async_compression::tokio::bufread::ZstdEncoder;
6use bytes::Bytes;
7use futures_util::StreamExt;
8use objectstore_types::Metadata;
9use reqwest::Body;
10use serde::Deserialize;
11use tokio::io::AsyncRead;
12use tokio_util::io::{ReaderStream, StreamReader};
13
14pub use objectstore_types::{Compression, ExpirationPolicy};
15
16use crate::{ClientStream, Session};
17
18/// The response returned from the service after uploading an object.
19#[derive(Debug, Deserialize)]
20pub struct PutResponse {
21    /// The key of the object, as stored.
22    pub key: String,
23}
24
25pub(crate) enum PutBody {
26    Buffer(Bytes),
27    Stream(ClientStream),
28}
29
30impl fmt::Debug for PutBody {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        f.debug_tuple("PutBody").finish_non_exhaustive()
33    }
34}
35
36impl Session {
37    fn put_body(&self, body: PutBody) -> PutBuilder {
38        let metadata = Metadata {
39            expiration_policy: self.scope.usecase().expiration_policy(),
40            compression: Some(self.scope.usecase().compression()),
41            ..Default::default()
42        };
43
44        PutBuilder {
45            session: self.clone(),
46            metadata,
47            key: None,
48            body,
49        }
50    }
51
52    /// Creates or replaces an object using a [`Bytes`]-like payload.
53    pub fn put(&self, body: impl Into<Bytes>) -> PutBuilder {
54        self.put_body(PutBody::Buffer(body.into()))
55    }
56
57    /// Creates or replaces an object using a streaming payload.
58    pub fn put_stream(&self, body: ClientStream) -> PutBuilder {
59        self.put_body(PutBody::Stream(body))
60    }
61
62    /// Creates or replaces an object using an [`AsyncRead`] payload.
63    pub fn put_read<R>(&self, body: R) -> PutBuilder
64    where
65        R: AsyncRead + Send + Sync + 'static,
66    {
67        let stream = ReaderStream::new(body).boxed();
68        self.put_body(PutBody::Stream(stream))
69    }
70}
71
72/// A [`put`](Session::put) request builder.
73#[derive(Debug)]
74pub struct PutBuilder {
75    session: Session,
76    metadata: Metadata,
77    key: Option<String>,
78    body: PutBody,
79}
80
81impl PutBuilder {
82    /// Sets an explicit object key.
83    ///
84    /// If a key is specified, the object will be stored under that key. Otherwise, the Objectstore
85    /// server will automatically assign a random key, which is then returned from this request.
86    pub fn key(mut self, key: impl Into<String>) -> Self {
87        self.key = Some(key.into()).filter(|k| !k.is_empty());
88        self
89    }
90
91    /// Sets an explicit compression algorithm to be used for this payload.
92    ///
93    /// [`None`] should be used if no compression should be performed by the client,
94    /// either because the payload is uncompressible (such as a media format), or if the user
95    /// will handle any kind of compression, without the clients knowledge.
96    ///
97    /// By default, the compression algorithm set on this Session's Usecase is used.
98    pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
99        self.metadata.compression = compression.into();
100        self
101    }
102
103    /// Sets the expiration policy of the object to be uploaded.
104    ///
105    /// By default, the expiration policy set on this Session's Usecase is used.
106    pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
107        self.metadata.expiration_policy = expiration_policy;
108        self
109    }
110
111    /// Sets the content type of the object to be uploaded.
112    ///
113    /// You can use the utility function [`crate::utils::guess_mime_type`] to attempt to guess a
114    /// `content_type` based on magic bytes.
115    pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
116        self.metadata.content_type = content_type.into();
117        self
118    }
119
120    /// This sets the custom metadata to the provided map.
121    ///
122    /// It will clear any previously set metadata.
123    pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
124        self.metadata.custom = metadata.into();
125        self
126    }
127
128    /// Appends they `key`/`value` to the custom metadata of this object.
129    pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
130        self.metadata.custom.insert(key.into(), value.into());
131        self
132    }
133}
134
135// TODO: instead of a separate `send` method, it would be nice to just implement `IntoFuture`.
136// However, `IntoFuture` needs to define the resulting future as an associated type,
137// and "impl trait in associated type position" is not yet stable :-(
138impl PutBuilder {
139    /// Sends the built put request to the upstream service.
140    pub async fn send(self) -> crate::Result<PutResponse> {
141        let method = match self.key {
142            Some(_) => reqwest::Method::PUT,
143            None => reqwest::Method::POST,
144        };
145
146        let mut builder = self
147            .session
148            .request(method, self.key.as_deref().unwrap_or_default());
149
150        let body = match (self.metadata.compression, self.body) {
151            (Some(Compression::Zstd), PutBody::Buffer(bytes)) => {
152                let cursor = Cursor::new(bytes);
153                let encoder = ZstdEncoder::new(cursor);
154                let stream = ReaderStream::new(encoder);
155                Body::wrap_stream(stream)
156            }
157            (Some(Compression::Zstd), PutBody::Stream(stream)) => {
158                let stream = StreamReader::new(stream);
159                let encoder = ZstdEncoder::new(stream);
160                let stream = ReaderStream::new(encoder);
161                Body::wrap_stream(stream)
162            }
163            (None, PutBody::Buffer(bytes)) => bytes.into(),
164            (None, PutBody::Stream(stream)) => Body::wrap_stream(stream),
165            // _ => todo!("compression algorithms other than `zstd` are currently not supported"),
166        };
167
168        builder = builder.headers(self.metadata.to_headers("", false)?);
169
170        let response = builder.body(body).send().await?;
171        Ok(response.error_for_status()?.json().await?)
172    }
173}