Add support for writing object metadata with PutObject (#1062)

* Add support for writing object metadata with PutObject Signed-off-by: Simon Beal <[email protected]> * Make changes from code review Signed-off-by: Simon Beal <[email protected]> * Fix merge conflicts Signed-off-by: Simon Beal <[email protected]> --------- Signed-off-by: Simon Beal <[email protected]>
awslabs · Oct 16, 2024 · 6a8a483 · 6a8a483
1 parent e98a5c2
commit 6a8a483
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 3 deletions.
diff --git a/mountpoint-s3-client/src/mock_client.rs b/mountpoint-s3-client/src/mock_client.rs
@@ -373,6 +373,7 @@ pub struct MockObject {
     last_modified: OffsetDateTime,
     etag: ETag,
     parts: Option<MockObjectParts>,
+    object_metadata: HashMap<String, String>,
 }
 
 impl MockObject {
@@ -391,6 +392,7 @@ impl MockObject {
             last_modified: OffsetDateTime::now_utc(),
             etag,
             parts: None,
+            object_metadata: HashMap::new(),
         }
     }
 
@@ -403,6 +405,7 @@ impl MockObject {
             last_modified: OffsetDateTime::now_utc(),
             etag,
             parts: None,
+            object_metadata: HashMap::new(),
         }
     }
 
@@ -425,6 +428,7 @@ impl MockObject {
             last_modified: OffsetDateTime::now_utc(),
             etag,
             parts: None,
+            object_metadata: HashMap::new(),
         }
     }
 
@@ -436,6 +440,10 @@ impl MockObject {
         self.storage_class = storage_class;
     }
 
+    pub fn set_object_metadata(&mut self, object_metadata: HashMap<String, String>) {
+        self.object_metadata = object_metadata;
+    }
+
     pub fn set_restored(&mut self, restore_status: Option<RestoreStatus>) {
         self.restore_status = restore_status;
     }
