objectstore_client/
put.rs1use std::fmt;
2use std::io::{self, Cursor};
3use std::path::PathBuf;
4use std::{borrow::Cow, collections::BTreeMap};
5
6use async_compression::tokio::bufread::ZstdEncoder;
7use bytes::Bytes;
8use futures_util::StreamExt;
9use objectstore_types::metadata::Metadata;
10use reqwest::Body;
11use serde::Deserialize;
12use tokio::fs::File;
13use tokio::io::{AsyncRead, BufReader};
14use tokio_util::io::{ReaderStream, StreamReader};
15
16pub use objectstore_types::metadata::{Compression, ExpirationPolicy};
17
18use crate::{ClientStream, ObjectKey, Session};
19
20#[derive(Debug, Deserialize)]
22pub struct PutResponse {
23 pub key: ObjectKey,
25}
26
27pub(crate) enum PutBody {
28 Buffer(Bytes),
29 Stream(ClientStream),
30 File(File),
31 Path(PathBuf),
32}
33
34impl fmt::Debug for PutBody {
35 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36 f.debug_tuple("PutBody").finish_non_exhaustive()
37 }
38}
39
40impl Session {
41 fn put_body(&self, body: PutBody) -> PutBuilder {
42 let metadata = Metadata {
43 expiration_policy: self.scope.usecase().expiration_policy(),
44 compression: self.scope.usecase().compression(),
45 ..Default::default()
46 };
47
48 PutBuilder {
49 session: self.clone(),
50 metadata,
51 key: None,
52 body,
53 }
54 }
55
56 pub fn put(&self, body: impl Into<Bytes>) -> PutBuilder {
58 self.put_body(PutBody::Buffer(body.into()))
59 }
60
61 pub fn put_stream(&self, body: ClientStream) -> PutBuilder {
63 self.put_body(PutBody::Stream(body))
64 }
65
66 pub fn put_read<R>(&self, body: R) -> PutBuilder
68 where
69 R: AsyncRead + Send + Sync + 'static,
70 {
71 let stream = ReaderStream::new(body).boxed();
72 self.put_body(PutBody::Stream(stream))
73 }
74
75 pub fn put_file(&self, file: File) -> PutBuilder {
83 self.put_body(PutBody::File(file))
84 }
85
86 pub fn put_path(&self, path: impl Into<PathBuf>) -> PutBuilder {
97 self.put_body(PutBody::Path(path.into()))
98 }
99}
100
101#[derive(Debug)]
103pub struct PutBuilder {
104 pub(crate) session: Session,
105 pub(crate) metadata: Metadata,
106 pub(crate) key: Option<ObjectKey>,
107 pub(crate) body: PutBody,
108}
109
110impl PutBuilder {
111 pub fn key(mut self, key: impl Into<ObjectKey>) -> Self {
116 self.key = Some(key.into()).filter(|k| !k.is_empty());
117 self
118 }
119
120 pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
128 self.metadata.compression = compression.into();
129 self
130 }
131
132 pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
136 self.metadata.expiration_policy = expiration_policy;
137 self
138 }
139
140 pub fn content_type(mut self, content_type: impl Into<Cow<'static, str>>) -> Self {
145 self.metadata.content_type = content_type.into();
146 self
147 }
148
149 pub fn origin(mut self, origin: impl Into<String>) -> Self {
167 self.metadata.origin = Some(origin.into());
168 self
169 }
170
171 pub fn filename(mut self, filename: impl Into<String>) -> Self {
177 self.metadata.filename = Some(filename.into());
178 self
179 }
180
181 pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
185 self.metadata.custom = metadata.into();
186 self
187 }
188
189 pub fn append_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
191 self.metadata.custom.insert(key.into(), value.into());
192 self
193 }
194}
195
196pub(crate) async fn maybe_compress(
198 body: PutBody,
199 compression: Option<Compression>,
200) -> io::Result<Body> {
201 Ok(match (compression, body) {
202 (Some(Compression::Zstd), PutBody::Buffer(bytes)) => {
203 let cursor = Cursor::new(bytes);
204 let encoder = ZstdEncoder::new(cursor);
205 let stream = ReaderStream::new(encoder);
206 Body::wrap_stream(stream)
207 }
208 (Some(Compression::Zstd), PutBody::Stream(stream)) => {
209 let stream = StreamReader::new(stream);
210 let encoder = ZstdEncoder::new(stream);
211 let stream = ReaderStream::new(encoder);
212 Body::wrap_stream(stream)
213 }
214 (Some(Compression::Zstd), PutBody::File(file)) => {
215 let reader = BufReader::new(file);
216 let encoder = ZstdEncoder::new(reader);
217 let stream = ReaderStream::new(encoder);
218 Body::wrap_stream(stream)
219 }
220 (Some(Compression::Zstd), PutBody::Path(file)) => {
221 let file = File::open(file).await?;
222 let reader = BufReader::new(file);
223 let encoder = ZstdEncoder::new(reader);
224 let stream = ReaderStream::new(encoder);
225 Body::wrap_stream(stream)
226 }
227 (None, PutBody::Buffer(bytes)) => bytes.into(),
228 (None, PutBody::Stream(stream)) => Body::wrap_stream(stream),
229 (None, PutBody::File(file)) => {
230 let stream = ReaderStream::new(file);
231 Body::wrap_stream(stream)
232 }
233 (None, PutBody::Path(path)) => {
234 let stream = ReaderStream::new(File::open(path).await?);
235 Body::wrap_stream(stream)
236 }
237 })
238}
239
240impl PutBuilder {
244 pub async fn send(self) -> crate::Result<PutResponse> {
246 let method = match self.key {
247 Some(_) => reqwest::Method::PUT,
248 None => reqwest::Method::POST,
249 };
250
251 let mut builder = self
252 .session
253 .request(method, self.key.as_deref().unwrap_or_default())?;
254
255 let body = maybe_compress(self.body, self.metadata.compression).await?;
256
257 builder = builder.headers(self.metadata.to_headers("")?);
258
259 let response = builder.body(body).send().await?;
260 Ok(response.error_for_status()?.json().await?)
261 }
262}