1use 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#[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 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 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 pub fn key(&self) -> &str {
101 &self.key
102 }
103
104 pub fn token(&self) -> &SessionToken {
106 &self.token
107 }
108
109 pub fn granularity(&self) -> Option<u64> {
114 self.granularity
115 }
116
117 pub fn progress(&self) -> UploadProgressBuilder {
119 UploadProgressBuilder {
120 upload: self.clone(),
121 }
122 }
123
124 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 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 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 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 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#[derive(Debug)]
206pub struct CreateResumableUploadBuilder {
207 session: Session,
208 total_length: u64,
209 key: Option<ObjectKey>,
210 metadata: Metadata,
211}
212
213impl CreateResumableUploadBuilder {
214 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 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 pub fn expiration_policy(mut self, expiration_policy: ExpirationPolicy) -> Self {
228 self.metadata.expiration_policy = expiration_policy;
229 self
230 }
231
232 pub fn compression(mut self, compression: impl Into<Option<Compression>>) -> Self {
240 self.metadata.compression = compression.into();
241 self
242 }
243
244 pub fn origin(mut self, origin: impl Into<String>) -> Self {
246 self.metadata.origin = Some(origin.into());
247 self
248 }
249
250 pub fn filename(mut self, filename: impl Into<String>) -> Self {
252 self.metadata.filename = Some(filename.into());
253 self
254 }
255
256 pub fn set_metadata(mut self, metadata: impl Into<BTreeMap<String, String>>) -> Self {
258 self.metadata.custom = metadata.into();
259 self
260 }
261
262 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 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#[derive(Debug)]
317pub struct UploadProgressBuilder {
318 upload: ResumableUpload,
319}
320
321impl UploadProgressBuilder {
322 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
339pub 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 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#[derive(Debug)]
396pub struct CancelUploadBuilder {
397 upload: ResumableUpload,
398}
399
400impl CancelUploadBuilder {
401 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}