Conditional Put (apache#4879) (apache#4984)

* Add version to PutResult * Conditional Put (apache#4879) * Don't support HttpStore * Add R2 Support * Update Azure StatusCode * Fixes * Clippy * Clippy * PutRequestBuilder * Clippy * Add stress test * Clippy
ritchie46 · Oct 27, 2023 · e3cce56 · e3cce56
1 parent 5601b7a
commit e3cce56
Show file tree

Hide file tree

Showing 22 changed files with 791 additions and 302 deletions.
diff --git a/object_store/src/aws/builder.rs b/object_store/src/aws/builder.rs
@@ -20,7 +20,8 @@ use crate::aws::credential::{
     InstanceCredentialProvider, TaskCredentialProvider, WebIdentityProvider,
 };
 use crate::aws::{
-    AmazonS3, AwsCredential, AwsCredentialProvider, Checksum, S3CopyIfNotExists, STORE,
+    AmazonS3, AwsCredential, AwsCredentialProvider, Checksum, S3ConditionalPut, S3CopyIfNotExists,
+    STORE,
 };
 use crate::client::TokenCredentialProvider;
 use crate::config::ConfigValue;
@@ -152,6 +153,8 @@ pub struct AmazonS3Builder {
     skip_signature: ConfigValue<bool>,
     /// Copy if not exists
     copy_if_not_exists: Option<ConfigValue<S3CopyIfNotExists>>,
+    /// Put precondition
+    conditional_put: Option<ConfigValue<S3ConditionalPut>>,
 }
 
 /// Configuration keys for [`AmazonS3Builder`]
@@ -288,6 +291,11 @@ pub enum AmazonS3ConfigKey {
     /// See [`S3CopyIfNotExists`]
     CopyIfNotExists,
 
+    /// Configure how to provide conditional put operations
+    ///
+    /// See [`S3ConditionalPut`]
+    ConditionalPut,
+
     /// Skip signing request
     SkipSignature,
 
@@ -312,7 +320,8 @@ impl AsRef<str> for AmazonS3ConfigKey {
             Self::Checksum => "aws_checksum_algorithm",
             Self::ContainerCredentialsRelativeUri => "aws_container_credentials_relative_uri",
             Self::SkipSignature => "aws_skip_signature",
-            Self::CopyIfNotExists => "copy_if_not_exists",
+            Self::CopyIfNotExists => "aws_copy_if_not_exists",
+            Self::ConditionalPut => "aws_conditional_put",
             Self::Client(opt) => opt.as_ref(),
         }
     }
@@ -339,7 +348,8 @@ impl FromStr for AmazonS3ConfigKey {
             "aws_checksum_algorithm" | "checksum_algorithm" => Ok(Self::Checksum),
             "aws_container_credentials_relative_uri" => Ok(Self::ContainerCredentialsRelativeUri),
             "aws_skip_signature" | "skip_signature" => Ok(Self::SkipSignature),
-            "copy_if_not_exists" => Ok(Self::CopyIfNotExists),
+            "aws_copy_if_not_exists" | "copy_if_not_exists" => Ok(Self::CopyIfNotExists),
+            "aws_conditional_put" | "conditional_put" => Ok(Self::ConditionalPut),
             // Backwards compatibility
             "aws_allow_http" => Ok(Self::Client(ClientConfigKey::AllowHttp)),
             _ => match s.parse() {
@@ -446,6 +456,9 @@ impl AmazonS3Builder {
             AmazonS3ConfigKey::CopyIfNotExists => {
                 self.copy_if_not_exists = Some(ConfigValue::Deferred(value.into()))
             }
+            AmazonS3ConfigKey::ConditionalPut => {
+                self.conditional_put = Some(ConfigValue::Deferred(value.into()))
+            }
         };
         self
     }
@@ -509,6 +522,9 @@ impl AmazonS3Builder {
             AmazonS3ConfigKey::CopyIfNotExists => {
                 self.copy_if_not_exists.as_ref().map(ToString::to_string)
             }
+            AmazonS3ConfigKey::ConditionalPut => {
+                self.conditional_put.as_ref().map(ToString::to_string)
+            }
         }
     }
 
@@ -713,6 +729,12 @@ impl AmazonS3Builder {
         self
     }
 
+    /// Configure how to provide conditional put operations
+    pub fn with_conditional_put(mut self, config: S3ConditionalPut) -> Self {
+        self.conditional_put = Some(config.into());
+        self
+    }
+
     /// Create a [`AmazonS3`] instance from the provided values,
     /// consuming `self`.
     pub fn build(mut self) -> Result<AmazonS3> {
@@ -724,6 +746,7 @@ impl AmazonS3Builder {
         let region = self.region.context(MissingRegionSnafu)?;
         let checksum = self.checksum_algorithm.map(|x| x.get()).transpose()?;
         let copy_if_not_exists = self.copy_if_not_exists.map(|x| x.get()).transpose()?;
+        let put_precondition = self.conditional_put.map(|x| x.get()).transpose()?;
 
         let credentials = if let Some(credentials) = self.credentials {
             credentials
@@ -830,6 +853,7 @@ impl AmazonS3Builder {
             skip_signature: self.skip_signature.get()?,
             checksum,
             copy_if_not_exists,
+            conditional_put: put_precondition,
         };
 
         let client = Arc::new(S3Client::new(config)?);