Skip to main content

polars_io/cloud/cloud_writer/
multipart_upload.rs

1use polars_error::PolarsResult;
2
3use crate::cloud::ObjectStoreErrorContext;
4
5/// Wrapper for [`object_store::MultipartUpload`] that handles error conversion.
6pub struct PlMultipartUpload {
7    inner: Box<dyn object_store::MultipartUpload>,
8    error_cx: ObjectStoreErrorContext,
9}
10
11impl PlMultipartUpload {
12    pub fn new(
13        inner: Box<dyn object_store::MultipartUpload>,
14        error_cx: ObjectStoreErrorContext,
15    ) -> Self {
16        Self { inner, error_cx }
17    }
18
19    pub fn put(
20        &mut self,
21        payload: object_store::PutPayload,
22    ) -> impl Future<Output = PolarsResult<()>> + Send + 'static {
23        let fut = self.inner.put_part(payload);
24        let error_cx = self.error_cx.clone();
25
26        async move { fut.await.map_err(|e| error_cx.attach_err_info(e).into()) }
27    }
28
29    pub async fn finish(&mut self) -> PolarsResult<object_store::PutResult> {
30        self.inner
31            .complete()
32            .await
33            .map_err(|e| self.error_cx.clone().attach_err_info(e).into())
34    }
35}