added multipart uploads for other object storage sources

Author: parmesant

src/storage/azure_blob.rs

@@ -41,6 +41,7 @@ use object_store::{
     BackoffConfig, ClientOptions, ObjectMeta, ObjectStore, PutPayload, RetryConfig,
 };
 use relative_path::{RelativePath, RelativePathBuf};
+use tokio::{fs::OpenOptions, io::AsyncReadExt};
 use tracing::{error, info};
 use url::Url;
 
@@ -53,8 +54,8 @@ use crate::{
 use super::{
     metrics_layer::MetricLayer, object_storage::parseable_json_path, to_object_store_path,
     ObjectStorage, ObjectStorageError, ObjectStorageProvider, CONNECT_TIMEOUT_SECS,
-    PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, SCHEMA_FILE_NAME, STREAM_METADATA_FILE_NAME,
-    STREAM_ROOT_DIRECTORY,
+    MIN_MULTIPART_UPLOAD_SIZE, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, SCHEMA_FILE_NAME,
+    STREAM_METADATA_FILE_NAME, STREAM_ROOT_DIRECTORY,
 };
 
 #[derive(Debug, Clone, clap::Args)]
@@ -378,6 +379,65 @@ impl BlobStore {
         res
     }
 
+    async fn _upload_multipart(
+        &self,
+        key: &RelativePath,
+        path: &Path,
+    ) -> Result<(), ObjectStorageError> {
+        let mut file = OpenOptions::new().read(true).open(path).await?;
+        let location = &to_object_store_path(key);
+
+        let mut async_writer = self.client.put_multipart(location).await?;
+
+        let meta = file.metadata().await?;
+        let total_size = meta.len() as usize;
+        if total_size < MIN_MULTIPART_UPLOAD_SIZE {
+            let mut data = Vec::new();
+            file.read_to_end(&mut data).await?;
+            self.client.put(location, data.into()).await?;
+            // async_writer.put_part(data.into()).await?;
+            // async_writer.complete().await?;
+            return Ok(());
+        } else {
+            let mut data = Vec::new();
+            file.read_to_end(&mut data).await?;
+
+            // let mut upload_parts = Vec::new();
+
+            let has_final_partial_part = total_size % MIN_MULTIPART_UPLOAD_SIZE > 0;
+            let num_full_parts = total_size / MIN_MULTIPART_UPLOAD_SIZE;
+            let total_parts = num_full_parts + if has_final_partial_part { 1 } else { 0 };
+
+            // Upload each part
+            for part_number in 0..(total_parts) {
+                let start_pos = part_number * MIN_MULTIPART_UPLOAD_SIZE;
+                let end_pos = if part_number == num_full_parts && has_final_partial_part {
+                    // Last part might be smaller than 5MB (which is allowed)
+                    total_size
+                } else {
+                    // All other parts must be at least 5MB
+                    start_pos + MIN_MULTIPART_UPLOAD_SIZE
+                };
+
+                // Extract this part's data
+                let part_data = data[start_pos..end_pos].to_vec();
+
+                // Upload the part
+                async_writer.put_part(part_data.into()).await?;
+
+                // upload_parts.push(part_number as u64 + 1);
+            }
+            match async_writer.complete().await {
+                Ok(_) => {}
+                Err(err) => {
+                    error!("Failed to complete multipart upload. {:?}", err);
+                    async_writer.abort().await?;
+                }
+            };
+        }
+        Ok(())
+    }
+
     // TODO: introduce parallel, multipart-uploads if required
     // async fn _upload_multipart(&self, key: &str, path: &Path) -> Result<(), ObjectStorageError> {
     //     let mut buf = vec![0u8; MULTIPART_UPLOAD_SIZE / 2];
@@ -426,10 +486,10 @@ impl BlobStore {
 impl ObjectStorage for BlobStore {
     async fn upload_multipart(
         &self,
-        _key: &RelativePath,
-        _path: &Path,
+        key: &RelativePath,
+        path: &Path,
     ) -> Result<(), ObjectStorageError> {
-        unimplemented!()
+        self._upload_multipart(key, path).await
     }
     async fn get_buffered_reader(
         &self,

src/storage/localfs.rs

@@ -30,7 +30,10 @@ use fs_extra::file::CopyOptions;
 use futures::{stream::FuturesUnordered, TryStreamExt};
 use object_store::{buffered::BufReader, ObjectMeta};
 use relative_path::{RelativePath, RelativePathBuf};
-use tokio::fs::{self, DirEntry};
+use tokio::{
+    fs::{self, DirEntry, OpenOptions},
+    io::AsyncReadExt,
+};
 use tokio_stream::wrappers::ReadDirStream;
 
 use crate::{
@@ -106,10 +109,13 @@ impl LocalFS {
 impl ObjectStorage for LocalFS {
     async fn upload_multipart(
         &self,
-        _key: &RelativePath,
-        _path: &Path,
+        key: &RelativePath,
+        path: &Path,
     ) -> Result<(), ObjectStorageError> {
-        unimplemented!()
+        let mut file = OpenOptions::new().read(true).open(path).await?;
+        let mut data = Vec::new();
+        file.read_to_end(&mut data).await?;
+        self.put_object(key, data.into()).await
     }
     async fn get_buffered_reader(
         &self,

src/storage/mod.rs

@@ -77,6 +77,7 @@ pub const CURRENT_SCHEMA_VERSION: &str = "v6";
 const CONNECT_TIMEOUT_SECS: u64 = 5;
 const REQUEST_TIMEOUT_SECS: u64 = 300;
 
+pub const MIN_MULTIPART_UPLOAD_SIZE: usize = 25 * 1024 * 1024;
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ObjectStoreFormat {
     /// Version of schema registry

src/storage/s3.rs

@@ -55,14 +55,13 @@ use crate::{
 use super::{
     metrics_layer::MetricLayer, object_storage::parseable_json_path, to_object_store_path,
     ObjectStorage, ObjectStorageError, ObjectStorageProvider, CONNECT_TIMEOUT_SECS,
-    PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, SCHEMA_FILE_NAME, STREAM_METADATA_FILE_NAME,
-    STREAM_ROOT_DIRECTORY,
+    MIN_MULTIPART_UPLOAD_SIZE, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, SCHEMA_FILE_NAME,
+    STREAM_METADATA_FILE_NAME, STREAM_ROOT_DIRECTORY,
 };
 
 // in bytes
 // const MULTIPART_UPLOAD_SIZE: usize = 1024 * 1024 * 100;
 const AWS_CONTAINER_CREDENTIALS_RELATIVE_URI: &str = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI";
-const MIN_MULTIPART_UPLOAD_SIZE: usize = 25 * 1024 * 1024;
 
 #[derive(Debug, Clone, clap::Args)]
 #[command(
@@ -560,7 +559,7 @@ impl S3 {
                 // upload_parts.push(part_number as u64 + 1);
             }
             match async_writer.complete().await {
-                Ok(_) => {},
+                Ok(_) => {}
                 Err(err) => {
                     error!("Failed to complete multipart upload. {:?}", err);
                     async_writer.abort().await?;