@@ -731,6 +739,7 @@ impl ObjectClient for MockClient {
 
         let mut object: MockObject = contents.into();
         object.set_storage_class(params.storage_class.clone());
+        object.set_object_metadata(params.object_metadata.clone());
         let etag = object.etag.clone();
         add_object(&self.objects, key, object);
         Ok(PutObjectResult {
@@ -870,6 +879,7 @@ impl MockPutObjectRequest {
         let buffer = std::mem::take(&mut self.buffer);
         let mut object: MockObject = buffer.into();
         object.set_storage_class(self.params.storage_class.clone());
+        object.set_object_metadata(self.params.object_metadata.clone());
         // For S3 Standard, part attributes are only available when additional checksums are used
         if self.params.trailing_checksums == PutObjectTrailingChecksums::Enabled {
             object.parts = Some(MockObjectParts::Parts(parts));
@@ -1505,8 +1515,10 @@ mod tests {
             ..Default::default()
         });
 
+        let object_metadata = HashMap::from([("foo".to_string(), "bar".to_string())]);
+        let put_object_params = PutObjectParams::new().object_metadata(object_metadata.clone());
         let mut put_request = client
-            .put_object("test_bucket", "key1", &Default::default())
+            .put_object("test_bucket", "key1", &put_object_params)
             .await
             .expect("put_object failed");
 
@@ -1534,6 +1546,7 @@ mod tests {
             next_offset += body.len() as u64;
             assert_eq!(body, obj.read(offset, body.len()));
         }
+        assert_eq!(object_metadata, get_request.object.object_metadata);
     }
 
     #[tokio::test]
@@ -1546,8 +1559,10 @@ mod tests {
         });
 
         let content = vec![42u8; 512];
+        let object_metadata = HashMap::from([("foo".to_string(), "bar".to_string())]);
+        let put_object_params = PutObjectSingleParams::new().object_metadata(object_metadata.clone());
         let _put_result = client
-            .put_object_single("test_bucket", "key1", &Default::default(), &content)
+            .put_object_single("test_bucket", "key1", &put_object_params, &content)
             .await
             .expect("put_object failed");
 
@@ -1556,6 +1571,7 @@ mod tests {
             .await
             .expect("get_object failed");
 
+        assert_eq!(object_metadata, get_request.object.object_metadata);
         // Check that the result of get_object is correct.
         let actual = get_request.collect().await.expect("failed to collect body");
         assert_eq!(&content, &*actual);

diff --git a/mountpoint-s3-client/src/object_client.rs b/mountpoint-s3-client/src/object_client.rs
@@ -7,6 +7,7 @@ use async_trait::async_trait;
 use auto_impl::auto_impl;
 use futures::Stream;
 use mountpoint_s3_crt::s3::client::BufferPoolUsageStats;
+use std::collections::HashMap;
 use thiserror::Error;
 use time::OffsetDateTime;
 
@@ -270,6 +271,8 @@ pub struct PutObjectParams {
     pub ssekms_key_id: Option<String>,
     /// Custom headers to add to the request
     pub custom_headers: Vec<(String, String)>,
+    /// User-defined object metadata
+    pub object_metadata: HashMap<String, String>,
 }
 
 impl PutObjectParams {
@@ -307,6 +310,12 @@ impl PutObjectParams {
         self.custom_headers.push((name, value));
         self
     }
+
+    /// Set user defined object metadata.
+    pub fn object_metadata(mut self, value: HashMap<String, String>) -> Self {
+        self.object_metadata = value;
+        self
+    }
 }
 
 /// How CRC32c checksums are used for parts of a multi-part PutObject request
@@ -345,6 +354,8 @@ pub struct PutObjectSingleParams {
     pub ssekms_key_id: Option<String>,
     /// Custom headers to add to the request
     pub custom_headers: Vec<(String, String)>,
+    /// User-defined object metadata
+    pub object_metadata: HashMap<String, String>,
 }
 
 impl PutObjectSingleParams {
@@ -382,6 +393,12 @@ impl PutObjectSingleParams {
         self.custom_headers.push((name, value));
         self
     }
+
+    /// Set user defined object metadata.
+    pub fn object_metadata(mut self, value: HashMap<String, String>) -> Self {
+        self.object_metadata = value;
+        self
+    }
 }
 
 /// A checksum used by the object client for integrity checks on uploads.

diff --git a/mountpoint-s3-client/src/s3_crt_client/put_object.rs b/mountpoint-s3-client/src/s3_crt_client/put_object.rs
@@ -45,6 +45,11 @@ impl S3CrtClient {
         };
         message.set_checksum_config(checksum_config);
 
+        for (name, value) in &params.object_metadata {
+            message
+                .set_header(&Header::new(format!("x-amz-meta-{}", name), value))
+                .map_err(S3RequestError::construction_failure)?
+        }
         for (name, value) in &params.custom_headers {
             message
                 .inner
@@ -130,6 +135,11 @@ impl S3CrtClient {
                     .set_checksum_header(checksum)
                     .map_err(S3RequestError::construction_failure)?;
             }
+            for (name, value) in &params.object_metadata {
+                message
+                    .set_header(&Header::new(format!("x-amz-meta-{}", name), value))
+                    .map_err(S3RequestError::construction_failure)?
+            }
             for (name, value) in &params.custom_headers {
                 message
                     .inner

diff --git a/mountpoint-s3-client/tests/put_object.rs b/mountpoint-s3-client/tests/put_object.rs
@@ -2,6 +2,7 @@
 
 pub mod common;
 
+use std::collections::HashMap;
 use std::time::Duration;
 
 use common::*;
@@ -375,6 +376,52 @@ async fn test_put_checksums(trailing_checksums: PutObjectTrailingChecksums) {
     }
 }
 
+#[test_case(HashMap::new(); "Empty")]
+#[test_case(HashMap::from([("foo".to_string(), "bar".to_string()), ("a".to_string(), "b".to_string())]); "ASCII")]
+#[tokio::test]
+async fn test_put_user_object_metadata_happy(object_metadata: HashMap<String, String>) {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_user_object_metadata_happy");
+    let client_config = S3ClientConfig::new().endpoint_config(EndpointConfig::new(&get_test_region()));
+    let client = S3CrtClient::new(client_config).expect("could not create test client");
+    let key = format!("{prefix}hello");
+
+    let params = PutObjectParams::new().object_metadata(object_metadata.clone());
+
+    let mut request = client
+        .put_object(&bucket, &key, &params)
+        .await
+        .expect("put_object should succeed");
+
+    request.write(b"data").await.unwrap();
+    request.complete().await.unwrap();
+
+    let sdk_client = get_test_sdk_client().await;
+    let output = sdk_client.head_object().bucket(&bucket).key(key).send().await.unwrap();
+
+    match output.metadata() {
+        Some(returned_object_metadata) => {
+            assert_eq!(&object_metadata, returned_object_metadata);
+        }
+        None => {
+            assert!(object_metadata.is_empty());
+        }
+    }
+}
+
+#[test_case(HashMap::from([("£".to_string(), "£".to_string())]); "UTF-8")]
+#[tokio::test]
+async fn test_put_user_object_metadata_bad_header(object_metadata: HashMap<String, String>) {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_user_object_metadata_bad_header");
+    let client_config = S3ClientConfig::new().endpoint_config(EndpointConfig::new(&get_test_region()));
+    let client = S3CrtClient::new(client_config).expect("could not create test client");
+    let key = format!("{prefix}hello");
+
+    let params = PutObjectParams::new().object_metadata(object_metadata.clone());
+
+    let mut request = client.put_object(&bucket, &key, &params).await.unwrap();
+    request.write(b"data").await.expect_err("header parsing should fail");
+}
+
 #[test_case(true; "pass review")]
 #[test_case(false; "fail review")]
 #[tokio::test]
@@ -450,7 +497,7 @@ async fn check_get_object<Client: ObjectClient>(
 // S3 Express One Zone is a distinct storage class and can't be overridden
 #[cfg(not(feature = "s3express_tests"))]
 async fn test_put_object_storage_class(storage_class: &str) {
-    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_object_abort");
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_object_storage_class");
     let client = get_test_client();
     let key = format!("{prefix}hello");
 

diff --git a/mountpoint-s3-client/tests/put_object_single.rs b/mountpoint-s3-client/tests/put_object_single.rs
@@ -2,6 +2,8 @@
 
 pub mod common;
 
+use std::collections::HashMap;
+
 use common::*;
 use mountpoint_s3_client::checksums::{crc32c, crc32c_to_base64};
 use mountpoint_s3_client::config::{EndpointConfig, S3ClientConfig};
@@ -117,6 +119,49 @@ async fn test_put_checksums(checksum_algorithm: Option<ChecksumAlgorithm>) {
     }
 }
 
+#[test_case(HashMap::new(); "Empty")]
+#[test_case(HashMap::from([("foo".to_string(), "bar".to_string()), ("a".to_string(), "b".to_string())]); "ASCII")]
+#[tokio::test]
+async fn test_put_user_object_metadata_happy(object_metadata: HashMap<String, String>) {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_user_object_metadata_happy");
+    let client_config = S3ClientConfig::new().endpoint_config(EndpointConfig::new(&get_test_region()));
+    let client = S3CrtClient::new(client_config).expect("could not create test client");
+    let key = format!("{prefix}hello");
+
+    let params = PutObjectSingleParams::new().object_metadata(object_metadata.clone());
+    client
+        .put_object_single(&bucket, &key, &params, b"data")
+        .await
+        .expect("put_object should succeed");
+
+    let sdk_client = get_test_sdk_client().await;
+    let output = sdk_client.head_object().bucket(&bucket).key(key).send().await.unwrap();
+
+    match output.metadata() {
+        Some(returned_object_metadata) => {
+            assert_eq!(&object_metadata, returned_object_metadata);
+        }
+        None => {
+            assert!(object_metadata.is_empty());
+        }
+    }
+}
+
+#[test_case(HashMap::from([("£".to_string(), "£".to_string())]); "UTF-8")]
+#[tokio::test]
+async fn test_put_user_object_metadata_bad_header(object_metadata: HashMap<String, String>) {
+    let (bucket, prefix) = get_test_bucket_and_prefix("test_put_user_object_metadata_bad_header");
+    let client_config = S3ClientConfig::new().endpoint_config(EndpointConfig::new(&get_test_region()));
+    let client = S3CrtClient::new(client_config).expect("could not create test client");
+    let key = format!("{prefix}hello");
+
+    let params = PutObjectSingleParams::new().object_metadata(object_metadata.clone());
+    client
+        .put_object_single(&bucket, &key, &params, b"data")
+        .await
+        .expect_err("header parsing should fail");
+}
+
 #[test_case("INTELLIGENT_TIERING")]
 #[test_case("GLACIER")]
 #[tokio::test]