(arrow-rs) branch master updated: Fix PutPayloadMut::push not updating content_length (#5743) (#5744)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 02be02bb51f Fix PutPayloadMut::push not updating content_length 
(#5743) (#5744)
02be02bb51f is described below

commit 02be02bb51f5c4db5a75419297c4b07b098cf6ba
Author: Raphael Taylor-Davies <1781103+tustv...@users.noreply.github.com>
AuthorDate: Fri May 10 17:59:21 2024 +0100

    Fix PutPayloadMut::push not updating content_length (#5743) (#5744)
---
 object_store/src/integration.rs | 10 ++++++++++
 object_store/src/payload.rs     | 16 +++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/object_store/src/integration.rs b/object_store/src/integration.rs
index 9a7d117158c..d08c4509f36 100644
--- a/object_store/src/integration.rs
+++ b/object_store/src/integration.rs
@@ -789,6 +789,16 @@ pub async fn stream_get(storage: &DynObjectStore) {
     let bytes_written = 
storage.get(&location).await.unwrap().bytes().await.unwrap();
     assert_eq!(bytes_expected, bytes_written);
 
+    let location = Path::from("test_dir/test_put_part.txt");
+    let upload = storage.put_multipart(&location).await.unwrap();
+    let mut write = WriteMultipart::new(upload);
+    write.put(vec![0; 2].into());
+    write.put(vec![3; 4].into());
+    write.finish().await.unwrap();
+
+    let meta = storage.head(&location).await.unwrap();
+    assert_eq!(meta.size, 6);
+
     // We can abort an empty write
     let location = Path::from("test_dir/test_abort_upload.txt");
     let mut upload = storage.put_multipart(&location).await.unwrap();
diff --git a/object_store/src/payload.rs b/object_store/src/payload.rs
index 486bea3ea91..d71f016bcd0 100644
--- a/object_store/src/payload.rs
+++ b/object_store/src/payload.rs
@@ -252,7 +252,8 @@ impl PutPayloadMut {
             let completed = std::mem::take(&mut self.in_progress);
             self.completed.push(completed.into())
         }
-        self.completed.push(bytes)
+        self.len += bytes.len();
+        self.completed.push(bytes);
     }
 
     /// Returns `true` if this [`PutPayloadMut`] contains no bytes
@@ -311,4 +312,17 @@ mod test {
         assert_eq!(chunks[4].len(), 20);
         assert_eq!(chunks[5].len(), 6);
     }
+
+    #[test]
+    fn test_content_length() {
+        let mut chunk = PutPayloadMut::new();
+        chunk.push(vec![0; 23].into());
+        assert_eq!(chunk.content_length(), 23);
+        chunk.extend_from_slice(&[0; 4]);
+        assert_eq!(chunk.content_length(), 27);
+        chunk.push(vec![0; 121].into());
+        assert_eq!(chunk.content_length(), 148);
+        let payload = chunk.freeze();
+        assert_eq!(payload.content_length(), 148);
+    }
 }