Support native S3 conditional writes

.github/workflows/object_store.yml

@@ -161,6 +161,12 @@ jobs:
       - name: Run object_store tests
         run: cargo test --features=aws,azure,gcp,http
 
+      - name: Run object_store tests (AWS native conditional put)
+        run: cargo test --features=aws
+        env:
+          AWS_CONDITIONAL_PUT: etag-put-if-not-exists
+          AWS_COPY_IF_NOT_EXISTS: multipart
+
       - name: GCS Output
         if: ${{ !cancelled() }}
         run: docker logs $GCS_CONTAINER

object_store/src/aws/client.rs

@@ -28,8 +28,8 @@ use crate::client::header::{get_put_result, get_version};
 use crate::client::list::ListClient;
 use crate::client::retry::RetryExt;
 use crate::client::s3::{
-    CompleteMultipartUpload, CompleteMultipartUploadResult, InitiateMultipartUploadResult,
-    ListResponse,
+    CompleteMultipartUpload, CompleteMultipartUploadResult, CopyPartResult,
+    InitiateMultipartUploadResult, ListResponse,
 };
 use crate::client::GetOptionsExt;
 use crate::multipart::PartId;
@@ -98,8 +98,11 @@ pub(crate) enum Error {
     #[snafu(display("Error getting create multipart response body: {}", source))]
     CreateMultipartResponseBody { source: reqwest::Error },
 
-    #[snafu(display("Error performing complete multipart request: {}", source))]
-    CompleteMultipartRequest { source: crate::client::retry::Error },
+    #[snafu(display("Error performing complete multipart request: {}: {}", path, source))]
+    CompleteMultipartRequest {
+        source: crate::client::retry::Error,
+        path: String,
+    },
 
     #[snafu(display("Error getting complete multipart response body: {}", source))]
     CompleteMultipartResponseBody { source: reqwest::Error },
@@ -118,13 +121,32 @@ pub(crate) enum Error {
 
 impl From<Error> for crate::Error {
     fn from(err: Error) -> Self {
-        Self::Generic {
-            store: STORE,
-            source: Box::new(err),
+        match err {
+            Error::CompleteMultipartRequest { source, path } => source.error(STORE, path),
+            _ => Self::Generic {
+                store: STORE,
+                source: Box::new(err),
+            },
         }
     }
 }
 
+pub(crate) enum PutPartPayload<'a> {
+    Part(PutPayload),
+    Copy(&'a Path),
+}
+
+impl Default for PutPartPayload<'_> {
+    fn default() -> Self {
+        Self::Part(PutPayload::default())
+    }
+}
+
+pub(crate) enum CompleteMultipartMode {
+    Overwrite,
+    Create,
+}
+
 #[derive(Deserialize)]
 #[serde(rename_all = "PascalCase", rename = "DeleteResult")]
 struct BatchDeleteResponse {
@@ -605,15 +627,24 @@ impl S3Client {
         path: &Path,
         upload_id: &MultipartId,
         part_idx: usize,
-        data: PutPayload,
+        data: PutPartPayload<'_>,
     ) -> Result<PartId> {
+        let is_copy = matches!(data, PutPartPayload::Copy(_));
         let part = (part_idx + 1).to_string();
 
         let mut request = self
             .request(Method::PUT, path)
-            .with_payload(data)
             .query(&[("partNumber", &part), ("uploadId", upload_id)])
             .idempotent(true);
+
+        request = match data {
+            PutPartPayload::Part(payload) => request.with_payload(payload),
+            PutPartPayload::Copy(path) => request.header(
+                "x-amz-copy-source",
+                &format!("{}/{}", self.config.bucket, encode_path(path)),
+            ),
+        };
+
         if self
             .config
             .encryption_headers
@@ -625,21 +656,48 @@ impl S3Client {
         }
         let response = request.send().await?;
 
-        let content_id = get_etag(response.headers()).context(MetadataSnafu)?;
+        let content_id = match is_copy {
+            false => get_etag(response.headers()).context(MetadataSnafu)?,
+            true => {
+                let response = response
+                    .bytes()
+                    .await
+                    .context(CreateMultipartResponseBodySnafu)?;
+                let response: CopyPartResult = quick_xml::de::from_reader(response.reader())
+                    .context(InvalidMultipartResponseSnafu)?;
+                response.e_tag
+            }
+        };
         Ok(PartId { content_id })
     }
 
+    pub(crate) async fn abort_multipart(&self, location: &Path, upload_id: &str) -> Result<()> {
+        self.request(Method::DELETE, location)
+            .query(&[("uploadId", upload_id)])
+            .with_encryption_headers()
+            .send()
+            .await?;
+
+        Ok(())
+    }
+
     pub(crate) async fn complete_multipart(
         &self,
         location: &Path,
         upload_id: &str,
         parts: Vec<PartId>,
+        mode: CompleteMultipartMode,
     ) -> Result<PutResult> {
         let parts = if parts.is_empty() {
             // If no parts were uploaded, upload an empty part
             // otherwise the completion request will fail
             let part = self
-                .put_part(location, &upload_id.to_string(), 0, PutPayload::default())
+                .put_part(
+                    location,
+                    &upload_id.to_string(),
+                    0,
+                    PutPartPayload::default(),
+                )
                 .await?;
             vec![part]
         } else {
@@ -651,18 +709,27 @@ impl S3Client {
         let credential = self.config.get_session_credential().await?;
         let url = self.config.path_url(location);
 
-        let response = self
+        let request = self
             .client
             .request(Method::POST, url)
             .query(&[("uploadId", upload_id)])
             .body(body)
-            .with_aws_sigv4(credential.authorizer(), None)
+            .with_aws_sigv4(credential.authorizer(), None);
+
+        let request = match mode {
+            CompleteMultipartMode::Overwrite => request,
+            CompleteMultipartMode::Create => request.header("If-None-Match", "*"),
+        };
+
+        let response = request
             .retryable(&self.config.retry_config)
             .idempotent(true)
             .retry_error_body(true)
             .send()
             .await
-            .context(CompleteMultipartRequestSnafu)?;
+            .context(CompleteMultipartRequestSnafu {
+                path: location.as_ref(),
+            })?;
 
         let version = get_version(response.headers(), VERSION_HEADER).context(MetadataSnafu)?;
 

object_store/src/aws/mod.rs

@@ -36,7 +36,7 @@ use reqwest::{Method, StatusCode};
 use std::{sync::Arc, time::Duration};
 use url::Url;
 
-use crate::aws::client::{RequestError, S3Client};
+use crate::aws::client::{CompleteMultipartMode, PutPartPayload, RequestError, S3Client};
 use crate::client::get::GetClientExt;
 use crate::client::list::ListClientExt;
 use crate::client::CredentialProvider;
@@ -169,7 +169,10 @@ impl ObjectStore for AmazonS3 {
         match (opts.mode, &self.client.config.conditional_put) {
             (PutMode::Overwrite, _) => request.idempotent(true).do_put().await,
             (PutMode::Create | PutMode::Update(_), None) => Err(Error::NotImplemented),
-            (PutMode::Create, Some(S3ConditionalPut::ETagMatch)) => {
+            (
+                PutMode::Create,
+                Some(S3ConditionalPut::ETagMatch | S3ConditionalPut::ETagPutIfNotExists),
+            ) => {
                 match request.header(&IF_NONE_MATCH, "*").do_put().await {
                     // Technically If-None-Match should return NotModified but some stores,
                     // such as R2, instead return PreconditionFailed
@@ -193,6 +196,7 @@ impl ObjectStore for AmazonS3 {
                     source: "ETag required for conditional put".to_string().into(),
                 })?;
                 match put {
+                    S3ConditionalPut::ETagPutIfNotExists => Err(Error::NotImplemented),
                     S3ConditionalPut::ETagMatch => {
                         request.header(&IF_MATCH, etag.as_str()).do_put().await
                     }
@@ -293,6 +297,47 @@ impl ObjectStore for AmazonS3 {
         let (k, v, status) = match &self.client.config.copy_if_not_exists {
             Some(S3CopyIfNotExists::Header(k, v)) => (k, v, StatusCode::PRECONDITION_FAILED),
             Some(S3CopyIfNotExists::HeaderWithStatus(k, v, status)) => (k, v, *status),
+            Some(S3CopyIfNotExists::Multipart) => {
+                let upload_id = self
+                    .client
+                    .create_multipart(to, PutMultipartOpts::default())
+                    .await?;
+
+                let res = async {
+                    let part_id = self
+                        .client
+                        .put_part(to, &upload_id, 0, PutPartPayload::Copy(from))
+                        .await?;
+                    match self
+                        .client
+                        .complete_multipart(
+                            to,
+                            &upload_id,
+                            vec![part_id],
+                            CompleteMultipartMode::Create,
+                        )
+                        .await
+                    {
+                        Err(e @ Error::Precondition { .. }) => Err(Error::AlreadyExists {
+                            path: to.to_string(),
+                            source: Box::new(e),
+                        }),
+                        Ok(_) => Ok(()),
+                        Err(e) => Err(e),
+                    }
+                }
+                .await;
+
+                // If the multipart upload failed, make a best effort attempt to
+                // clean it up. It's the caller's responsibility to add a
+                // lifecycle rule if guaranteed cleanup is required, as we
+                // cannot protect against an ill-timed process crash.
+                if res.is_err() {
+                    let _ = self.client.abort_multipart(to, &upload_id).await;
+                }
+
+                return res;
+            }
             Some(S3CopyIfNotExists::Dynamo(lock)) => {
                 return lock.copy_if_not_exists(&self.client, from, to).await
             }
@@ -340,7 +385,12 @@ impl MultipartUpload for S3MultiPartUpload {
         Box::pin(async move {
             let part = state
                 .client
-                .put_part(&state.location, &state.upload_id, idx, data)
+                .put_part(
+                    &state.location,
+                    &state.upload_id,
+                    idx,
+                    PutPartPayload::Part(data),
+                )
                 .await?;
             state.parts.put(idx, part);
             Ok(())
@@ -352,7 +402,12 @@ impl MultipartUpload for S3MultiPartUpload {
 
         self.state
             .client
-            .complete_multipart(&self.state.location, &self.state.upload_id, parts)
+            .complete_multipart(
+                &self.state.location,
+                &self.state.upload_id,
+                parts,
+                CompleteMultipartMode::Overwrite,
+            )
             .await
     }
 
@@ -384,7 +439,9 @@ impl MultipartStore for AmazonS3 {
         part_idx: usize,
         data: PutPayload,
     ) -> Result<PartId> {
-        self.client.put_part(path, id, part_idx, data).await
+        self.client
+            .put_part(path, id, part_idx, PutPartPayload::Part(data))
+            .await
     }
 
     async fn complete_multipart(
@@ -393,7 +450,9 @@ impl MultipartStore for AmazonS3 {
         id: &MultipartId,
         parts: Vec<PartId>,
     ) -> Result<PutResult> {
-        self.client.complete_multipart(path, id, parts).await
+        self.client
+            .complete_multipart(path, id, parts, CompleteMultipartMode::Overwrite)
+            .await
     }
 
     async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()> {
@@ -427,7 +486,6 @@ mod tests {
         let integration = config.build().unwrap();
         let config = &integration.client.config;
         let test_not_exists = config.copy_if_not_exists.is_some();
-        let test_conditional_put = config.conditional_put.is_some();
 
         put_get_delete_list(&integration).await;
         get_opts(&integration).await;
@@ -458,8 +516,9 @@ mod tests {
         if test_not_exists {
             copy_if_not_exists(&integration).await;
         }
-        if test_conditional_put {
-            put_opts(&integration, true).await;
+        if let Some(conditional_put) = &config.conditional_put {
+            let supports_update = !matches!(conditional_put, S3ConditionalPut::ETagPutIfNotExists);
+            put_opts(&integration, supports_update).await;
         }
 
         // run integration test with unsigned payload enabled