From 76009ee1e14e9ad4c308fbf2f085d98dcca7505b Mon Sep 17 00:00:00 2001 From: Shawn Chang Date: Thu, 19 Mar 2026 14:53:33 -0700 Subject: [PATCH 01/13] Implement Storage for object-store --- Cargo.lock | 16 ++ Cargo.toml | 1 + crates/storage/object_store/Cargo.toml | 45 ++++ crates/storage/object_store/src/lib.rs | 328 +++++++++++++++++++++++++ crates/storage/object_store/src/s3.rs | 134 ++++++++++ 5 files changed, 524 insertions(+) create mode 100644 crates/storage/object_store/Cargo.toml create mode 100644 crates/storage/object_store/src/lib.rs create mode 100644 crates/storage/object_store/src/s3.rs diff --git a/Cargo.lock b/Cargo.lock index 010f8985eb..3eee2c20d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3984,6 +3984,22 @@ dependencies = [ "tracing", ] +[[package]] +name = "iceberg-storage-object_store" +version = "0.9.0" +dependencies = [ + "async-trait", + "bytes", + "dashmap", + "futures", + "iceberg", + "object_store", + "serde", + "tokio", + "typetag", + "url", +] + [[package]] name = "iceberg-storage-opendal" version = "0.10.1" diff --git a/Cargo.toml b/Cargo.toml index d1083a3dc7..9019efb9d7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -107,6 +107,7 @@ iceberg-catalog-s3tables = { version = "0.10.0", path = "./crates/catalog/s3tabl iceberg-catalog-sql = { version = "0.10.0", path = "./crates/catalog/sql" } iceberg-datafusion = { version = "0.10.0", path = "./crates/integrations/datafusion" } iceberg-property-macro = { version = "0.10.0", path = "./crates/property-macro" } +iceberg-storage-object_store = { version = "0.10.0", path = "./crates/storage/object_store" } iceberg-storage-opendal = { version = "0.10.0", path = "./crates/storage/opendal" } indicatif = "0.18" itertools = "0.13" diff --git a/crates/storage/object_store/Cargo.toml b/crates/storage/object_store/Cargo.toml new file mode 100644 index 0000000000..4a78aacd29 --- /dev/null +++ b/crates/storage/object_store/Cargo.toml @@ -0,0 +1,45 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "iceberg-storage-object_store" +edition = { workspace = true } +version = { workspace = true } +license = { workspace = true } +repository = { workspace = true } + +categories = ["database"] +description = "Apache Iceberg object_store storage implementation" +keywords = ["iceberg", "object-store", "storage", "s3"] + +[features] +default = ["object_store-s3"] +object_store-s3 = ["object_store/aws"] + +[dependencies] +async-trait = { workspace = true } +bytes = { workspace = true } +dashmap = { workspace = true } +futures = { workspace = true } +iceberg = { workspace = true } +object_store = { version = "0.12" } +serde = { workspace = true } +typetag = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs new file mode 100644 index 0000000000..3d6a2748f5 --- /dev/null +++ b/crates/storage/object_store/src/lib.rs @@ -0,0 +1,328 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `object_store`-based storage implementation for Apache Iceberg. +//! +//! This crate provides [`ObjectStoreStorage`] and [`ObjectStoreStorageFactory`], +//! which implement the [`Storage`](iceberg::io::Storage) and +//! [`StorageFactory`](iceberg::io::StorageFactory) traits from the `iceberg` crate +//! using the [`object_store`](https://docs.rs/object_store) crate as the backend. +//! +//! Currently only S3 storage is supported (via the `object_store-s3` feature flag, +//! enabled by default). + +#[cfg(feature = "object_store-s3")] +mod s3; + +use std::ops::Range; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use dashmap::DashMap; +use futures::StreamExt; +use futures::stream::BoxStream; +#[cfg(feature = "object_store-s3")] +use iceberg::io::S3Config; +use iceberg::io::{ + FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageConfig, + StorageFactory, +}; +use iceberg::{Error, ErrorKind, Result}; +use object_store::path::Path as ObjectStorePath; +use object_store::{ObjectStore, PutPayload, WriteMultipart}; +#[cfg(feature = "object_store-s3")] +use s3::{build_s3_store, parse_s3_url}; +use serde::{Deserialize, Serialize}; + +/// Convert an `object_store::Error` into an `iceberg::Error`. +fn from_object_store_error(e: object_store::Error) -> Error { + Error::new(ErrorKind::Unexpected, "Failure in doing io operation").with_source(e) +} + +/// `object_store`-based storage factory. +/// +/// Use this factory with `FileIOBuilder::new(factory)` to create FileIO instances +/// backed by the `object_store` crate. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum ObjectStoreStorageFactory { + /// S3 storage factory. + #[cfg(feature = "object_store-s3")] + S3, +} + +#[typetag::serde(name = "ObjectStoreStorageFactory")] +impl StorageFactory for ObjectStoreStorageFactory { + #[allow(unused_variables)] + fn build(&self, config: &StorageConfig) -> Result> { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorageFactory::S3 => { + let s3_config = S3Config::try_from(config)?; + Ok(Arc::new(ObjectStoreStorage::S3 { + config: Arc::new(s3_config), + store_cache: Arc::new(DashMap::new()), + })) + } + } + } +} + +/// `object_store`-based storage implementation. +/// +/// Stores are cached per bucket to avoid rebuilding the client on every operation. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum ObjectStoreStorage { + /// S3 storage variant. + #[cfg(feature = "object_store-s3")] + S3 { + /// Parsed S3 configuration from iceberg core. + config: Arc, + /// Per-bucket store cache. + #[serde(skip, default)] + store_cache: Arc>>, + }, +} + +impl ObjectStoreStorage { + /// Get or create a cached store and extract the relative `ObjectStorePath`. + fn get_store_and_path(&self, path: &str) -> Result<(Arc, ObjectStorePath)> { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3 { + config, + store_cache, + } => { + let (_scheme, bucket, relative) = parse_s3_url(path)?; + + let store = store_cache + .entry(bucket.to_string()) + .or_try_insert_with(|| build_s3_store(config, bucket))? + .value() + .clone(); + + Ok((store, ObjectStorePath::from(relative))) + } + } + } +} + +#[typetag::serde(name = "ObjectStoreStorage")] +#[async_trait] +impl Storage for ObjectStoreStorage { + async fn exists(&self, path: &str) -> Result { + let (store, object_path) = self.get_store_and_path(path)?; + match store.head(&object_path).await { + Ok(_) => Ok(true), + Err(object_store::Error::NotFound { .. }) => Ok(false), + Err(e) => Err(from_object_store_error(e)), + } + } + + async fn metadata(&self, path: &str) -> Result { + let (store, object_path) = self.get_store_and_path(path)?; + let meta = store + .head(&object_path) + .await + .map_err(from_object_store_error)?; + Ok(FileMetadata { + size: meta.size as u64, + }) + } + + async fn read(&self, path: &str) -> Result { + let (store, object_path) = self.get_store_and_path(path)?; + let result = store + .get(&object_path) + .await + .map_err(from_object_store_error)?; + result.bytes().await.map_err(from_object_store_error) + } + + async fn reader(&self, path: &str) -> Result> { + let (store, object_path) = self.get_store_and_path(path)?; + Ok(Box::new(ObjectStoreReader { + store, + path: object_path, + })) + } + + async fn write(&self, path: &str, bs: Bytes) -> Result<()> { + let (store, object_path) = self.get_store_and_path(path)?; + store + .put(&object_path, PutPayload::from_bytes(bs)) + .await + .map_err(from_object_store_error)?; + Ok(()) + } + + async fn writer(&self, path: &str) -> Result> { + let (store, object_path) = self.get_store_and_path(path)?; + let upload = store + .put_multipart(&object_path) + .await + .map_err(from_object_store_error)?; + let writer = WriteMultipart::new(upload); + Ok(Box::new(ObjectStoreWriter { + writer: Some(writer), + })) + } + + async fn delete(&self, path: &str) -> Result<()> { + let (store, object_path) = self.get_store_and_path(path)?; + store + .delete(&object_path) + .await + .map_err(from_object_store_error)?; + Ok(()) + } + + async fn delete_prefix(&self, path: &str) -> Result<()> { + let (store, object_path) = self.get_store_and_path(path)?; + let prefix = if object_path.as_ref().ends_with('/') { + object_path + } else { + ObjectStorePath::from(format!("{}/", object_path.as_ref())) + }; + + let mut list_stream = store.list(Some(&prefix)); + while let Some(entry) = list_stream.next().await { + let entry = entry.map_err(from_object_store_error)?; + store + .delete(&entry.location) + .await + .map_err(from_object_store_error)?; + } + Ok(()) + } + + async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> { + while let Some(path) = paths.next().await { + let (store, object_path) = self.get_store_and_path(&path)?; + store + .delete(&object_path) + .await + .map_err(from_object_store_error)?; + } + Ok(()) + } + + fn new_input(&self, path: &str) -> Result { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> Result { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } +} + +/// Reader that implements `FileRead` using `object_store`. +struct ObjectStoreReader { + store: Arc, + path: ObjectStorePath, +} + +#[async_trait] +impl FileRead for ObjectStoreReader { + async fn read(&self, range: Range) -> Result { + let opts = object_store::GetOptions { + range: Some((range.start..range.end).into()), + ..Default::default() + }; + let result = self + .store + .get_opts(&self.path, opts) + .await + .map_err(from_object_store_error)?; + result.bytes().await.map_err(from_object_store_error) + } +} + +/// Writer that implements `FileWrite` using `object_store` multipart upload. +struct ObjectStoreWriter { + writer: Option, +} + +#[async_trait] +impl FileWrite for ObjectStoreWriter { + async fn write(&mut self, bs: Bytes) -> Result<()> { + let writer = self + .writer + .as_mut() + .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Writer has already been closed"))?; + writer.write(&bs); + Ok(()) + } + + async fn close(&mut self) -> Result<()> { + let writer = self + .writer + .take() + .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Writer has already been closed"))?; + writer.finish().await.map_err(from_object_store_error)?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(feature = "object_store-s3")] + fn make_s3_storage() -> ObjectStoreStorage { + ObjectStoreStorage::S3 { + config: Arc::new(S3Config::default()), + store_cache: Arc::new(DashMap::new()), + } + } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_store_cache_reuses_store() { + let storage = make_s3_storage(); + let (store1, _) = storage + .get_store_and_path("s3://test-bucket/file1.parquet") + .unwrap(); + let (store2, _) = storage + .get_store_and_path("s3://test-bucket/file2.parquet") + .unwrap(); + assert!(Arc::ptr_eq(&store1, &store2)); + } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_store_cache_different_buckets() { + let storage = make_s3_storage(); + let (store1, _) = storage + .get_store_and_path("s3://bucket-a/file.parquet") + .unwrap(); + let (store2, _) = storage + .get_store_and_path("s3://bucket-b/file.parquet") + .unwrap(); + assert!(!Arc::ptr_eq(&store1, &store2)); + } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_relative_path_extraction() { + let storage = make_s3_storage(); + let (_, path) = storage + .get_store_and_path("s3://my-bucket/data/file.parquet") + .unwrap(); + assert_eq!(path.as_ref(), "data/file.parquet"); + } +} diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs new file mode 100644 index 0000000000..8ff03ded67 --- /dev/null +++ b/crates/storage/object_store/src/s3.rs @@ -0,0 +1,134 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use iceberg::io::S3Config; +use iceberg::{Error, ErrorKind, Result}; +use object_store::ObjectStore; +use object_store::aws::AmazonS3Builder; +use url::Url; + +/// Parse an absolute S3 URL into (scheme, bucket, relative_path). +/// +/// Accepts `s3://` and `s3a://` schemes. +pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { + let url = Url::parse(path).map_err(|e| { + Error::new(ErrorKind::DataInvalid, format!("Invalid URL: {path}")).with_source(e) + })?; + + let scheme = &path[..url.scheme().len()]; + match scheme { + "s3" | "s3a" => {} + _ => { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Unsupported S3 scheme: {scheme} in url: {path}"), + )); + } + } + + let bucket_str = url.host_str().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid s3 url: {path}, missing bucket"), + ) + })?; + + let prefix_len = scheme.len() + "://".len() + bucket_str.len() + "/".len(); + let relative = if path.len() > prefix_len { + &path[prefix_len..] + } else { + "" + }; + + let bucket_start = scheme.len() + "://".len(); + let bucket = &path[bucket_start..bucket_start + bucket_str.len()]; + + Ok((scheme, bucket, relative)) +} + +/// Build an `AmazonS3` store from iceberg's `S3Config` for a given bucket. +pub(crate) fn build_s3_store(config: &S3Config, bucket: &str) -> Result> { + let mut builder = AmazonS3Builder::new().with_bucket_name(bucket); + + if let Some(ref endpoint) = config.endpoint { + builder = builder.with_endpoint(endpoint); + if endpoint.starts_with("http://") { + builder = builder.with_allow_http(true); + } + } + if let Some(ref access_key_id) = config.access_key_id { + builder = builder.with_access_key_id(access_key_id); + } + if let Some(ref secret_access_key) = config.secret_access_key { + builder = builder.with_secret_access_key(secret_access_key); + } + if let Some(ref session_token) = config.session_token { + builder = builder.with_token(session_token); + } + if let Some(ref region) = config.region { + builder = builder.with_region(region); + } + if config.enable_virtual_host_style { + builder = builder.with_virtual_hosted_style_request(true); + } + if config.allow_anonymous { + builder = builder.with_skip_signature(true); + } + + let store = builder.build().map_err(|e| { + Error::new(ErrorKind::Unexpected, "Failed to build S3 object store").with_source(e) + })?; + Ok(Arc::new(store)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_s3_url() { + let (scheme, bucket, relative) = + parse_s3_url("s3://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(scheme, "s3"); + assert_eq!(bucket, "my-bucket"); + assert_eq!(relative, "path/to/file.parquet"); + } + + #[test] + fn test_parse_s3a_url() { + let (scheme, bucket, relative) = + parse_s3_url("s3a://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(scheme, "s3a"); + assert_eq!(bucket, "my-bucket"); + assert_eq!(relative, "path/to/file.parquet"); + } + + #[test] + fn test_parse_s3_url_unsupported_scheme() { + assert!(parse_s3_url("gs://my-bucket/file.parquet").is_err()); + } + + #[test] + fn test_parse_s3_url_bucket_only() { + let (scheme, bucket, relative) = parse_s3_url("s3://my-bucket/").unwrap(); + assert_eq!(scheme, "s3"); + assert_eq!(bucket, "my-bucket"); + assert_eq!(relative, ""); + } +} From 5e18076219756fb90bc26fe90afc571c08830e2b Mon Sep 17 00:00:00 2001 From: dron Date: Mon, 7 Sep 2026 15:45:46 +0530 Subject: [PATCH 02/13] feat(storage): modernize object_store backend, support s3n, and add concurrent deletes - Hoist `object_store` 0.13 to workspace dependencies to align with DataFusion. - Support `s3n://` scheme alongside `s3://` and `s3a://` in `parse_s3_url`. - Optimize `delete_stream` with `try_for_each_concurrent` instead of sequential loop. - Add unit tests for `s3n://` URL parsing and FileIO/Storage serialization roundtrips. - Wire crate workspace lints and publish flag. --- Cargo.lock | 7 ++-- Cargo.toml | 1 + crates/storage/object_store/Cargo.toml | 7 +++- crates/storage/object_store/src/lib.rs | 58 +++++++++++++++++++++----- crates/storage/object_store/src/s3.rs | 15 +++++-- 5 files changed, 70 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3eee2c20d0..8cf8a714ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1484,7 +1484,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -3986,7 +3986,7 @@ dependencies = [ [[package]] name = "iceberg-storage-object_store" -version = "0.9.0" +version = "0.10.1" dependencies = [ "async-trait", "bytes", @@ -3995,6 +3995,7 @@ dependencies = [ "iceberg", "object_store", "serde", + "serde_json", "tokio", "typetag", "url", @@ -8557,7 +8558,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 9019efb9d7..30911f7948 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -122,6 +122,7 @@ mockito = "1" motore-macros = "0.4.3" murmur3 = "0.5.2" once_cell = "1.20" +object_store = "0.13" opendal = "0.58" ordered-float = "4" parquet = "59.2" diff --git a/crates/storage/object_store/Cargo.toml b/crates/storage/object_store/Cargo.toml index 4a78aacd29..49e6e59d29 100644 --- a/crates/storage/object_store/Cargo.toml +++ b/crates/storage/object_store/Cargo.toml @@ -21,6 +21,7 @@ edition = { workspace = true } version = { workspace = true } license = { workspace = true } repository = { workspace = true } +publish = true categories = ["database"] description = "Apache Iceberg object_store storage implementation" @@ -36,10 +37,14 @@ bytes = { workspace = true } dashmap = { workspace = true } futures = { workspace = true } iceberg = { workspace = true } -object_store = { version = "0.12" } +object_store = { workspace = true } serde = { workspace = true } typetag = { workspace = true } url = { workspace = true } [dev-dependencies] +serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } + +[lints] +workspace = true diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 3d6a2748f5..419f7ca819 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -34,8 +34,8 @@ use std::sync::Arc; use async_trait::async_trait; use bytes::Bytes; use dashmap::DashMap; -use futures::StreamExt; use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; #[cfg(feature = "object_store-s3")] use iceberg::io::S3Config; use iceberg::io::{ @@ -44,7 +44,7 @@ use iceberg::io::{ }; use iceberg::{Error, ErrorKind, Result}; use object_store::path::Path as ObjectStorePath; -use object_store::{ObjectStore, PutPayload, WriteMultipart}; +use object_store::{ObjectStore, ObjectStoreExt, PutPayload, WriteMultipart}; #[cfg(feature = "object_store-s3")] use s3::{build_s3_store, parse_s3_url}; use serde::{Deserialize, Serialize}; @@ -210,15 +210,18 @@ impl Storage for ObjectStoreStorage { Ok(()) } - async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> { - while let Some(path) = paths.next().await { - let (store, object_path) = self.get_store_and_path(&path)?; - store - .delete(&object_path) - .await - .map_err(from_object_store_error)?; - } - Ok(()) + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { + paths + .map(Ok) + .try_for_each_concurrent(16, |path| async move { + let (store, object_path) = self.get_store_and_path(&path)?; + store + .delete(&object_path) + .await + .map_err(from_object_store_error)?; + Ok(()) + }) + .await } fn new_input(&self, path: &str) -> Result { @@ -325,4 +328,37 @@ mod tests { .unwrap(); assert_eq!(path.as_ref(), "data/file.parquet"); } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_storage_serialization_roundtrip() { + let storage = make_s3_storage(); + let serialized = serde_json::to_string(&storage).unwrap(); + let deserialized: ObjectStoreStorage = serde_json::from_str(&serialized).unwrap(); + match deserialized { + ObjectStoreStorage::S3 { config, .. } => { + assert_eq!(config, Arc::new(S3Config::default())); + } + } + } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_storage_factory_serialization_roundtrip() { + let factory = ObjectStoreStorageFactory::S3; + let serialized = serde_json::to_string(&factory).unwrap(); + let deserialized: ObjectStoreStorageFactory = serde_json::from_str(&serialized).unwrap(); + assert!(matches!(deserialized, ObjectStoreStorageFactory::S3)); + } + + #[cfg(feature = "object_store-s3")] + #[test] + fn test_file_io_serialization_roundtrip() { + use iceberg::io::FileIOBuilder; + let factory = Arc::new(ObjectStoreStorageFactory::S3); + let file_io = FileIOBuilder::new(factory).build(); + let bytes = file_io.serialize_all().unwrap(); + let deserialized = iceberg::io::FileIO::deserialize_all(&bytes).unwrap(); + assert_eq!(file_io.config(), deserialized.config()); + } } diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 8ff03ded67..0b385ea4fc 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -19,13 +19,13 @@ use std::sync::Arc; use iceberg::io::S3Config; use iceberg::{Error, ErrorKind, Result}; -use object_store::ObjectStore; use object_store::aws::AmazonS3Builder; +use object_store::ObjectStore; use url::Url; /// Parse an absolute S3 URL into (scheme, bucket, relative_path). /// -/// Accepts `s3://` and `s3a://` schemes. +/// Accepts `s3://` and `s3a://` `s3n://` schemes. pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { let url = Url::parse(path).map_err(|e| { Error::new(ErrorKind::DataInvalid, format!("Invalid URL: {path}")).with_source(e) @@ -33,7 +33,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { let scheme = &path[..url.scheme().len()]; match scheme { - "s3" | "s3a" => {} + "s3" | "s3a" | "s3n" => {} _ => { return Err(Error::new( ErrorKind::DataInvalid, @@ -131,4 +131,13 @@ mod tests { assert_eq!(bucket, "my-bucket"); assert_eq!(relative, ""); } + + #[test] + fn test_parse_s3n_url() { + let (schema, bucket, relative) = + parse_s3_url("s3n://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(schema, "s3n"); + assert_eq!(bucket, "my-bucket"); + assert_eq!(relative, "path/to/file.parquet"); + } } From b22898d0cf391f1e290d21c8042792605e48bc24 Mon Sep 17 00:00:00 2001 From: dron Date: Mon, 7 Sep 2026 16:15:41 +0530 Subject: [PATCH 03/13] feat(storage): zero-copy writes and empty bucket validation in object_store --- crates/storage/object_store/src/lib.rs | 2 +- crates/storage/object_store/src/s3.rs | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 419f7ca819..e08d2e79b4 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -267,7 +267,7 @@ impl FileWrite for ObjectStoreWriter { .writer .as_mut() .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Writer has already been closed"))?; - writer.write(&bs); + writer.put(bs); Ok(()) } diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 0b385ea4fc..52ca1ff3c6 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -49,6 +49,13 @@ pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { ) })?; + if bucket_str.is_empty() { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Invalid s3 url: {path}, missing bucket"), + )); + }; + let prefix_len = scheme.len() + "://".len() + bucket_str.len() + "/".len(); let relative = if path.len() > prefix_len { &path[prefix_len..] @@ -140,4 +147,10 @@ mod tests { assert_eq!(bucket, "my-bucket"); assert_eq!(relative, "path/to/file.parquet"); } + + #[test] + fn test_parse_s3_url_empty_bucket() { + assert!(parse_s3_url("s3:///path/to/file.parquet").is_err()); + assert!(parse_s3_url("s3://").is_err()); + } } From 6457dbd0a6c8a2af2d704ff4eab263354f6a366e Mon Sep 17 00:00:00 2001 From: dron Date: Mon, 7 Sep 2026 17:01:13 +0530 Subject: [PATCH 04/13] chore: add public-api.txt, format toml with taplo, and fix doc links --- Cargo.lock | 1 + Cargo.toml | 2 +- crates/storage/object_store/Cargo.toml | 8 +- crates/storage/object_store/LICENSE | 1 + crates/storage/object_store/NOTICE | 1 + crates/storage/object_store/public-api.txt | 65 ++++++ crates/storage/object_store/src/lib.rs | 178 ++++++++++------ crates/storage/object_store/src/s3.rs | 201 ++++++++++++++---- .../object_store/tests/file_io_s3_test.rs | 159 ++++++++++++++ 9 files changed, 508 insertions(+), 108 deletions(-) create mode 120000 crates/storage/object_store/LICENSE create mode 120000 crates/storage/object_store/NOTICE create mode 100644 crates/storage/object_store/public-api.txt create mode 100644 crates/storage/object_store/tests/file_io_s3_test.rs diff --git a/Cargo.lock b/Cargo.lock index 8cf8a714ff..957182bde9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3993,6 +3993,7 @@ dependencies = [ "dashmap", "futures", "iceberg", + "iceberg_test_utils", "object_store", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 30911f7948..b51040bc69 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -121,8 +121,8 @@ mockall = "0.13.1" mockito = "1" motore-macros = "0.4.3" murmur3 = "0.5.2" -once_cell = "1.20" object_store = "0.13" +once_cell = "1.20" opendal = "0.58" ordered-float = "4" parquet = "59.2" diff --git a/crates/storage/object_store/Cargo.toml b/crates/storage/object_store/Cargo.toml index 49e6e59d29..758246f439 100644 --- a/crates/storage/object_store/Cargo.toml +++ b/crates/storage/object_store/Cargo.toml @@ -16,12 +16,12 @@ # under the License. [package] -name = "iceberg-storage-object_store" edition = { workspace = true } -version = { workspace = true } license = { workspace = true } -repository = { workspace = true } +name = "iceberg-storage-object_store" publish = true +repository = { workspace = true } +version = { workspace = true } categories = ["database"] description = "Apache Iceberg object_store storage implementation" @@ -39,10 +39,12 @@ futures = { workspace = true } iceberg = { workspace = true } object_store = { workspace = true } serde = { workspace = true } +tokio = { workspace = true } typetag = { workspace = true } url = { workspace = true } [dev-dependencies] +iceberg_test_utils = { path = "../../test_utils", features = ["tests"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/crates/storage/object_store/LICENSE b/crates/storage/object_store/LICENSE new file mode 120000 index 0000000000..5853aaea53 --- /dev/null +++ b/crates/storage/object_store/LICENSE @@ -0,0 +1 @@ +../../../LICENSE \ No newline at end of file diff --git a/crates/storage/object_store/NOTICE b/crates/storage/object_store/NOTICE new file mode 120000 index 0000000000..295f6bdb3a --- /dev/null +++ b/crates/storage/object_store/NOTICE @@ -0,0 +1 @@ +../../../NOTICE \ No newline at end of file diff --git a/crates/storage/object_store/public-api.txt b/crates/storage/object_store/public-api.txt new file mode 100644 index 0000000000..f323bb8ac2 --- /dev/null +++ b/crates/storage/object_store/public-api.txt @@ -0,0 +1,65 @@ +pub mod iceberg_storage_object_store +pub enum iceberg_storage_object_store::ObjectStoreStorage +pub iceberg_storage_object_store::ObjectStoreStorage::S3(iceberg_storage_object_store::S3Storage) +impl core::clone::Clone for iceberg_storage_object_store::ObjectStoreStorage +pub fn iceberg_storage_object_store::ObjectStoreStorage::clone(&self) -> iceberg_storage_object_store::ObjectStoreStorage +impl core::fmt::Debug for iceberg_storage_object_store::ObjectStoreStorage +pub fn iceberg_storage_object_store::ObjectStoreStorage::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg::io::storage::Storage for iceberg_storage_object_store::ObjectStoreStorage +pub fn iceberg_storage_object_store::ObjectStoreStorage::delete<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::delete_prefix<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::delete_stream<'life0, 'async_trait>(&'life0 self, paths: futures_core::stream::BoxStream<'static, alloc::string::String>) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::exists<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::metadata<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::new_input(&self, path: &str) -> iceberg::error::Result +pub fn iceberg_storage_object_store::ObjectStoreStorage::new_output(&self, path: &str) -> iceberg::error::Result +pub fn iceberg_storage_object_store::ObjectStoreStorage::read<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::reader<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::write<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str, bs: bytes::bytes::Bytes) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_storage_object_store::ObjectStoreStorage::writer<'life0, 'life1, 'async_trait>(&'life0 self, path: &'life1 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +impl serde_core::ser::Serialize for iceberg_storage_object_store::ObjectStoreStorage +pub fn iceberg_storage_object_store::ObjectStoreStorage::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::ObjectStoreStorage +pub fn iceberg_storage_object_store::ObjectStoreStorage::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +impl core::marker::Freeze for iceberg_storage_object_store::ObjectStoreStorage +impl core::marker::Send for iceberg_storage_object_store::ObjectStoreStorage +impl core::marker::Sync for iceberg_storage_object_store::ObjectStoreStorage +impl core::marker::Unpin for iceberg_storage_object_store::ObjectStoreStorage +impl core::marker::UnsafeUnpin for iceberg_storage_object_store::ObjectStoreStorage +impl !core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::ObjectStoreStorage +impl !core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::ObjectStoreStorage +pub enum iceberg_storage_object_store::ObjectStoreStorageFactory +pub iceberg_storage_object_store::ObjectStoreStorageFactory::S3 +impl core::clone::Clone for iceberg_storage_object_store::ObjectStoreStorageFactory +pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::clone(&self) -> iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::fmt::Debug for iceberg_storage_object_store::ObjectStoreStorageFactory +pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg::io::storage::StorageFactory for iceberg_storage_object_store::ObjectStoreStorageFactory +pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::build(&self, config: &iceberg::io::storage::config::StorageConfig) -> iceberg::error::Result> +impl serde_core::ser::Serialize for iceberg_storage_object_store::ObjectStoreStorageFactory +pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::ObjectStoreStorageFactory +pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +impl core::marker::Freeze for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::marker::Send for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::marker::Sync for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::marker::Unpin for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::marker::UnsafeUnpin for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::ObjectStoreStorageFactory +impl core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::ObjectStoreStorageFactory +pub struct iceberg_storage_object_store::S3Storage +impl core::clone::Clone for iceberg_storage_object_store::S3Storage +pub fn iceberg_storage_object_store::S3Storage::clone(&self) -> iceberg_storage_object_store::S3Storage +impl core::fmt::Debug for iceberg_storage_object_store::S3Storage +pub fn iceberg_storage_object_store::S3Storage::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl serde_core::ser::Serialize for iceberg_storage_object_store::S3Storage +pub fn iceberg_storage_object_store::S3Storage::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::S3Storage +pub fn iceberg_storage_object_store::S3Storage::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +impl core::marker::Freeze for iceberg_storage_object_store::S3Storage +impl core::marker::Send for iceberg_storage_object_store::S3Storage +impl core::marker::Sync for iceberg_storage_object_store::S3Storage +impl core::marker::Unpin for iceberg_storage_object_store::S3Storage +impl core::marker::UnsafeUnpin for iceberg_storage_object_store::S3Storage +impl !core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::S3Storage +impl !core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::S3Storage diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index e08d2e79b4..351029e0dc 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -18,8 +18,8 @@ //! `object_store`-based storage implementation for Apache Iceberg. //! //! This crate provides [`ObjectStoreStorage`] and [`ObjectStoreStorageFactory`], -//! which implement the [`Storage`](iceberg::io::Storage) and -//! [`StorageFactory`](iceberg::io::StorageFactory) traits from the `iceberg` crate +//! which implement the [`Storage`] and +//! [`StorageFactory`] traits from the `iceberg` crate //! using the [`object_store`](https://docs.rs/object_store) crate as the backend. //! //! Currently only S3 storage is supported (via the `object_store-s3` feature flag, @@ -54,6 +54,11 @@ fn from_object_store_error(e: object_store::Error) -> Error { Error::new(ErrorKind::Unexpected, "Failure in doing io operation").with_source(e) } +/// Convert `object_store::ObjectMeta` into `iceberg::io::FileMetadata`. +fn to_file_metadata(meta: object_store::ObjectMeta) -> FileMetadata { + FileMetadata { size: meta.size } +} + /// `object_store`-based storage factory. /// /// Use this factory with `FileIOBuilder::new(factory)` to create FileIO instances @@ -73,15 +78,25 @@ impl StorageFactory for ObjectStoreStorageFactory { #[cfg(feature = "object_store-s3")] ObjectStoreStorageFactory::S3 => { let s3_config = S3Config::try_from(config)?; - Ok(Arc::new(ObjectStoreStorage::S3 { + Ok(Arc::new(ObjectStoreStorage::S3(S3Storage { config: Arc::new(s3_config), store_cache: Arc::new(DashMap::new()), - })) + }))) } } } } +type StoreCache = Arc>>; + +/// `object_store` S3 storage state. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct S3Storage { + config: Arc, + #[serde(skip, default)] + store_cache: StoreCache, +} + /// `object_store`-based storage implementation. /// /// Stores are cached per bucket to avoid rebuilding the client on every operation. @@ -89,33 +104,42 @@ impl StorageFactory for ObjectStoreStorageFactory { pub enum ObjectStoreStorage { /// S3 storage variant. #[cfg(feature = "object_store-s3")] - S3 { - /// Parsed S3 configuration from iceberg core. - config: Arc, - /// Per-bucket store cache. - #[serde(skip, default)] - store_cache: Arc>>, - }, + S3(S3Storage), +} + +struct StoreAndPath { + store: Arc, + path: ObjectStorePath, } impl ObjectStoreStorage { /// Get or create a cached store and extract the relative `ObjectStorePath`. - fn get_store_and_path(&self, path: &str) -> Result<(Arc, ObjectStorePath)> { + fn get_store_and_path(&self, path: &str) -> Result { match self { #[cfg(feature = "object_store-s3")] - ObjectStoreStorage::S3 { - config, - store_cache, - } => { - let (_scheme, bucket, relative) = parse_s3_url(path)?; - - let store = store_cache - .entry(bucket.to_string()) - .or_try_insert_with(|| build_s3_store(config, bucket))? + ObjectStoreStorage::S3(s3) => { + let parsed = parse_s3_url(path)?; + + let store = s3 + .store_cache + .entry(parsed.bucket.clone()) + .or_try_insert_with(|| build_s3_store(&s3.config, &parsed.bucket))? .value() .clone(); - Ok((store, ObjectStorePath::from(relative))) + let object_path = + ObjectStorePath::from_url_path(&parsed.relative).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid URL path: {}", parsed.relative), + ) + .with_source(e) + })?; + + Ok(StoreAndPath { + store, + path: object_path, + }) } } } @@ -125,8 +149,8 @@ impl ObjectStoreStorage { #[async_trait] impl Storage for ObjectStoreStorage { async fn exists(&self, path: &str) -> Result { - let (store, object_path) = self.get_store_and_path(path)?; - match store.head(&object_path).await { + let target = self.get_store_and_path(path)?; + match target.store.head(&target.path).await { Ok(_) => Ok(true), Err(object_store::Error::NotFound { .. }) => Ok(false), Err(e) => Err(from_object_store_error(e)), @@ -134,46 +158,48 @@ impl Storage for ObjectStoreStorage { } async fn metadata(&self, path: &str) -> Result { - let (store, object_path) = self.get_store_and_path(path)?; - let meta = store - .head(&object_path) + let target = self.get_store_and_path(path)?; + let meta = target + .store + .head(&target.path) .await .map_err(from_object_store_error)?; - Ok(FileMetadata { - size: meta.size as u64, - }) + Ok(to_file_metadata(meta)) } async fn read(&self, path: &str) -> Result { - let (store, object_path) = self.get_store_and_path(path)?; - let result = store - .get(&object_path) + let target = self.get_store_and_path(path)?; + let result = target + .store + .get(&target.path) .await .map_err(from_object_store_error)?; result.bytes().await.map_err(from_object_store_error) } async fn reader(&self, path: &str) -> Result> { - let (store, object_path) = self.get_store_and_path(path)?; + let target = self.get_store_and_path(path)?; Ok(Box::new(ObjectStoreReader { - store, - path: object_path, + store: target.store, + path: target.path, })) } async fn write(&self, path: &str, bs: Bytes) -> Result<()> { - let (store, object_path) = self.get_store_and_path(path)?; - store - .put(&object_path, PutPayload::from_bytes(bs)) + let target = self.get_store_and_path(path)?; + target + .store + .put(&target.path, PutPayload::from_bytes(bs)) .await .map_err(from_object_store_error)?; Ok(()) } async fn writer(&self, path: &str) -> Result> { - let (store, object_path) = self.get_store_and_path(path)?; - let upload = store - .put_multipart(&object_path) + let target = self.get_store_and_path(path)?; + let upload = target + .store + .put_multipart(&target.path) .await .map_err(from_object_store_error)?; let writer = WriteMultipart::new(upload); @@ -183,26 +209,28 @@ impl Storage for ObjectStoreStorage { } async fn delete(&self, path: &str) -> Result<()> { - let (store, object_path) = self.get_store_and_path(path)?; - store - .delete(&object_path) + let target = self.get_store_and_path(path)?; + target + .store + .delete(&target.path) .await .map_err(from_object_store_error)?; Ok(()) } async fn delete_prefix(&self, path: &str) -> Result<()> { - let (store, object_path) = self.get_store_and_path(path)?; - let prefix = if object_path.as_ref().ends_with('/') { - object_path + let target = self.get_store_and_path(path)?; + let prefix = if target.path.as_ref().ends_with('/') { + target.path } else { - ObjectStorePath::from(format!("{}/", object_path.as_ref())) + ObjectStorePath::from(format!("{}/", target.path.as_ref())) }; - let mut list_stream = store.list(Some(&prefix)); + let mut list_stream = target.store.list(Some(&prefix)); while let Some(entry) = list_stream.next().await { let entry = entry.map_err(from_object_store_error)?; - store + target + .store .delete(&entry.location) .await .map_err(from_object_store_error)?; @@ -214,9 +242,10 @@ impl Storage for ObjectStoreStorage { paths .map(Ok) .try_for_each_concurrent(16, |path| async move { - let (store, object_path) = self.get_store_and_path(&path)?; - store - .delete(&object_path) + let target = self.get_store_and_path(&path)?; + target + .store + .delete(&target.path) .await .map_err(from_object_store_error)?; Ok(()) @@ -260,6 +289,18 @@ struct ObjectStoreWriter { writer: Option, } +impl Drop for ObjectStoreWriter { + fn drop(&mut self) { + if let Some(writer) = self.writer.take() + && let Ok(handle) = tokio::runtime::Handle::try_current() + { + handle.spawn(async move { + let _ = writer.abort().await; + }); + } + } +} + #[async_trait] impl FileWrite for ObjectStoreWriter { async fn write(&mut self, bs: Bytes) -> Result<()> { @@ -287,46 +328,51 @@ mod tests { #[cfg(feature = "object_store-s3")] fn make_s3_storage() -> ObjectStoreStorage { - ObjectStoreStorage::S3 { + ObjectStoreStorage::S3(S3Storage { config: Arc::new(S3Config::default()), store_cache: Arc::new(DashMap::new()), - } + }) } #[cfg(feature = "object_store-s3")] #[test] fn test_store_cache_reuses_store() { let storage = make_s3_storage(); - let (store1, _) = storage + let target1 = storage .get_store_and_path("s3://test-bucket/file1.parquet") .unwrap(); - let (store2, _) = storage + let target2 = storage .get_store_and_path("s3://test-bucket/file2.parquet") .unwrap(); - assert!(Arc::ptr_eq(&store1, &store2)); + assert!(Arc::ptr_eq(&target1.store, &target2.store)); } #[cfg(feature = "object_store-s3")] #[test] fn test_store_cache_different_buckets() { let storage = make_s3_storage(); - let (store1, _) = storage + let target1 = storage .get_store_and_path("s3://bucket-a/file.parquet") .unwrap(); - let (store2, _) = storage + let target2 = storage .get_store_and_path("s3://bucket-b/file.parquet") .unwrap(); - assert!(!Arc::ptr_eq(&store1, &store2)); + assert!(!Arc::ptr_eq(&target1.store, &target2.store)); } #[cfg(feature = "object_store-s3")] #[test] fn test_relative_path_extraction() { let storage = make_s3_storage(); - let (_, path) = storage + let target = storage .get_store_and_path("s3://my-bucket/data/file.parquet") .unwrap(); - assert_eq!(path.as_ref(), "data/file.parquet"); + assert_eq!(target.path.as_ref(), "data/file.parquet"); + + let target_encoded = storage + .get_store_and_path("s3://my-bucket/data%20dir/file.parquet") + .unwrap(); + assert_eq!(target_encoded.path.as_ref(), "data dir/file.parquet"); } #[cfg(feature = "object_store-s3")] @@ -336,8 +382,8 @@ mod tests { let serialized = serde_json::to_string(&storage).unwrap(); let deserialized: ObjectStoreStorage = serde_json::from_str(&serialized).unwrap(); match deserialized { - ObjectStoreStorage::S3 { config, .. } => { - assert_eq!(config, Arc::new(S3Config::default())); + ObjectStoreStorage::S3(s3) => { + assert_eq!(s3.config, Arc::new(S3Config::default())); } } } diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 52ca1ff3c6..5f33a5a07c 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -15,23 +15,32 @@ // specific language governing permissions and limitations // under the License. +use std::str::FromStr; use std::sync::Arc; use iceberg::io::S3Config; use iceberg::{Error, ErrorKind, Result}; -use object_store::aws::AmazonS3Builder; use object_store::ObjectStore; +use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; use url::Url; -/// Parse an absolute S3 URL into (scheme, bucket, relative_path). +/// Parsed components of an S3 URL. +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct ParsedS3Url { + pub(crate) scheme: String, + pub(crate) bucket: String, + pub(crate) relative: String, +} + +/// Parse an absolute S3 URL into [`ParsedS3Url`]. /// -/// Accepts `s3://` and `s3a://` `s3n://` schemes. -pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { +/// Accepts `s3://`, `s3a://`, and `s3n://` schemes. +pub(crate) fn parse_s3_url(path: &str) -> Result { let url = Url::parse(path).map_err(|e| { Error::new(ErrorKind::DataInvalid, format!("Invalid URL: {path}")).with_source(e) })?; - let scheme = &path[..url.scheme().len()]; + let scheme = url.scheme(); match scheme { "s3" | "s3a" | "s3n" => {} _ => { @@ -42,35 +51,50 @@ pub(crate) fn parse_s3_url(path: &str) -> Result<(&str, &str, &str)> { } } - let bucket_str = url.host_str().ok_or_else(|| { + let bucket = url.host_str().ok_or_else(|| { Error::new( ErrorKind::DataInvalid, format!("Invalid s3 url: {path}, missing bucket"), ) })?; - if bucket_str.is_empty() { + if bucket.is_empty() { return Err(Error::new( ErrorKind::DataInvalid, - format!("Invalid s3 url: {path}, missing bucket"), + format!("Empty s3 url: {path}, missing bucket"), )); - }; - - let prefix_len = scheme.len() + "://".len() + bucket_str.len() + "/".len(); - let relative = if path.len() > prefix_len { - &path[prefix_len..] - } else { - "" - }; + } - let bucket_start = scheme.len() + "://".len(); - let bucket = &path[bucket_start..bucket_start + bucket_str.len()]; + let relative = url.path().trim_start_matches('/'); - Ok((scheme, bucket, relative)) + Ok(ParsedS3Url { + scheme: scheme.to_string(), + bucket: bucket.to_string(), + relative: relative.to_string(), + }) } /// Build an `AmazonS3` store from iceberg's `S3Config` for a given bucket. pub(crate) fn build_s3_store(config: &S3Config, bucket: &str) -> Result> { + if config.role_arn.is_some() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "S3 assume-role (role_arn) is not supported by object_store backend", + )); + } + if config.disable_ec2_metadata { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "S3 disable_ec2_metadata is not supported by object_store backend", + )); + } + if config.disable_config_load { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "S3 disable_config_load is not supported by object_store backend", + )); + } + let mut builder = AmazonS3Builder::new().with_bucket_name(bucket); if let Some(ref endpoint) = config.endpoint { @@ -98,6 +122,37 @@ pub(crate) fn build_s3_store(config: &S3Config, bucket: &str) -> Result { + let key = config + .server_side_encryption_aws_kms_key_id + .as_deref() + .unwrap_or_default(); + builder = builder.with_sse_kms_encryption(key); + } + "AES256" => { + builder = builder.with_config( + AmazonS3ConfigKey::from_str("aws_server_side_encryption").map_err(|e| { + Error::new(ErrorKind::Unexpected, "Failed to parse S3 config key") + .with_source(e) + })?, + "AES256", + ); + } + other => { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("Unsupported server side encryption type: {other}"), + )); + } + } + } + + if let Some(ref custom_key) = config.server_side_encryption_customer_key { + builder = builder.with_ssec_encryption(custom_key); + } + let store = builder.build().map_err(|e| { Error::new(ErrorKind::Unexpected, "Failed to build S3 object store").with_source(e) })?; @@ -110,20 +165,18 @@ mod tests { #[test] fn test_parse_s3_url() { - let (scheme, bucket, relative) = - parse_s3_url("s3://my-bucket/path/to/file.parquet").unwrap(); - assert_eq!(scheme, "s3"); - assert_eq!(bucket, "my-bucket"); - assert_eq!(relative, "path/to/file.parquet"); + let parsed = parse_s3_url("s3://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "path/to/file.parquet"); } #[test] fn test_parse_s3a_url() { - let (scheme, bucket, relative) = - parse_s3_url("s3a://my-bucket/path/to/file.parquet").unwrap(); - assert_eq!(scheme, "s3a"); - assert_eq!(bucket, "my-bucket"); - assert_eq!(relative, "path/to/file.parquet"); + let parsed = parse_s3_url("s3a://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3a"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "path/to/file.parquet"); } #[test] @@ -133,19 +186,42 @@ mod tests { #[test] fn test_parse_s3_url_bucket_only() { - let (scheme, bucket, relative) = parse_s3_url("s3://my-bucket/").unwrap(); - assert_eq!(scheme, "s3"); - assert_eq!(bucket, "my-bucket"); - assert_eq!(relative, ""); + let parsed = parse_s3_url("s3://my-bucket/").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, ""); } #[test] fn test_parse_s3n_url() { - let (schema, bucket, relative) = - parse_s3_url("s3n://my-bucket/path/to/file.parquet").unwrap(); - assert_eq!(schema, "s3n"); - assert_eq!(bucket, "my-bucket"); - assert_eq!(relative, "path/to/file.parquet"); + let parsed = parse_s3_url("s3n://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3n"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "path/to/file.parquet"); + } + + #[test] + fn test_parse_s3_url_uppercase_scheme() { + let parsed = parse_s3_url("S3://my-bucket/path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "path/to/file.parquet"); + } + + #[test] + fn test_parse_s3_url_percent_encoded_bucket() { + let parsed = parse_s3_url("s3://my%2Dbucket/path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my%2Dbucket"); + assert_eq!(parsed.relative, "path/to/file.parquet"); + } + + #[test] + fn test_parse_s3_url_percent_encoded_path() { + let parsed = parse_s3_url("s3://my-bucket/path%20with%20spaces/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "path%20with%20spaces/file.parquet"); } #[test] @@ -153,4 +229,53 @@ mod tests { assert!(parse_s3_url("s3:///path/to/file.parquet").is_err()); assert!(parse_s3_url("s3://").is_err()); } + + #[test] + fn test_build_s3_store_kms_encryption() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption("aws:kms") + .server_side_encryption_aws_kms_key_id("arn:aws:kms:us-east-1:123456789012:key/test") + .build(); + assert!(build_s3_store(&config, "my-bucket").is_ok()); + } + + #[test] + fn test_build_s3_store_ssec_encryption() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption_customer_key("MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU2Nzg5MDE=") + .build(); + assert!(build_s3_store(&config, "my-bucket").is_ok()); + } + + #[test] + fn test_build_s3_store_unsupported_role_arn() { + let config = S3Config::builder() + .region("us-east-1") + .role_arn("arn:aws:iam::123456789012:role/test-role") + .build(); + let err = build_s3_store(&config, "my-bucket").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + } + + #[test] + fn test_build_s3_store_unsupported_disable_ec2_metadata() { + let config = S3Config::builder() + .region("us-east-1") + .disable_ec2_metadata(true) + .build(); + let err = build_s3_store(&config, "my-bucket").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + } + + #[test] + fn test_build_s3_store_unsupported_disable_config_load() { + let config = S3Config::builder() + .region("us-east-1") + .disable_config_load(true) + .build(); + let err = build_s3_store(&config, "my-bucket").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + } } diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs new file mode 100644 index 0000000000..e2805f49de --- /dev/null +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -0,0 +1,159 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Integration tests for FileIO S3 using object_store backend. +//! +//! These tests assume Docker containers are started externally via `make docker-up`. +//! Each test uses unique file paths based on module path to avoid conflicts. + +#[cfg(feature = "object_store-s3")] +mod tests { + use std::sync::Arc; + + use bytes::Bytes; + use futures::StreamExt; + use iceberg::io::{ + FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, + S3_SECRET_ACCESS_KEY, + }; + use iceberg_storage_object_store::ObjectStoreStorageFactory; + use iceberg_test_utils::{get_minio_endpoint, normalize_test_name_with_parts, set_up}; + + async fn get_file_io() -> FileIO { + set_up(); + + let minio_endpoint = get_minio_endpoint(); + + FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) + .with_props(vec![ + (S3_ENDPOINT, minio_endpoint), + (S3_ACCESS_KEY_ID, "admin".to_string()), + (S3_SECRET_ACCESS_KEY, "password".to_string()), + (S3_REGION, "us-east-1".to_string()), + (S3_PATH_STYLE_ACCESS, "true".to_string()), + ]) + .build() + } + + fn roundtrip_file_io(file_io: &FileIO) -> FileIO { + let serialized = file_io.serialize_all().unwrap(); + FileIO::deserialize_all(&serialized).unwrap() + } + + #[tokio::test] + async fn test_file_io_s3_serialization_roundtrip() { + let file_io = roundtrip_file_io(&get_file_io().await); + let path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_serialization_roundtrip") + ); + + let _ = file_io.delete(&path).await; + file_io + .new_output(&path) + .unwrap() + .write(Bytes::from_static(b"roundtrip")) + .await + .unwrap(); + assert_eq!( + file_io.new_input(&path).unwrap().read().await.unwrap(), + Bytes::from_static(b"roundtrip") + ); + file_io.delete(&path).await.unwrap(); + assert!(!file_io.exists(&path).await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_exists() { + let file_io = get_file_io().await; + assert!(!file_io.exists("s3://bucket2/any").await.unwrap()); + assert!(file_io.exists("s3://bucket1/").await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_output() { + let file_io = get_file_io().await; + let output_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_output") + ); + let _ = file_io.delete(&output_path).await; + assert!(!file_io.exists(&output_path).await.unwrap()); + let output_file = file_io.new_output(&output_path).unwrap(); + { + output_file.write("123".into()).await.unwrap(); + } + assert!(file_io.exists(&output_path).await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_input() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_input") + ); + let output_file = file_io.new_output(&file_path).unwrap(); + { + output_file.write("test_input".into()).await.unwrap(); + } + + let input_file = file_io.new_input(&file_path).unwrap(); + { + let buffer = input_file.read().await.unwrap(); + assert_eq!(buffer, "test_input".as_bytes()); + } + } + + #[tokio::test] + async fn test_file_io_s3_delete_stream() { + let file_io = get_file_io().await; + + let paths: Vec = (0..5) + .map(|i| { + format!( + "s3://bucket1/{}/file-{i}", + normalize_test_name_with_parts!("test_file_io_s3_delete_stream") + ) + }) + .collect(); + for path in &paths { + let _ = file_io.delete(path).await; + file_io + .new_output(path) + .unwrap() + .write("delete-me".into()) + .await + .unwrap(); + assert!(file_io.exists(path).await.unwrap()); + } + + let stream = futures::stream::iter(paths.clone()).boxed(); + file_io.delete_stream(stream).await.unwrap(); + + for path in &paths { + assert!(!file_io.exists(path).await.unwrap()); + } + } + + #[tokio::test] + async fn test_file_io_s3_delete_stream_empty() { + let file_io = get_file_io().await; + let stream = futures::stream::empty().boxed(); + file_io.delete_stream(stream).await.unwrap(); + } +} From ed8a3913607c822523d58a9ecdd7ab61652ea2c2 Mon Sep 17 00:00:00 2001 From: Sruhvx Date: Wed, 16 Sep 2026 20:21:21 +0530 Subject: [PATCH 05/13] fix(ci): update rustls to 0.23.45, unpublish object_store, and fix public-api baseline --- Cargo.lock | 18 +++++++++--------- Cargo.toml | 4 ++++ crates/storage/object_store/Cargo.toml | 2 +- crates/storage/object_store/public-api.txt | 21 --------------------- 4 files changed, 14 insertions(+), 31 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 957182bde9..2c4166c5f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -627,9 +627,9 @@ dependencies = [ [[package]] name = "aws-lc-rs" -version = "1.17.1" +version = "1.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4342d8937fc7e5dd9b1c60292261c0670c882a2cd1719cfc11b1af41731e32ad" +checksum = "b281d307588d634de920874890732659e2e7672f72b5e10e81badc1a8a83621e" dependencies = [ "aws-lc-sys", "zeroize", @@ -637,9 +637,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.42.0" +version = "0.45.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d9ceb1da931507a12f4fccea479dccd00da1943e1b4ae72d8e502d707361444" +checksum = "9bff6c3b54fad79a2e60b8102caf565819711497c1f5f092f49508e2f5c31b27" dependencies = [ "cc", "cmake", @@ -6565,9 +6565,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.41" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "aws-lc-rs", "once_cell", @@ -6629,9 +6629,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "aws-lc-rs", "ring", @@ -7614,7 +7614,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index b51040bc69..8a1b49dd78 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -136,6 +136,10 @@ regex = "1.11.3" reqwest = { version = "0.12.12", default-features = false, features = ["json"] } roaring = { version = "0.11" } rstest = "0.26" +rustls = "0.23.45" +rustls-native-certs = "0.8" +rustls-pki-types = "1.15" +rustls-webpki = "0.103.15" serde = { version = "1.0.219", features = ["rc"] } serde_bytes = "0.11.17" serde_derive = "1.0.219" diff --git a/crates/storage/object_store/Cargo.toml b/crates/storage/object_store/Cargo.toml index 758246f439..2658e1ba07 100644 --- a/crates/storage/object_store/Cargo.toml +++ b/crates/storage/object_store/Cargo.toml @@ -19,7 +19,7 @@ edition = { workspace = true } license = { workspace = true } name = "iceberg-storage-object_store" -publish = true +publish = false repository = { workspace = true } version = { workspace = true } diff --git a/crates/storage/object_store/public-api.txt b/crates/storage/object_store/public-api.txt index f323bb8ac2..bba9e1c2d3 100644 --- a/crates/storage/object_store/public-api.txt +++ b/crates/storage/object_store/public-api.txt @@ -21,13 +21,6 @@ impl serde_core::ser::Serialize for iceberg_storage_object_store::ObjectStoreSto pub fn iceberg_storage_object_store::ObjectStoreStorage::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::ObjectStoreStorage pub fn iceberg_storage_object_store::ObjectStoreStorage::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> -impl core::marker::Freeze for iceberg_storage_object_store::ObjectStoreStorage -impl core::marker::Send for iceberg_storage_object_store::ObjectStoreStorage -impl core::marker::Sync for iceberg_storage_object_store::ObjectStoreStorage -impl core::marker::Unpin for iceberg_storage_object_store::ObjectStoreStorage -impl core::marker::UnsafeUnpin for iceberg_storage_object_store::ObjectStoreStorage -impl !core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::ObjectStoreStorage -impl !core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::ObjectStoreStorage pub enum iceberg_storage_object_store::ObjectStoreStorageFactory pub iceberg_storage_object_store::ObjectStoreStorageFactory::S3 impl core::clone::Clone for iceberg_storage_object_store::ObjectStoreStorageFactory @@ -40,13 +33,6 @@ impl serde_core::ser::Serialize for iceberg_storage_object_store::ObjectStoreSto pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::ObjectStoreStorageFactory pub fn iceberg_storage_object_store::ObjectStoreStorageFactory::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> -impl core::marker::Freeze for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::marker::Send for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::marker::Sync for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::marker::Unpin for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::marker::UnsafeUnpin for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::ObjectStoreStorageFactory -impl core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::ObjectStoreStorageFactory pub struct iceberg_storage_object_store::S3Storage impl core::clone::Clone for iceberg_storage_object_store::S3Storage pub fn iceberg_storage_object_store::S3Storage::clone(&self) -> iceberg_storage_object_store::S3Storage @@ -56,10 +42,3 @@ impl serde_core::ser::Serialize for iceberg_storage_object_store::S3Storage pub fn iceberg_storage_object_store::S3Storage::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_storage_object_store::S3Storage pub fn iceberg_storage_object_store::S3Storage::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> -impl core::marker::Freeze for iceberg_storage_object_store::S3Storage -impl core::marker::Send for iceberg_storage_object_store::S3Storage -impl core::marker::Sync for iceberg_storage_object_store::S3Storage -impl core::marker::Unpin for iceberg_storage_object_store::S3Storage -impl core::marker::UnsafeUnpin for iceberg_storage_object_store::S3Storage -impl !core::panic::unwind_safe::RefUnwindSafe for iceberg_storage_object_store::S3Storage -impl !core::panic::unwind_safe::UnwindSafe for iceberg_storage_object_store::S3Storage From fa5ff4b6bcc64ddba63d58d914b6399f3faa7eae Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 17:58:16 +0530 Subject: [PATCH 06/13] fix(storage): S3 KMS default key, URL percent-decode, batch delete_prefix, and multipart tests --- Cargo.lock | 2 + Cargo.toml | 1 + crates/storage/object_store/Cargo.toml | 4 +- crates/storage/object_store/src/lib.rs | 132 +++++-- crates/storage/object_store/src/s3.rs | 161 ++++++-- .../object_store/tests/file_io_s3_test.rs | 365 +++++++++++++++++- 6 files changed, 594 insertions(+), 71 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5637c88f57..6940c0ba91 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3996,9 +3996,11 @@ dependencies = [ "iceberg", "iceberg_test_utils", "object_store", + "percent-encoding", "serde", "serde_json", "tokio", + "tracing", "typetag", "url", ] diff --git a/Cargo.toml b/Cargo.toml index 8a1b49dd78..2adcf214c5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -126,6 +126,7 @@ once_cell = "1.20" opendal = "0.58" ordered-float = "4" parquet = "59.2" +percent-encoding = "2.3" pilota = "0.11.10" pretty_assertions = "1.4" proc-macro2 = "1" diff --git a/crates/storage/object_store/Cargo.toml b/crates/storage/object_store/Cargo.toml index 2658e1ba07..9b2c869a2c 100644 --- a/crates/storage/object_store/Cargo.toml +++ b/crates/storage/object_store/Cargo.toml @@ -38,8 +38,10 @@ dashmap = { workspace = true } futures = { workspace = true } iceberg = { workspace = true } object_store = { workspace = true } -serde = { workspace = true } +percent-encoding = { workspace = true } +serde = { workspace = true, features = ["derive"] } tokio = { workspace = true } +tracing = { workspace = true } typetag = { workspace = true } url = { workspace = true } diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 351029e0dc..20764777bc 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -49,9 +49,36 @@ use object_store::{ObjectStore, ObjectStoreExt, PutPayload, WriteMultipart}; use s3::{build_s3_store, parse_s3_url}; use serde::{Deserialize, Serialize}; -/// Convert an `object_store::Error` into an `iceberg::Error`. +/// Convert an `object_store::Error` into an `iceberg::Error`, +/// dispatching known variants to their corresponding `ErrorKind`. fn from_object_store_error(e: object_store::Error) -> Error { - Error::new(ErrorKind::Unexpected, "Failure in doing io operation").with_source(e) + let (kind, msg) = match &e { + object_store::Error::NotFound { path, .. } => ( + ErrorKind::DataInvalid, + format!("Object not found: {path}"), + ), + object_store::Error::AlreadyExists { path, .. } => ( + ErrorKind::DataInvalid, + format!("Object already exists: {path}"), + ), + object_store::Error::PermissionDenied { path, .. } => ( + ErrorKind::DataInvalid, + format!("Permission denied: {path}"), + ), + object_store::Error::Unauthenticated { path, .. } => ( + ErrorKind::DataInvalid, + format!("Unauthenticated: {path}"), + ), + object_store::Error::NotSupported { .. } => ( + ErrorKind::FeatureUnsupported, + "Operation not supported".to_string(), + ), + _ => ( + ErrorKind::Unexpected, + "Failure in doing io operation".to_string(), + ), + }; + Error::new(kind, msg).with_source(e) } /// Convert `object_store::ObjectMeta` into `iceberg::io::FileMetadata`. @@ -220,39 +247,53 @@ impl Storage for ObjectStoreStorage { async fn delete_prefix(&self, path: &str) -> Result<()> { let target = self.get_store_and_path(path)?; - let prefix = if target.path.as_ref().ends_with('/') { - target.path - } else { - ObjectStorePath::from(format!("{}/", target.path.as_ref())) - }; + let locations = target + .store + .list(Some(&target.path)) + .map_ok(|m| m.location) + .boxed(); + target + .store + .delete_stream(locations) + .try_collect::>() + .await + .map_err(from_object_store_error)?; + Ok(()) + } + + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { + // Collect and group by bucket so each store gets a single bulk DeleteObjects call. + let all_paths: Vec = paths.collect().await; + let mut grouped: std::collections::HashMap> = + std::collections::HashMap::new(); + let mut stores: std::collections::HashMap> = + std::collections::HashMap::new(); + + for path in all_paths { + let target = self.get_store_and_path(&path)?; + let bucket = match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(_) => { + let parsed = parse_s3_url(&path)?; + parsed.bucket + } + }; + stores.entry(bucket.clone()).or_insert(target.store); + grouped.entry(bucket).or_default().push(target.path); + } - let mut list_stream = target.store.list(Some(&prefix)); - while let Some(entry) = list_stream.next().await { - let entry = entry.map_err(from_object_store_error)?; - target - .store - .delete(&entry.location) + for (bucket, locations) in grouped { + let store = stores.remove(&bucket).expect("store must exist"); + let location_stream = futures::stream::iter(locations.into_iter().map(Ok)).boxed(); + store + .delete_stream(location_stream) + .try_collect::>() .await .map_err(from_object_store_error)?; } Ok(()) } - async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { - paths - .map(Ok) - .try_for_each_concurrent(16, |path| async move { - let target = self.get_store_and_path(&path)?; - target - .store - .delete(&target.path) - .await - .map_err(from_object_store_error)?; - Ok(()) - }) - .await - } - fn new_input(&self, path: &str) -> Result { Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) } @@ -291,12 +332,16 @@ struct ObjectStoreWriter { impl Drop for ObjectStoreWriter { fn drop(&mut self) { - if let Some(writer) = self.writer.take() - && let Ok(handle) = tokio::runtime::Handle::try_current() - { - handle.spawn(async move { - let _ = writer.abort().await; - }); + if let Some(writer) = self.writer.take() { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let _ = writer.abort().await; + }); + } else { + tracing::warn!( + "ObjectStoreWriter dropped outside a Tokio runtime; multipart upload abort skipped" + ); + } } } } @@ -407,4 +452,23 @@ mod tests { let deserialized = iceberg::io::FileIO::deserialize_all(&bytes).unwrap(); assert_eq!(file_io.config(), deserialized.config()); } + + #[tokio::test] + async fn test_writer_already_closed_errors() { + let mut writer = ObjectStoreWriter { writer: None }; + let write_err = writer.write(Bytes::from_static(b"data")).await.unwrap_err(); + assert_eq!(write_err.kind(), ErrorKind::Unexpected); + assert_eq!(write_err.message(), "Writer has already been closed"); + + let close_err = writer.close().await.unwrap_err(); + assert_eq!(close_err.kind(), ErrorKind::Unexpected); + assert_eq!(close_err.message(), "Writer has already been closed"); + } + + #[test] + fn test_writer_drop_outside_tokio_warns_and_does_not_panic() { + // A plain synchronous #[test] runs on an OS thread outside a Tokio runtime context + let writer = ObjectStoreWriter { writer: None }; + drop(writer); + } } diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 5f33a5a07c..2d36ea7d4e 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -22,6 +22,7 @@ use iceberg::io::S3Config; use iceberg::{Error, ErrorKind, Result}; use object_store::ObjectStore; use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; +use percent_encoding::percent_decode_str; use url::Url; /// Parsed components of an S3 URL. @@ -32,6 +33,8 @@ pub(crate) struct ParsedS3Url { pub(crate) relative: String, } + + /// Parse an absolute S3 URL into [`ParsedS3Url`]. /// /// Accepts `s3://`, `s3a://`, and `s3n://` schemes. @@ -40,7 +43,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { Error::new(ErrorKind::DataInvalid, format!("Invalid URL: {path}")).with_source(e) })?; - let scheme = url.scheme(); + let scheme = url.scheme(); match scheme { "s3" | "s3a" | "s3n" => {} _ => { @@ -65,6 +68,16 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { )); } + let bucket = percent_decode_str(bucket) + .decode_utf8() + .map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid percent-encoded bucket in s3 url: {path}"), + ) + .with_source(e) + })?; + let relative = url.path().trim_start_matches('/'); Ok(ParsedS3Url { @@ -74,6 +87,61 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { }) } +/// Parse a string into an [`AmazonS3ConfigKey`]. +fn parse_s3_config_key(key: &str) -> Result { + AmazonS3ConfigKey::from_str(key).map_err(|e| { + Error::new( + ErrorKind::Unexpected, + format!("Failed to parse S3 config key: {key}"), + ) + .with_source(e) + }) +} + +/// Configure Server-Side Encryption on `AmazonS3Builder` from `S3Config`. +/// +/// Uses string-based `with_config(parse_s3_config_key(...), ...)` because +/// `S3EncryptionConfigKey` is not re-exported from `object_store::aws` in 0.13.x. +/// The string keys (`"aws_server_side_encryption"`) are stable and used in +/// object_store's own test suite. +fn configure_sse(mut builder: AmazonS3Builder, config: &S3Config) -> Result { + if let Some(ref sse) = config.server_side_encryption { + match sse.as_str() { + "aws:kms" => match &config.server_side_encryption_aws_kms_key_id { + Some(key) => { + builder = builder.with_sse_kms_encryption(key); + } + None => { + builder = builder.with_config( + parse_s3_config_key("aws_server_side_encryption")?, + "aws:kms", + ); + } + }, + "AES256" => { + builder = builder.with_config( + parse_s3_config_key("aws_server_side_encryption")?, + "AES256", + ); + } + other => { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("Unsupported server side encryption type: {other}"), + )); + } + } + } + + if let Some(ref custom_key) = config.server_side_encryption_customer_key { + // Note: object_store's with_ssec_encryption automatically computes and sets + // the x-amz-server-side-encryption-customer-key-MD5 header from the decoded key. + builder = builder.with_ssec_encryption(custom_key); + } + + Ok(builder) +} + /// Build an `AmazonS3` store from iceberg's `S3Config` for a given bucket. pub(crate) fn build_s3_store(config: &S3Config, bucket: &str) -> Result> { if config.role_arn.is_some() { @@ -115,43 +183,12 @@ pub(crate) fn build_s3_store(config: &S3Config, bucket: &str) -> Result { - let key = config - .server_side_encryption_aws_kms_key_id - .as_deref() - .unwrap_or_default(); - builder = builder.with_sse_kms_encryption(key); - } - "AES256" => { - builder = builder.with_config( - AmazonS3ConfigKey::from_str("aws_server_side_encryption").map_err(|e| { - Error::new(ErrorKind::Unexpected, "Failed to parse S3 config key") - .with_source(e) - })?, - "AES256", - ); - } - other => { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - format!("Unsupported server side encryption type: {other}"), - )); - } - } - } - - if let Some(ref custom_key) = config.server_side_encryption_customer_key { - builder = builder.with_ssec_encryption(custom_key); - } + builder = configure_sse(builder, config)?; let store = builder.build().map_err(|e| { Error::new(ErrorKind::Unexpected, "Failed to build S3 object store").with_source(e) @@ -212,7 +249,7 @@ mod tests { fn test_parse_s3_url_percent_encoded_bucket() { let parsed = parse_s3_url("s3://my%2Dbucket/path/to/file.parquet").unwrap(); assert_eq!(parsed.scheme, "s3"); - assert_eq!(parsed.bucket, "my%2Dbucket"); + assert_eq!(parsed.bucket, "my-bucket"); assert_eq!(parsed.relative, "path/to/file.parquet"); } @@ -230,6 +267,60 @@ mod tests { assert!(parse_s3_url("s3://").is_err()); } + #[test] + fn test_parse_s3_config_key() { + assert!(parse_s3_config_key("aws_server_side_encryption").is_ok()); + let err = parse_s3_config_key("invalid_config_key_foo_bar").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::Unexpected); + } + + #[test] + fn test_configure_sse_kms_default_key() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption("aws:kms") + .build(); + assert!(configure_sse(AmazonS3Builder::new(), &config).is_ok()); + } + + #[test] + fn test_configure_sse_kms_custom_key() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption("aws:kms") + .server_side_encryption_aws_kms_key_id("arn:aws:kms:us-east-1:123456789012:key/test") + .build(); + assert!(configure_sse(AmazonS3Builder::new(), &config).is_ok()); + } + + #[test] + fn test_configure_sse_aes256() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption("AES256") + .build(); + assert!(configure_sse(AmazonS3Builder::new(), &config).is_ok()); + } + + #[test] + fn test_configure_sse_ssec() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption_customer_key("MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU2Nzg5MDE=") + .build(); + assert!(configure_sse(AmazonS3Builder::new(), &config).is_ok()); + } + + #[test] + fn test_configure_sse_unsupported() { + let config = S3Config::builder() + .region("us-east-1") + .server_side_encryption("unsupported_sse") + .build(); + let err = configure_sse(AmazonS3Builder::new(), &config).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + } + #[test] fn test_build_s3_store_kms_encryption() { let config = S3Config::builder() diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index e2805f49de..4cd886ac50 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -28,7 +28,7 @@ mod tests { use futures::StreamExt; use iceberg::io::{ FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, - S3_SECRET_ACCESS_KEY, + S3_SECRET_ACCESS_KEY, S3_SSE_KEY, S3_SSE_TYPE, }; use iceberg_storage_object_store::ObjectStoreStorageFactory; use iceberg_test_utils::{get_minio_endpoint, normalize_test_name_with_parts, set_up}; @@ -156,4 +156,367 @@ mod tests { let stream = futures::stream::empty().boxed(); file_io.delete_stream(stream).await.unwrap(); } + + #[tokio::test] + async fn test_file_io_s3_delete_stream_invalid_url() { + let file_io = get_file_io().await; + let stream = futures::stream::iter(vec!["invalid-url".to_string()]).boxed(); + let res = file_io.delete_stream(stream).await; + assert!(res.is_err()); + } + + #[tokio::test] + async fn test_file_io_s3_multipart_writer() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_multipart_writer") + ); + let _ = file_io.delete(&file_path).await; + + let output_file = file_io.new_output(&file_path).unwrap(); + let mut writer = output_file.writer().await.unwrap(); + + let chunk1 = Bytes::from_static(b"hello "); + let chunk2 = Bytes::from_static(b"multipart "); + let chunk3 = Bytes::from_static(b"world!"); + + writer.write(chunk1).await.unwrap(); + writer.write(chunk2).await.unwrap(); + writer.write(chunk3).await.unwrap(); + writer.close().await.unwrap(); + + assert!(file_io.exists(&file_path).await.unwrap()); + let input_file = file_io.new_input(&file_path).unwrap(); + let content = input_file.read().await.unwrap(); + assert_eq!(content, Bytes::from_static(b"hello multipart world!")); + + file_io.delete(&file_path).await.unwrap(); + assert!(!file_io.exists(&file_path).await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_multipart_writer_drop_aborts() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_multipart_writer_drop_aborts") + ); + let _ = file_io.delete(&file_path).await; + + let output_file = file_io.new_output(&file_path).unwrap(); + let mut writer = output_file.writer().await.unwrap(); + writer + .write(Bytes::from_static(b"uncommitted chunk")) + .await + .unwrap(); + + // Dropping writer without close() should abort the multipart upload + drop(writer); + + // Give background abort task a moment to execute + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + assert!(!file_io.exists(&file_path).await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_percent_encoded_bucket() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://my%2Dbucket/{}", + normalize_test_name_with_parts!("test_file_io_s3_percent_encoded_bucket") + ); + let canonical_path = format!( + "s3://my-bucket/{}", + normalize_test_name_with_parts!("test_file_io_s3_percent_encoded_bucket") + ); + + let _ = file_io.delete(&file_path).await; + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from_static(b"encoded-bucket-content")) + .await + .unwrap(); + + assert!(file_io.exists(&file_path).await.unwrap()); + assert!(file_io.exists(&canonical_path).await.unwrap()); + + let content = file_io + .new_input(&canonical_path) + .unwrap() + .read() + .await + .unwrap(); + assert_eq!(content, Bytes::from_static(b"encoded-bucket-content")); + + file_io.delete(&canonical_path).await.unwrap(); + assert!(!file_io.exists(&file_path).await.unwrap()); + } + + #[tokio::test] + async fn test_file_io_s3_sse_kms_default() { + set_up(); + let endpoint = get_minio_endpoint(); + + let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) + .with_props(vec![ + (S3_ENDPOINT, endpoint), + (S3_ACCESS_KEY_ID, "admin".to_string()), + (S3_SECRET_ACCESS_KEY, "password".to_string()), + (S3_REGION, "us-east-1".to_string()), + (S3_PATH_STYLE_ACCESS, "true".to_string()), + (S3_SSE_TYPE, "kms".to_string()), + ]) + .build(); + + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_sse_kms_default") + ); + + let _ = file_io.delete(&file_path).await; + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from_static(b"kms-encrypted-data")) + .await + .unwrap(); + + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!(content, Bytes::from_static(b"kms-encrypted-data")); + + file_io.delete(&file_path).await.unwrap(); + } + + #[tokio::test] + async fn test_file_io_s3_sse_kms_custom_key() { + set_up(); + let endpoint = get_minio_endpoint(); + + let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) + .with_props(vec![ + (S3_ENDPOINT, endpoint), + (S3_ACCESS_KEY_ID, "admin".to_string()), + (S3_SECRET_ACCESS_KEY, "password".to_string()), + (S3_REGION, "us-east-1".to_string()), + (S3_PATH_STYLE_ACCESS, "true".to_string()), + (S3_SSE_TYPE, "kms".to_string()), + ( + S3_SSE_KEY, + "arn:aws:kms:us-east-1:000000000000:key/a4644f9c-2149-414e-b6a6-a8e82cd6b69e" + .to_string(), + ), + ]) + .build(); + + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_sse_kms_custom_key") + ); + + let _ = file_io.delete(&file_path).await; + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from_static(b"kms-custom-key-encrypted-data")) + .await + .unwrap(); + + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!( + content, + Bytes::from_static(b"kms-custom-key-encrypted-data") + ); + + file_io.delete(&file_path).await.unwrap(); + } + + /// Writes 12 MiB (3 × 4 MiB chunks) to exercise real S3 multipart uploads + /// past the 10 MiB `WriteMultipart` buffer threshold, then reads back and + /// verifies byte-for-byte integrity. + #[tokio::test] + async fn test_file_io_s3_multipart_writer_past_threshold() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_multipart_writer_past_threshold") + ); + let _ = file_io.delete(&file_path).await; + + // 4 MiB chunk with deterministic pattern (repeating 0..=255) + const CHUNK_SIZE: usize = 4 * 1024 * 1024; + let pattern: Vec = (0..CHUNK_SIZE).map(|i| (i % 256) as u8).collect(); + let chunk = Bytes::from(pattern.clone()); + + let output_file = file_io.new_output(&file_path).unwrap(); + let mut writer = output_file.writer().await.unwrap(); + + // Write 3 chunks = 12 MiB total (past 10 MiB threshold) + for _ in 0..3 { + writer.write(chunk.clone()).await.unwrap(); + } + writer.close().await.unwrap(); + + // Read back and verify + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!(content.len(), 3 * CHUNK_SIZE); + for i in 0..3 { + assert_eq!( + &content[i * CHUNK_SIZE..(i + 1) * CHUNK_SIZE], + &pattern[..], + "chunk {i} mismatch" + ); + } + + file_io.delete(&file_path).await.unwrap(); + } + + /// Creates 15 files under `test_prefix/` and 2 under `other_prefix/`, + /// calls `delete_prefix` on `test_prefix/`, then asserts all 15 are gone + /// and the 2 outside the prefix are untouched. + #[tokio::test] + async fn test_file_io_s3_delete_prefix_bulk() { + let file_io = get_file_io().await; + let base = normalize_test_name_with_parts!("test_file_io_s3_delete_prefix_bulk"); + let target_prefix = format!("s3://bucket1/{base}/test_prefix"); + let other_prefix = format!("s3://bucket1/{base}/other_prefix"); + + // Create 15 files under test_prefix/ + for i in 0..15 { + let path = format!("{target_prefix}/file_{i}"); + let _ = file_io.delete(&path).await; + file_io + .new_output(&path) + .unwrap() + .write(Bytes::from(format!("data-{i}"))) + .await + .unwrap(); + } + + // Create 2 files under other_prefix/ + let keep_paths: Vec = (0..2) + .map(|i| format!("{other_prefix}/keep_{i}")) + .collect(); + for path in &keep_paths { + let _ = file_io.delete(path).await; + file_io + .new_output(path) + .unwrap() + .write(Bytes::from_static(b"keep-me")) + .await + .unwrap(); + } + + // Bulk delete under test_prefix/ + file_io.delete_prefix(&target_prefix).await.unwrap(); + + // Assert all 15 test_prefix files are gone + for i in 0..15 { + let path = format!("{target_prefix}/file_{i}"); + assert!( + !file_io.exists(&path).await.unwrap(), + "file_{i} should be deleted" + ); + } + + // Assert the 2 other_prefix files still exist + for path in &keep_paths { + assert!( + file_io.exists(path).await.unwrap(), + "{path} should still exist" + ); + } + + // Clean up + for path in &keep_paths { + file_io.delete(path).await.unwrap(); + } + } + + /// Writes a 1024-byte payload and verifies range reads return exact slices. + #[tokio::test] + async fn test_file_io_s3_range_reader() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_range_reader") + ); + let _ = file_io.delete(&file_path).await; + + // 1024 bytes: 0..=255 repeated 4 times + let payload: Vec = (0..1024).map(|i| (i % 256) as u8).collect(); + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from(payload.clone())) + .await + .unwrap(); + + let reader = file_io + .new_input(&file_path) + .unwrap() + .reader() + .await + .unwrap(); + + // Test various ranges + let r1 = reader.read(0..10).await.unwrap(); + assert_eq!(&r1[..], &payload[0..10]); + + let r2 = reader.read(10..50).await.unwrap(); + assert_eq!(&r2[..], &payload[10..50]); + + let r3 = reader.read(100..200).await.unwrap(); + assert_eq!(&r3[..], &payload[100..200]); + + // Cross the 256-byte pattern boundary + let r4 = reader.read(250..260).await.unwrap(); + assert_eq!(&r4[..], &payload[250..260]); + + // Last 10 bytes + let r5 = reader.read(1014..1024).await.unwrap(); + assert_eq!(&r5[..], &payload[1014..1024]); + + file_io.delete(&file_path).await.unwrap(); + } + + #[tokio::test] + async fn test_file_io_s3_sse_s3_aes256() { + set_up(); + let endpoint = get_minio_endpoint(); + + let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) + .with_props(vec![ + (S3_ENDPOINT, endpoint), + (S3_ACCESS_KEY_ID, "admin".to_string()), + (S3_SECRET_ACCESS_KEY, "password".to_string()), + (S3_REGION, "us-east-1".to_string()), + (S3_PATH_STYLE_ACCESS, "true".to_string()), + (S3_SSE_TYPE, "s3".to_string()), + ]) + .build(); + + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_sse_s3_aes256") + ); + + let _ = file_io.delete(&file_path).await; + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from_static(b"aes256-encrypted-data")) + .await + .unwrap(); + + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!(content, Bytes::from_static(b"aes256-encrypted-data")); + + file_io.delete(&file_path).await.unwrap(); + } } From eb7387f938a39ce71ef866eaf170ef0ca449f0ba Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 18:12:45 +0530 Subject: [PATCH 07/13] chore: fix cargo fmt formatting and match Swatinem/rust-cache in workflows --- .github/workflows/bindings_python_ci.yml | 2 +- .github/workflows/ci.yml | 8 +++--- .github/workflows/public-api.yml | 2 +- crates/storage/object_store/src/lib.rs | 21 +++++++-------- crates/storage/object_store/src/s3.rs | 26 +++++++------------ .../object_store/tests/file_io_s3_test.rs | 4 +-- 6 files changed, 26 insertions(+), 37 deletions(-) diff --git a/.github/workflows/bindings_python_ci.yml b/.github/workflows/bindings_python_ci.yml index 84ce59fca8..68586abff1 100644 --- a/.github/workflows/bindings_python_ci.yml +++ b/.github/workflows/bindings_python_ci.yml @@ -89,7 +89,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts if: runner.os != 'Linux' - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: key: bindings-python save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f9fce8c20c..571b46fcc9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -137,7 +137,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} @@ -158,7 +158,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} @@ -199,7 +199,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} @@ -228,7 +228,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: key: ${{ matrix.test-suite.name }} save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} diff --git a/.github/workflows/public-api.yml b/.github/workflows/public-api.yml index 9b19358fd0..b65b9cf404 100644 --- a/.github/workflows/public-api.yml +++ b/.github/workflows/public-api.yml @@ -45,7 +45,7 @@ jobs: uses: $/.github/actions/setup-builder - name: Cache Rust artifacts - uses: swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 20764777bc..b8e549c49d 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -53,22 +53,19 @@ use serde::{Deserialize, Serialize}; /// dispatching known variants to their corresponding `ErrorKind`. fn from_object_store_error(e: object_store::Error) -> Error { let (kind, msg) = match &e { - object_store::Error::NotFound { path, .. } => ( - ErrorKind::DataInvalid, - format!("Object not found: {path}"), - ), + object_store::Error::NotFound { path, .. } => { + (ErrorKind::DataInvalid, format!("Object not found: {path}")) + } object_store::Error::AlreadyExists { path, .. } => ( ErrorKind::DataInvalid, format!("Object already exists: {path}"), ), - object_store::Error::PermissionDenied { path, .. } => ( - ErrorKind::DataInvalid, - format!("Permission denied: {path}"), - ), - object_store::Error::Unauthenticated { path, .. } => ( - ErrorKind::DataInvalid, - format!("Unauthenticated: {path}"), - ), + object_store::Error::PermissionDenied { path, .. } => { + (ErrorKind::DataInvalid, format!("Permission denied: {path}")) + } + object_store::Error::Unauthenticated { path, .. } => { + (ErrorKind::DataInvalid, format!("Unauthenticated: {path}")) + } object_store::Error::NotSupported { .. } => ( ErrorKind::FeatureUnsupported, "Operation not supported".to_string(), diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 2d36ea7d4e..69b0f066fa 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -33,8 +33,6 @@ pub(crate) struct ParsedS3Url { pub(crate) relative: String, } - - /// Parse an absolute S3 URL into [`ParsedS3Url`]. /// /// Accepts `s3://`, `s3a://`, and `s3n://` schemes. @@ -43,7 +41,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { Error::new(ErrorKind::DataInvalid, format!("Invalid URL: {path}")).with_source(e) })?; - let scheme = url.scheme(); + let scheme = url.scheme(); match scheme { "s3" | "s3a" | "s3n" => {} _ => { @@ -68,15 +66,13 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { )); } - let bucket = percent_decode_str(bucket) - .decode_utf8() - .map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!("Invalid percent-encoded bucket in s3 url: {path}"), - ) - .with_source(e) - })?; + let bucket = percent_decode_str(bucket).decode_utf8().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid percent-encoded bucket in s3 url: {path}"), + ) + .with_source(e) + })?; let relative = url.path().trim_start_matches('/'); @@ -119,10 +115,8 @@ fn configure_sse(mut builder: AmazonS3Builder, config: &S3Config) -> Result { - builder = builder.with_config( - parse_s3_config_key("aws_server_side_encryption")?, - "AES256", - ); + builder = builder + .with_config(parse_s3_config_key("aws_server_side_encryption")?, "AES256"); } other => { return Err(Error::new( diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index 4cd886ac50..2fd100d949 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -398,9 +398,7 @@ mod tests { } // Create 2 files under other_prefix/ - let keep_paths: Vec = (0..2) - .map(|i| format!("{other_prefix}/keep_{i}")) - .collect(); + let keep_paths: Vec = (0..2).map(|i| format!("{other_prefix}/keep_{i}")).collect(); for path in &keep_paths { let _ = file_io.delete(path).await; file_io From 52b11c09e48e26e58edb779000cee7bcf5c267ca Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 18:26:40 +0530 Subject: [PATCH 08/13] test: use existing bucket1 with percent-encoding in file_io_s3_test --- crates/storage/object_store/tests/file_io_s3_test.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index 2fd100d949..552cb84e25 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -224,11 +224,11 @@ mod tests { async fn test_file_io_s3_percent_encoded_bucket() { let file_io = get_file_io().await; let file_path = format!( - "s3://my%2Dbucket/{}", + "s3://bucket%31/{}", normalize_test_name_with_parts!("test_file_io_s3_percent_encoded_bucket") ); let canonical_path = format!( - "s3://my-bucket/{}", + "s3://bucket1/{}", normalize_test_name_with_parts!("test_file_io_s3_percent_encoded_bucket") ); From 08e2338f1507c4f73e46c3ab97afb9263800ffae Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 18:39:38 +0530 Subject: [PATCH 09/13] test(s3): handle MinIO 501 on KMS integration tests --- .../object_store/tests/file_io_s3_test.rs | 56 ++++++++++++------- 1 file changed, 37 insertions(+), 19 deletions(-) diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index 552cb84e25..0ef95ab8f0 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -277,18 +277,27 @@ mod tests { ); let _ = file_io.delete(&file_path).await; - file_io + match file_io .new_output(&file_path) .unwrap() .write(Bytes::from_static(b"kms-encrypted-data")) .await - .unwrap(); - - assert!(file_io.exists(&file_path).await.unwrap()); - let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); - assert_eq!(content, Bytes::from_static(b"kms-encrypted-data")); - - file_io.delete(&file_path).await.unwrap(); + { + Ok(_) => { + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!(content, Bytes::from_static(b"kms-encrypted-data")); + file_io.delete(&file_path).await.unwrap(); + } + Err(e) + if e.to_string().contains("501") + || e.to_string().contains("NotImplemented") + || e.to_string().contains("KMS is not configured") => + { + // MinIO without KES does not configure KMS; passing 501 verifies header was sent + } + Err(e) => panic!("Unexpected error: {e:?}"), + } } #[tokio::test] @@ -318,21 +327,30 @@ mod tests { ); let _ = file_io.delete(&file_path).await; - file_io + match file_io .new_output(&file_path) .unwrap() .write(Bytes::from_static(b"kms-custom-key-encrypted-data")) .await - .unwrap(); - - assert!(file_io.exists(&file_path).await.unwrap()); - let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); - assert_eq!( - content, - Bytes::from_static(b"kms-custom-key-encrypted-data") - ); - - file_io.delete(&file_path).await.unwrap(); + { + Ok(_) => { + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!( + content, + Bytes::from_static(b"kms-custom-key-encrypted-data") + ); + file_io.delete(&file_path).await.unwrap(); + } + Err(e) + if e.to_string().contains("501") + || e.to_string().contains("NotImplemented") + || e.to_string().contains("KMS is not configured") => + { + // MinIO without KES does not configure KMS; passing 501 verifies header was sent + } + Err(e) => panic!("Unexpected error: {e:?}"), + } } /// Writes 12 MiB (3 × 4 MiB chunks) to exercise real S3 multipart uploads From 956743fa0a9086ac543fb727081a68557761cabe Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 18:48:42 +0530 Subject: [PATCH 10/13] test(s3): handle MinIO 501 on SSE-S3 AES256 integration test --- .../object_store/tests/file_io_s3_test.rs | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index 0ef95ab8f0..609bd6446a 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -522,17 +522,26 @@ mod tests { ); let _ = file_io.delete(&file_path).await; - file_io + match file_io .new_output(&file_path) .unwrap() .write(Bytes::from_static(b"aes256-encrypted-data")) .await - .unwrap(); - - assert!(file_io.exists(&file_path).await.unwrap()); - let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); - assert_eq!(content, Bytes::from_static(b"aes256-encrypted-data")); - - file_io.delete(&file_path).await.unwrap(); + { + Ok(_) => { + assert!(file_io.exists(&file_path).await.unwrap()); + let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); + assert_eq!(content, Bytes::from_static(b"aes256-encrypted-data")); + file_io.delete(&file_path).await.unwrap(); + } + Err(e) + if e.to_string().contains("501") + || e.to_string().contains("NotImplemented") + || e.to_string().contains("KMS is not configured") => + { + // MinIO without KES does not configure server-side encryption; passing 501 verifies header was sent + } + Err(e) => panic!("Unexpected error: {e:?}"), + } } } From d6d64b03fa7e3a2fd29e8e14bf6030b5131dc8f2 Mon Sep 17 00:00:00 2001 From: dron Date: Fri, 18 Sep 2026 18:48:42 +0530 Subject: [PATCH 11/13] fix(storage): multipart upload lifecycle aborts, error kind mappings, and bounded chunked delete streams Fixes drop abort handling outside Tokio runtime, adds empty range guard in ObjectStoreReader to avoid S3 416 errors, prevents memory unbounded buffering in delete_prefix and delete_stream by streaming deletions in configurable chunks (via s3.delete-batch-size with a default of 1000 and max-limit warning check), and drains delete streams via try_for_each. --- crates/storage/object_store/src/lib.rs | 467 ++++++++++++++++-- crates/storage/object_store/src/s3.rs | 19 +- .../object_store/tests/file_io_s3_test.rs | 25 - 3 files changed, 427 insertions(+), 84 deletions(-) diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index b8e549c49d..67a86279dd 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -32,7 +32,7 @@ use std::ops::Range; use std::sync::Arc; use async_trait::async_trait; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use dashmap::DashMap; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; @@ -44,7 +44,7 @@ use iceberg::io::{ }; use iceberg::{Error, ErrorKind, Result}; use object_store::path::Path as ObjectStorePath; -use object_store::{ObjectStore, ObjectStoreExt, PutPayload, WriteMultipart}; +use object_store::{MultipartUpload, ObjectStore, ObjectStoreExt, PutPayload, UploadPart}; #[cfg(feature = "object_store-s3")] use s3::{build_s3_store, parse_s3_url}; use serde::{Deserialize, Serialize}; @@ -54,17 +54,17 @@ use serde::{Deserialize, Serialize}; fn from_object_store_error(e: object_store::Error) -> Error { let (kind, msg) = match &e { object_store::Error::NotFound { path, .. } => { - (ErrorKind::DataInvalid, format!("Object not found: {path}")) + (ErrorKind::Unexpected, format!("Object not found: {path}")) } object_store::Error::AlreadyExists { path, .. } => ( - ErrorKind::DataInvalid, + ErrorKind::Unexpected, format!("Object already exists: {path}"), ), object_store::Error::PermissionDenied { path, .. } => { - (ErrorKind::DataInvalid, format!("Permission denied: {path}")) + (ErrorKind::Unexpected, format!("Permission denied: {path}")) } object_store::Error::Unauthenticated { path, .. } => { - (ErrorKind::DataInvalid, format!("Unauthenticated: {path}")) + (ErrorKind::Unexpected, format!("Unauthenticated: {path}")) } object_store::Error::NotSupported { .. } => ( ErrorKind::FeatureUnsupported, @@ -78,6 +78,43 @@ fn from_object_store_error(e: object_store::Error) -> Error { Error::new(kind, msg).with_source(e) } +/// Property key for configuring S3 bulk delete batch size. +pub const S3_DELETE_BATCH_SIZE: &str = "s3.delete-batch-size"; +/// Default batch size for S3 bulk deletions (matches AWS S3 DeleteObjects max). +pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000; +/// Maximum batch size allowed by the AWS S3 DeleteObjects API specification. +pub const S3_MAX_DELETE_BATCH_SIZE: usize = 1000; + +fn parse_delete_batch_size(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_DELETE_BATCH_SIZE) { + match val.parse::() { + Ok(parsed) if parsed > 0 => { + if parsed > S3_MAX_DELETE_BATCH_SIZE { + tracing::warn!( + configured = parsed, + limit = S3_MAX_DELETE_BATCH_SIZE, + "Configured s3.delete-batch-size exceeds AWS S3 hard limit of 1000; requests may fail with MalformedXML" + ); + } + parsed + } + _ => { + tracing::warn!( + val = %val, + "Invalid s3.delete-batch-size; falling back to default 1000" + ); + DEFAULT_DELETE_BATCH_SIZE + } + } + } else { + DEFAULT_DELETE_BATCH_SIZE + } +} + +fn default_delete_batch_size() -> usize { + DEFAULT_DELETE_BATCH_SIZE +} + /// Convert `object_store::ObjectMeta` into `iceberg::io::FileMetadata`. fn to_file_metadata(meta: object_store::ObjectMeta) -> FileMetadata { FileMetadata { size: meta.size } @@ -102,8 +139,16 @@ impl StorageFactory for ObjectStoreStorageFactory { #[cfg(feature = "object_store-s3")] ObjectStoreStorageFactory::S3 => { let s3_config = S3Config::try_from(config)?; + let delete_batch_size = parse_delete_batch_size(config); + tracing::info!( + batch_size = delete_batch_size, + "Initialized S3 storage with delete batch size {} (configure via '{}' to adjust)", + delete_batch_size, + S3_DELETE_BATCH_SIZE + ); Ok(Arc::new(ObjectStoreStorage::S3(S3Storage { config: Arc::new(s3_config), + delete_batch_size, store_cache: Arc::new(DashMap::new()), }))) } @@ -117,6 +162,8 @@ type StoreCache = Arc>>; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct S3Storage { config: Arc, + #[serde(default = "default_delete_batch_size")] + pub delete_batch_size: usize, #[serde(skip, default)] store_cache: StoreCache, } @@ -132,11 +179,25 @@ pub enum ObjectStoreStorage { } struct StoreAndPath { + bucket: String, store: Arc, path: ObjectStorePath, } +/// Helper for batching deletions per bucket. +struct BucketBatch { + store: Arc, + locations: Vec, +} + impl ObjectStoreStorage { + fn delete_batch_size(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.delete_batch_size, + } + } + /// Get or create a cached store and extract the relative `ObjectStorePath`. fn get_store_and_path(&self, path: &str) -> Result { match self { @@ -161,6 +222,7 @@ impl ObjectStoreStorage { })?; Ok(StoreAndPath { + bucket: parsed.bucket, store, path: object_path, }) @@ -226,9 +288,10 @@ impl Storage for ObjectStoreStorage { .put_multipart(&target.path) .await .map_err(from_object_store_error)?; - let writer = WriteMultipart::new(upload); Ok(Box::new(ObjectStoreWriter { - writer: Some(writer), + upload: Some(upload), + buffer: BytesMut::new(), + tasks: Vec::new(), })) } @@ -252,41 +315,42 @@ impl Storage for ObjectStoreStorage { target .store .delete_stream(locations) - .try_collect::>() + .try_for_each(|_| async { Ok(()) }) .await .map_err(from_object_store_error)?; Ok(()) } async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { - // Collect and group by bucket so each store gets a single bulk DeleteObjects call. - let all_paths: Vec = paths.collect().await; - let mut grouped: std::collections::HashMap> = - std::collections::HashMap::new(); - let mut stores: std::collections::HashMap> = - std::collections::HashMap::new(); - - for path in all_paths { - let target = self.get_store_and_path(&path)?; - let bucket = match self { - #[cfg(feature = "object_store-s3")] - ObjectStoreStorage::S3(_) => { - let parsed = parse_s3_url(&path)?; - parsed.bucket - } - }; - stores.entry(bucket.clone()).or_insert(target.store); - grouped.entry(bucket).or_default().push(target.path); - } + let batch_size = self.delete_batch_size(); + let mut chunk_stream = paths.chunks(batch_size); + + while let Some(chunk) = chunk_stream.next().await { + let mut batches: std::collections::HashMap = + std::collections::HashMap::new(); + + for path in chunk { + let target = self.get_store_and_path(&path)?; + batches + .entry(target.bucket) + .or_insert_with(|| BucketBatch { + store: target.store, + locations: Vec::new(), + }) + .locations + .push(target.path); + } - for (bucket, locations) in grouped { - let store = stores.remove(&bucket).expect("store must exist"); - let location_stream = futures::stream::iter(locations.into_iter().map(Ok)).boxed(); - store - .delete_stream(location_stream) - .try_collect::>() - .await - .map_err(from_object_store_error)?; + for (_bucket, batch) in batches { + let location_stream = + futures::stream::iter(batch.locations.into_iter().map(Ok)).boxed(); + batch + .store + .delete_stream(location_stream) + .try_for_each(|_| async { Ok(()) }) + .await + .map_err(from_object_store_error)?; + } } Ok(()) } @@ -309,6 +373,9 @@ struct ObjectStoreReader { #[async_trait] impl FileRead for ObjectStoreReader { async fn read(&self, range: Range) -> Result { + if range.is_empty() { + return Ok(Bytes::new()); + } let opts = object_store::GetOptions { range: Some((range.start..range.end).into()), ..Default::default() @@ -322,17 +389,66 @@ impl FileRead for ObjectStoreReader { } } +/// Minimum part size for S3 multipart upload (5 MiB). +const MIN_PART_SIZE: usize = 5 * 1024 * 1024; + /// Writer that implements `FileWrite` using `object_store` multipart upload. struct ObjectStoreWriter { - writer: Option, + upload: Option>, + buffer: BytesMut, + tasks: Vec, +} + +impl ObjectStoreWriter { + /// Flushes any buffered bytes as an in-flight part upload to S3. + fn flush_buffer( + buffer: &mut BytesMut, + tasks: &mut Vec, + upload: &mut Box, + ) { + if !buffer.is_empty() { + let part_data = std::mem::take(buffer).freeze(); + let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); + tasks.push(part_fut); + } + } + + /// Accumulates bytes into `buffer`, flushing 5 MiB parts when full. + fn append_bytes( + buffer: &mut BytesMut, + tasks: &mut Vec, + mut bs: Bytes, + upload: &mut Box, + ) { + while !bs.is_empty() { + let remaining = MIN_PART_SIZE.saturating_sub(buffer.len()); + if remaining == 0 { + Self::flush_buffer(buffer, tasks, upload); + continue; + } + + if bs.len() < remaining { + buffer.extend_from_slice(&bs); + return; + } + let chunk = bs.split_to(remaining); + buffer.extend_from_slice(&chunk); + Self::flush_buffer(buffer, tasks, upload); + } + } } impl Drop for ObjectStoreWriter { fn drop(&mut self) { - if let Some(writer) = self.writer.take() { + if let Some(mut upload) = self.upload.take() { if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { - let _ = writer.abort().await; + if let Err(e) = upload.abort().await { + tracing::warn!( + error = %e, + "Failed to abort multipart upload on drop" + ); + } }); } else { tracing::warn!( @@ -346,20 +462,55 @@ impl Drop for ObjectStoreWriter { #[async_trait] impl FileWrite for ObjectStoreWriter { async fn write(&mut self, bs: Bytes) -> Result<()> { - let writer = self - .writer - .as_mut() - .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Writer has already been closed"))?; - writer.put(bs); + let upload = self.upload.as_mut().ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "Writer has already been closed", + ) + })?; + Self::append_bytes(&mut self.buffer, &mut self.tasks, bs, upload); Ok(()) } async fn close(&mut self) -> Result<()> { - let writer = self - .writer - .take() - .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Writer has already been closed"))?; - writer.finish().await.map_err(from_object_store_error)?; + let mut upload = self.upload.take().ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "Writer has already been closed", + ) + })?; + + // Flush any remaining buffered data, or emit an empty part if 0 parts + // have been sent (S3 requires at least 1 part to complete a multipart upload). + if !self.buffer.is_empty() || self.tasks.is_empty() { + let part_data = std::mem::take(&mut self.buffer).freeze(); + let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); + self.tasks.push(part_fut); + } + + // Await all in-flight part uploads; abort if any part fails. + let tasks = std::mem::take(&mut self.tasks); + if let Err(e) = futures::future::try_join_all(tasks).await { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after part upload failure" + ); + } + return Err(from_object_store_error(e)); + } + + // Complete multipart upload; abort if S3 rejects completion. + if let Err(e) = upload.complete().await { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after complete failure" + ); + } + return Err(from_object_store_error(e)); + } + Ok(()) } } @@ -372,10 +523,46 @@ mod tests { fn make_s3_storage() -> ObjectStoreStorage { ObjectStoreStorage::S3(S3Storage { config: Arc::new(S3Config::default()), + delete_batch_size: DEFAULT_DELETE_BATCH_SIZE, store_cache: Arc::new(DashMap::new()), }) } + #[test] + fn test_parse_delete_batch_size_default() { + let config = StorageConfig::new(); + assert_eq!(parse_delete_batch_size(&config), DEFAULT_DELETE_BATCH_SIZE); + } + + #[test] + fn test_parse_delete_batch_size_custom() { + let mut props = std::collections::HashMap::new(); + props.insert(S3_DELETE_BATCH_SIZE.to_string(), "250".to_string()); + let config = StorageConfig::from_props(props); + assert_eq!(parse_delete_batch_size(&config), 250); + } + + #[test] + fn test_parse_delete_batch_size_exceeds_max() { + let mut props = std::collections::HashMap::new(); + props.insert(S3_DELETE_BATCH_SIZE.to_string(), "5000".to_string()); + let config = StorageConfig::from_props(props); + assert_eq!(parse_delete_batch_size(&config), 5000); + } + + #[test] + fn test_parse_delete_batch_size_invalid_fallback() { + let mut props = std::collections::HashMap::new(); + props.insert(S3_DELETE_BATCH_SIZE.to_string(), "invalid_number".to_string()); + let config = StorageConfig::from_props(props); + assert_eq!(parse_delete_batch_size(&config), DEFAULT_DELETE_BATCH_SIZE); + + let mut props_zero = std::collections::HashMap::new(); + props_zero.insert(S3_DELETE_BATCH_SIZE.to_string(), "0".to_string()); + let config_zero = StorageConfig::from_props(props_zero); + assert_eq!(parse_delete_batch_size(&config_zero), DEFAULT_DELETE_BATCH_SIZE); + } + #[cfg(feature = "object_store-s3")] #[test] fn test_store_cache_reuses_store() { @@ -450,22 +637,194 @@ mod tests { assert_eq!(file_io.config(), deserialized.config()); } + #[tokio::test] + async fn test_reader_empty_range() { + let store = Arc::new(object_store::memory::InMemory::new()); + let path = ObjectStorePath::from("data/test.parquet"); + let reader = ObjectStoreReader { store, path }; + + let result = reader.read(0..0).await.unwrap(); + assert!(result.is_empty()); + + let result = reader.read(10..10).await.unwrap(); + assert!(result.is_empty()); + } + #[tokio::test] async fn test_writer_already_closed_errors() { - let mut writer = ObjectStoreWriter { writer: None }; + let mut writer = ObjectStoreWriter { + upload: None, + buffer: BytesMut::new(), + tasks: Vec::new(), + }; let write_err = writer.write(Bytes::from_static(b"data")).await.unwrap_err(); - assert_eq!(write_err.kind(), ErrorKind::Unexpected); + assert_eq!(write_err.kind(), ErrorKind::PreconditionFailed); assert_eq!(write_err.message(), "Writer has already been closed"); let close_err = writer.close().await.unwrap_err(); - assert_eq!(close_err.kind(), ErrorKind::Unexpected); + assert_eq!(close_err.kind(), ErrorKind::PreconditionFailed); assert_eq!(close_err.message(), "Writer has already been closed"); } + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use object_store::PutResult; + + #[derive(Debug)] + struct MockMultipartUpload { + parts_count: Arc, + completed: Arc, + aborted: Arc, + fail_part: bool, + fail_complete: bool, + } + + #[async_trait] + impl MultipartUpload for MockMultipartUpload { + fn put_part(&mut self, _data: PutPayload) -> UploadPart { + let count = self.parts_count.clone(); + let fail = self.fail_part; + Box::pin(async move { + if fail { + Err(object_store::Error::Generic { + store: "mock", + source: "part failed".into(), + }) + } else { + count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + }) + } + + async fn complete(&mut self) -> object_store::Result { + if self.fail_complete { + Err(object_store::Error::Generic { + store: "mock", + source: "complete failed".into(), + }) + } else { + self.completed.store(true, Ordering::SeqCst); + Ok(PutResult { + e_tag: None, + version: None, + }) + } + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.aborted.store(true, Ordering::SeqCst); + Ok(()) + } + } + + fn make_mock_writer( + fail_part: bool, + fail_complete: bool, + ) -> ( + ObjectStoreWriter, + Arc, + Arc, + Arc, + ) { + let parts_count = Arc::new(AtomicUsize::new(0)); + let completed = Arc::new(AtomicBool::new(false)); + let aborted = Arc::new(AtomicBool::new(false)); + + let upload = Box::new(MockMultipartUpload { + parts_count: parts_count.clone(), + completed: completed.clone(), + aborted: aborted.clone(), + fail_part, + fail_complete, + }); + + let writer = ObjectStoreWriter { + upload: Some(upload), + buffer: BytesMut::new(), + tasks: Vec::new(), + }; + + (writer, parts_count, completed, aborted) + } + #[test] fn test_writer_drop_outside_tokio_warns_and_does_not_panic() { - // A plain synchronous #[test] runs on an OS thread outside a Tokio runtime context - let writer = ObjectStoreWriter { writer: None }; + // A plain synchronous #[test] runs on an OS thread outside a Tokio runtime context. + // Passing an active upload verifies that Handle::try_current() failing logs a warning without panicking. + let (writer, _parts, completed, aborted) = make_mock_writer(false, false); drop(writer); + + assert!(!completed.load(Ordering::SeqCst)); + assert!(!aborted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_writer_close_success() { + let (mut writer, parts, completed, aborted) = make_mock_writer(false, false); + writer + .write(Bytes::from(vec![0u8; 6 * 1024 * 1024])) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert_eq!(parts.load(Ordering::SeqCst), 2); // 5 MiB + 1 MiB leftover part + assert!(completed.load(Ordering::SeqCst)); + assert!(!aborted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_writer_close_empty_file_emits_one_part() { + let (mut writer, parts, completed, aborted) = make_mock_writer(false, false); + // Closing a fresh writer with 0 bytes must emit exactly 1 part for S3 validity + writer.close().await.unwrap(); + + assert_eq!(parts.load(Ordering::SeqCst), 1); + assert!(completed.load(Ordering::SeqCst)); + assert!(!aborted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_writer_part_failure_aborts() { + let (mut writer, _parts, completed, aborted) = make_mock_writer(true, false); + writer + .write(Bytes::from(vec![0u8; 6 * 1024 * 1024])) + .await + .unwrap(); + let err = writer.close().await.unwrap_err(); + + assert_eq!(err.kind(), ErrorKind::Unexpected); + assert!(!completed.load(Ordering::SeqCst)); + assert!( + aborted.load(Ordering::SeqCst), + "abort() MUST be called when part upload fails" + ); + } + + #[tokio::test] + async fn test_writer_complete_failure_aborts() { + let (mut writer, _parts, completed, aborted) = make_mock_writer(false, true); + writer.write(Bytes::from_static(b"hello")).await.unwrap(); + let err = writer.close().await.unwrap_err(); + + assert_eq!(err.kind(), ErrorKind::Unexpected); + assert!(!completed.load(Ordering::SeqCst)); + assert!( + aborted.load(Ordering::SeqCst), + "abort() MUST be called when complete() fails" + ); + } + + #[tokio::test] + async fn test_writer_drop_aborts() { + let (writer, _parts, completed, aborted) = make_mock_writer(false, false); + drop(writer); + + // Allow background tokio task to run + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!(!completed.load(Ordering::SeqCst)); + assert!( + aborted.load(Ordering::SeqCst), + "abort() MUST be called on drop" + ); } } diff --git a/crates/storage/object_store/src/s3.rs b/crates/storage/object_store/src/s3.rs index 69b0f066fa..7547ccdaa8 100644 --- a/crates/storage/object_store/src/s3.rs +++ b/crates/storage/object_store/src/s3.rs @@ -46,7 +46,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { "s3" | "s3a" | "s3n" => {} _ => { return Err(Error::new( - ErrorKind::DataInvalid, + ErrorKind::FeatureUnsupported, format!("Unsupported S3 scheme: {scheme} in url: {path}"), )); } @@ -74,7 +74,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { .with_source(e) })?; - let relative = url.path().trim_start_matches('/'); + let relative = url.path().strip_prefix('/').unwrap_or(url.path()); Ok(ParsedS3Url { scheme: scheme.to_string(), @@ -87,7 +87,7 @@ pub(crate) fn parse_s3_url(path: &str) -> Result { fn parse_s3_config_key(key: &str) -> Result { AmazonS3ConfigKey::from_str(key).map_err(|e| { Error::new( - ErrorKind::Unexpected, + ErrorKind::DataInvalid, format!("Failed to parse S3 config key: {key}"), ) .with_source(e) @@ -212,7 +212,8 @@ mod tests { #[test] fn test_parse_s3_url_unsupported_scheme() { - assert!(parse_s3_url("gs://my-bucket/file.parquet").is_err()); + let err = parse_s3_url("gs://my-bucket/file.parquet").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); } #[test] @@ -255,6 +256,14 @@ mod tests { assert_eq!(parsed.relative, "path%20with%20spaces/file.parquet"); } + #[test] + fn test_parse_s3_url_double_slash_path() { + let parsed = parse_s3_url("s3://my-bucket//path/to/file.parquet").unwrap(); + assert_eq!(parsed.scheme, "s3"); + assert_eq!(parsed.bucket, "my-bucket"); + assert_eq!(parsed.relative, "/path/to/file.parquet"); + } + #[test] fn test_parse_s3_url_empty_bucket() { assert!(parse_s3_url("s3:///path/to/file.parquet").is_err()); @@ -265,7 +274,7 @@ mod tests { fn test_parse_s3_config_key() { assert!(parse_s3_config_key("aws_server_side_encryption").is_ok()); let err = parse_s3_config_key("invalid_config_key_foo_bar").unwrap_err(); - assert_eq!(err.kind(), ErrorKind::Unexpected); + assert_eq!(err.kind(), ErrorKind::DataInvalid); } #[test] diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index 609bd6446a..df25f0784b 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -195,31 +195,6 @@ mod tests { assert!(!file_io.exists(&file_path).await.unwrap()); } - #[tokio::test] - async fn test_file_io_s3_multipart_writer_drop_aborts() { - let file_io = get_file_io().await; - let file_path = format!( - "s3://bucket1/{}", - normalize_test_name_with_parts!("test_file_io_s3_multipart_writer_drop_aborts") - ); - let _ = file_io.delete(&file_path).await; - - let output_file = file_io.new_output(&file_path).unwrap(); - let mut writer = output_file.writer().await.unwrap(); - writer - .write(Bytes::from_static(b"uncommitted chunk")) - .await - .unwrap(); - - // Dropping writer without close() should abort the multipart upload - drop(writer); - - // Give background abort task a moment to execute - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - assert!(!file_io.exists(&file_path).await.unwrap()); - } - #[tokio::test] async fn test_file_io_s3_percent_encoded_bucket() { let file_io = get_file_io().await; From cc6edad9f6cbb49ba2c18530c3a0473f7eb72426 Mon Sep 17 00:00:00 2001 From: dron Date: Thu, 24 Sep 2026 21:54:04 +0530 Subject: [PATCH 12/13] fix(storage): bound multipart upload concurrency and fail-fast during write --- crates/storage/object_store/src/lib.rs | 129 ++++++++++++++++++++----- 1 file changed, 105 insertions(+), 24 deletions(-) diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 67a86279dd..6742089d02 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -34,7 +34,7 @@ use std::sync::Arc; use async_trait::async_trait; use bytes::{Bytes, BytesMut}; use dashmap::DashMap; -use futures::stream::BoxStream; +use futures::stream::{BoxStream, FuturesUnordered}; use futures::{StreamExt, TryStreamExt}; #[cfg(feature = "object_store-s3")] use iceberg::io::S3Config; @@ -291,7 +291,9 @@ impl Storage for ObjectStoreStorage { Ok(Box::new(ObjectStoreWriter { upload: Some(upload), buffer: BytesMut::new(), - tasks: Vec::new(), + tasks: FuturesUnordered::new(), + parts_submitted: 0, + bytes_written: 0, })) } @@ -392,38 +394,46 @@ impl FileRead for ObjectStoreReader { /// Minimum part size for S3 multipart upload (5 MiB). const MIN_PART_SIZE: usize = 5 * 1024 * 1024; +/// Default maximum concurrent in-flight part uploads. +const MAX_CONCURRENT_PART_UPLOADS: usize = 8; + /// Writer that implements `FileWrite` using `object_store` multipart upload. struct ObjectStoreWriter { upload: Option>, buffer: BytesMut, - tasks: Vec, + tasks: FuturesUnordered, + parts_submitted: usize, + bytes_written: u64, } impl ObjectStoreWriter { /// Flushes any buffered bytes as an in-flight part upload to S3. fn flush_buffer( buffer: &mut BytesMut, - tasks: &mut Vec, + tasks: &mut FuturesUnordered, + parts_submitted: &mut usize, upload: &mut Box, ) { if !buffer.is_empty() { let part_data = std::mem::take(buffer).freeze(); let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); tasks.push(part_fut); + *parts_submitted += 1; } } /// Accumulates bytes into `buffer`, flushing 5 MiB parts when full. fn append_bytes( buffer: &mut BytesMut, - tasks: &mut Vec, + tasks: &mut FuturesUnordered, + parts_submitted: &mut usize, mut bs: Bytes, upload: &mut Box, ) { while !bs.is_empty() { let remaining = MIN_PART_SIZE.saturating_sub(buffer.len()); if remaining == 0 { - Self::flush_buffer(buffer, tasks, upload); + Self::flush_buffer(buffer, tasks, parts_submitted, upload); continue; } @@ -433,7 +443,7 @@ impl ObjectStoreWriter { } let chunk = bs.split_to(remaining); buffer.extend_from_slice(&chunk); - Self::flush_buffer(buffer, tasks, upload); + Self::flush_buffer(buffer, tasks, parts_submitted, upload); } } } @@ -468,11 +478,36 @@ impl FileWrite for ObjectStoreWriter { "Writer has already been closed", ) })?; - Self::append_bytes(&mut self.buffer, &mut self.tasks, bs, upload); + self.bytes_written += bs.len() as u64; + Self::append_bytes( + &mut self.buffer, + &mut self.tasks, + &mut self.parts_submitted, + bs, + upload, + ); + + // Throttle in-flight uploads: fail fast if any part upload fails + while self.tasks.len() >= MAX_CONCURRENT_PART_UPLOADS { + if let Some(res) = self.tasks.next().await { + if let Err(e) = res { + if let Some(mut upload) = self.upload.take() { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after part upload failure" + ); + } + } + return Err(from_object_store_error(e)); + } + } + } + Ok(()) } - async fn close(&mut self) -> Result<()> { + async fn close(&mut self) -> Result { let mut upload = self.upload.take().ok_or_else(|| { Error::new( ErrorKind::PreconditionFailed, @@ -482,22 +517,24 @@ impl FileWrite for ObjectStoreWriter { // Flush any remaining buffered data, or emit an empty part if 0 parts // have been sent (S3 requires at least 1 part to complete a multipart upload). - if !self.buffer.is_empty() || self.tasks.is_empty() { + if !self.buffer.is_empty() || self.parts_submitted == 0 { let part_data = std::mem::take(&mut self.buffer).freeze(); let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); self.tasks.push(part_fut); + self.parts_submitted += 1; } // Await all in-flight part uploads; abort if any part fails. - let tasks = std::mem::take(&mut self.tasks); - if let Err(e) = futures::future::try_join_all(tasks).await { - if let Err(abort_err) = upload.abort().await { - tracing::warn!( - error = %abort_err, - "Failed to abort multipart upload after part upload failure" - ); + while let Some(res) = self.tasks.next().await { + if let Err(e) = res { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after part upload failure" + ); + } + return Err(from_object_store_error(e)); } - return Err(from_object_store_error(e)); } // Complete multipart upload; abort if S3 rejects completion. @@ -511,7 +548,9 @@ impl FileWrite for ObjectStoreWriter { return Err(from_object_store_error(e)); } - Ok(()) + Ok(FileMetadata { + size: self.bytes_written, + }) } } @@ -655,13 +694,17 @@ mod tests { let mut writer = ObjectStoreWriter { upload: None, buffer: BytesMut::new(), - tasks: Vec::new(), + tasks: FuturesUnordered::new(), + parts_submitted: 0, + bytes_written: 0, }; let write_err = writer.write(Bytes::from_static(b"data")).await.unwrap_err(); assert_eq!(write_err.kind(), ErrorKind::PreconditionFailed); assert_eq!(write_err.message(), "Writer has already been closed"); - let close_err = writer.close().await.unwrap_err(); + let close_res = writer.close().await; + assert!(close_res.is_err()); + let close_err = close_res.err().unwrap(); assert_eq!(close_err.kind(), ErrorKind::PreconditionFailed); assert_eq!(close_err.message(), "Writer has already been closed"); } @@ -741,7 +784,9 @@ mod tests { let writer = ObjectStoreWriter { upload: Some(upload), buffer: BytesMut::new(), - tasks: Vec::new(), + tasks: FuturesUnordered::new(), + parts_submitted: 0, + bytes_written: 0, }; (writer, parts_count, completed, aborted) @@ -790,7 +835,9 @@ mod tests { .write(Bytes::from(vec![0u8; 6 * 1024 * 1024])) .await .unwrap(); - let err = writer.close().await.unwrap_err(); + let close_res = writer.close().await; + assert!(close_res.is_err()); + let err = close_res.err().unwrap(); assert_eq!(err.kind(), ErrorKind::Unexpected); assert!(!completed.load(Ordering::SeqCst)); @@ -804,7 +851,9 @@ mod tests { async fn test_writer_complete_failure_aborts() { let (mut writer, _parts, completed, aborted) = make_mock_writer(false, true); writer.write(Bytes::from_static(b"hello")).await.unwrap(); - let err = writer.close().await.unwrap_err(); + let close_res = writer.close().await; + assert!(close_res.is_err()); + let err = close_res.err().unwrap(); assert_eq!(err.kind(), ErrorKind::Unexpected); assert!(!completed.load(Ordering::SeqCst)); @@ -814,6 +863,38 @@ mod tests { ); } + #[tokio::test] + async fn test_writer_bounded_concurrency_past_limit() { + let (mut writer, parts, completed, aborted) = make_mock_writer(false, false); + // Write 10 parts (50 MiB) which exceeds MAX_CONCURRENT_PART_UPLOADS (8) + writer + .write(Bytes::from(vec![0u8; 10 * 5 * 1024 * 1024])) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert_eq!(parts.load(Ordering::SeqCst), 10); + assert!(completed.load(Ordering::SeqCst)); + assert!(!aborted.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_writer_fail_fast_during_write() { + let (mut writer, _parts, completed, aborted) = make_mock_writer(true, false); + // Writing past MAX_CONCURRENT_PART_UPLOADS triggers in-flight task awaiting inside write() + let err = writer + .write(Bytes::from(vec![0u8; 10 * 5 * 1024 * 1024])) + .await + .unwrap_err(); + + assert_eq!(err.kind(), ErrorKind::Unexpected); + assert!(!completed.load(Ordering::SeqCst)); + assert!( + aborted.load(Ordering::SeqCst), + "abort() MUST be called immediately on write() failure" + ); + } + #[tokio::test] async fn test_writer_drop_aborts() { let (writer, _parts, completed, aborted) = make_mock_writer(false, false); From 67f3677f690bb74e6894da7d06d99bd04a2bfd25 Mon Sep 17 00:00:00 2001 From: dron Date: Thu, 24 Sep 2026 21:54:04 +0530 Subject: [PATCH 13/13] fix(storage): bound multipart upload concurrency and fail-fast during write --- crates/storage/object_store/src/lib.rs | 186 ++++++++++++------ .../object_store/tests/file_io_s3_test.rs | 75 +++---- 2 files changed, 155 insertions(+), 106 deletions(-) diff --git a/crates/storage/object_store/src/lib.rs b/crates/storage/object_store/src/lib.rs index 6742089d02..aaae42cec7 100644 --- a/crates/storage/object_store/src/lib.rs +++ b/crates/storage/object_store/src/lib.rs @@ -78,30 +78,21 @@ fn from_object_store_error(e: object_store::Error) -> Error { Error::new(kind, msg).with_source(e) } -/// Property key for configuring S3 bulk delete batch size. -pub const S3_DELETE_BATCH_SIZE: &str = "s3.delete-batch-size"; -/// Default batch size for S3 bulk deletions (matches AWS S3 DeleteObjects max). -pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000; +/// Property key for configuring S3 bulk delete batch size (matches Iceberg Java specification). +pub const S3_DELETE_BATCH_SIZE: &str = "s3.delete.batch-size"; +/// Default batch size for S3 bulk deletions (matches Iceberg Java specification). +pub const DEFAULT_DELETE_BATCH_SIZE: usize = 250; /// Maximum batch size allowed by the AWS S3 DeleteObjects API specification. pub const S3_MAX_DELETE_BATCH_SIZE: usize = 1000; fn parse_delete_batch_size(config: &StorageConfig) -> usize { if let Some(val) = config.get(S3_DELETE_BATCH_SIZE) { match val.parse::() { - Ok(parsed) if parsed > 0 => { - if parsed > S3_MAX_DELETE_BATCH_SIZE { - tracing::warn!( - configured = parsed, - limit = S3_MAX_DELETE_BATCH_SIZE, - "Configured s3.delete-batch-size exceeds AWS S3 hard limit of 1000; requests may fail with MalformedXML" - ); - } - parsed - } + Ok(parsed) if parsed > 0 => parsed.clamp(1, S3_MAX_DELETE_BATCH_SIZE), _ => { tracing::warn!( val = %val, - "Invalid s3.delete-batch-size; falling back to default 1000" + "Invalid s3.delete.batch-size; falling back to default 250" ); DEFAULT_DELETE_BATCH_SIZE } @@ -163,7 +154,7 @@ type StoreCache = Arc>>; pub struct S3Storage { config: Arc, #[serde(default = "default_delete_batch_size")] - pub delete_batch_size: usize, + delete_batch_size: usize, #[serde(skip, default)] store_cache: StoreCache, } @@ -407,44 +398,71 @@ struct ObjectStoreWriter { } impl ObjectStoreWriter { - /// Flushes any buffered bytes as an in-flight part upload to S3. - fn flush_buffer( + /// Best-effort abort on upload error, converting the underlying + /// object_store error to an iceberg Error while logging abort failures. + async fn abort_and_wrap( + upload: &mut Box, + e: object_store::Error, + context: &'static str, + ) -> Error { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after {context}" + ); + } + from_object_store_error(e) + } + + /// Flushes any buffered bytes as an in-flight part upload to S3, + /// throttling concurrency so that at most `MAX_CONCURRENT_PART_UPLOADS` + /// uploads are in flight at any given time. + async fn flush_buffer( buffer: &mut BytesMut, tasks: &mut FuturesUnordered, parts_submitted: &mut usize, upload: &mut Box, - ) { + ) -> Result<()> { if !buffer.is_empty() { + // Drain in-flight tasks down below cap before queuing the next chunk + while tasks.len() >= MAX_CONCURRENT_PART_UPLOADS { + if let Some(Err(e)) = tasks.next().await { + return Err(Self::abort_and_wrap(upload, e, "part upload failure").await); + } + } + let part_data = std::mem::take(buffer).freeze(); let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); tasks.push(part_fut); *parts_submitted += 1; } + Ok(()) } /// Accumulates bytes into `buffer`, flushing 5 MiB parts when full. - fn append_bytes( + async fn append_bytes( buffer: &mut BytesMut, tasks: &mut FuturesUnordered, parts_submitted: &mut usize, mut bs: Bytes, upload: &mut Box, - ) { + ) -> Result<()> { while !bs.is_empty() { let remaining = MIN_PART_SIZE.saturating_sub(buffer.len()); if remaining == 0 { - Self::flush_buffer(buffer, tasks, parts_submitted, upload); + Self::flush_buffer(buffer, tasks, parts_submitted, upload).await?; continue; } if bs.len() < remaining { buffer.extend_from_slice(&bs); - return; + return Ok(()); } let chunk = bs.split_to(remaining); buffer.extend_from_slice(&chunk); - Self::flush_buffer(buffer, tasks, parts_submitted, upload); + Self::flush_buffer(buffer, tasks, parts_submitted, upload).await?; } + Ok(()) } } @@ -485,26 +503,8 @@ impl FileWrite for ObjectStoreWriter { &mut self.parts_submitted, bs, upload, - ); - - // Throttle in-flight uploads: fail fast if any part upload fails - while self.tasks.len() >= MAX_CONCURRENT_PART_UPLOADS { - if let Some(res) = self.tasks.next().await { - if let Err(e) = res { - if let Some(mut upload) = self.upload.take() { - if let Err(abort_err) = upload.abort().await { - tracing::warn!( - error = %abort_err, - "Failed to abort multipart upload after part upload failure" - ); - } - } - return Err(from_object_store_error(e)); - } - } - } - - Ok(()) + ) + .await } async fn close(&mut self) -> Result { @@ -527,25 +527,13 @@ impl FileWrite for ObjectStoreWriter { // Await all in-flight part uploads; abort if any part fails. while let Some(res) = self.tasks.next().await { if let Err(e) = res { - if let Err(abort_err) = upload.abort().await { - tracing::warn!( - error = %abort_err, - "Failed to abort multipart upload after part upload failure" - ); - } - return Err(from_object_store_error(e)); + return Err(Self::abort_and_wrap(&mut upload, e, "part upload failure").await); } } // Complete multipart upload; abort if S3 rejects completion. if let Err(e) = upload.complete().await { - if let Err(abort_err) = upload.abort().await { - tracing::warn!( - error = %abort_err, - "Failed to abort multipart upload after complete failure" - ); - } - return Err(from_object_store_error(e)); + return Err(Self::abort_and_wrap(&mut upload, e, "complete failure").await); } Ok(FileMetadata { @@ -576,9 +564,9 @@ mod tests { #[test] fn test_parse_delete_batch_size_custom() { let mut props = std::collections::HashMap::new(); - props.insert(S3_DELETE_BATCH_SIZE.to_string(), "250".to_string()); + props.insert(S3_DELETE_BATCH_SIZE.to_string(), "500".to_string()); let config = StorageConfig::from_props(props); - assert_eq!(parse_delete_batch_size(&config), 250); + assert_eq!(parse_delete_batch_size(&config), 500); } #[test] @@ -586,20 +574,26 @@ mod tests { let mut props = std::collections::HashMap::new(); props.insert(S3_DELETE_BATCH_SIZE.to_string(), "5000".to_string()); let config = StorageConfig::from_props(props); - assert_eq!(parse_delete_batch_size(&config), 5000); + assert_eq!(parse_delete_batch_size(&config), S3_MAX_DELETE_BATCH_SIZE); } #[test] fn test_parse_delete_batch_size_invalid_fallback() { let mut props = std::collections::HashMap::new(); - props.insert(S3_DELETE_BATCH_SIZE.to_string(), "invalid_number".to_string()); + props.insert( + S3_DELETE_BATCH_SIZE.to_string(), + "invalid_number".to_string(), + ); let config = StorageConfig::from_props(props); assert_eq!(parse_delete_batch_size(&config), DEFAULT_DELETE_BATCH_SIZE); let mut props_zero = std::collections::HashMap::new(); props_zero.insert(S3_DELETE_BATCH_SIZE.to_string(), "0".to_string()); let config_zero = StorageConfig::from_props(props_zero); - assert_eq!(parse_delete_batch_size(&config_zero), DEFAULT_DELETE_BATCH_SIZE); + assert_eq!( + parse_delete_batch_size(&config_zero), + DEFAULT_DELETE_BATCH_SIZE + ); } #[cfg(feature = "object_store-s3")] @@ -710,6 +704,7 @@ mod tests { } use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use object_store::PutResult; #[derive(Debug)] @@ -908,4 +903,69 @@ mod tests { "abort() MUST be called on drop" ); } + + #[tokio::test] + async fn test_writer_bounded_concurrency_in_flight_cap() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + + #[derive(Debug)] + struct TrackingMockUpload { + in_flight: Arc, + max_in_flight: Arc, + } + + #[async_trait] + impl MultipartUpload for TrackingMockUpload { + fn put_part(&mut self, _data: PutPayload) -> UploadPart { + let in_flight = self.in_flight.clone(); + let max_in_flight = self.max_in_flight.clone(); + Box::pin(async move { + let current = in_flight.fetch_add(1, Ordering::SeqCst) + 1; + max_in_flight.fetch_max(current, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + in_flight.fetch_sub(1, Ordering::SeqCst); + Ok(()) + }) + } + + async fn complete(&mut self) -> object_store::Result { + Ok(PutResult { + e_tag: None, + version: None, + }) + } + + async fn abort(&mut self) -> object_store::Result<()> { + Ok(()) + } + } + + let upload = Box::new(TrackingMockUpload { + in_flight: in_flight.clone(), + max_in_flight: max_in_flight.clone(), + }); + + let mut writer = ObjectStoreWriter { + upload: Some(upload), + buffer: BytesMut::new(), + tasks: FuturesUnordered::new(), + parts_submitted: 0, + bytes_written: 0, + }; + + // Write 20 parts (100 MiB) in a single write call + writer + .write(Bytes::from(vec![0u8; 20 * 5 * 1024 * 1024])) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert!( + max_in_flight.load(Ordering::SeqCst) <= MAX_CONCURRENT_PART_UPLOADS, + "max concurrent in-flight uploads exceeded {} (was {})", + MAX_CONCURRENT_PART_UPLOADS, + max_in_flight.load(Ordering::SeqCst) + ); + } } diff --git a/crates/storage/object_store/tests/file_io_s3_test.rs b/crates/storage/object_store/tests/file_io_s3_test.rs index df25f0784b..3b90669143 100644 --- a/crates/storage/object_store/tests/file_io_s3_test.rs +++ b/crates/storage/object_store/tests/file_io_s3_test.rs @@ -31,16 +31,16 @@ mod tests { S3_SECRET_ACCESS_KEY, S3_SSE_KEY, S3_SSE_TYPE, }; use iceberg_storage_object_store::ObjectStoreStorageFactory; - use iceberg_test_utils::{get_minio_endpoint, normalize_test_name_with_parts, set_up}; + use iceberg_test_utils::{get_object_store_endpoint, normalize_test_name_with_parts, set_up}; async fn get_file_io() -> FileIO { set_up(); - let minio_endpoint = get_minio_endpoint(); + let endpoint = get_object_store_endpoint(); FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) .with_props(vec![ - (S3_ENDPOINT, minio_endpoint), + (S3_ENDPOINT, endpoint), (S3_ACCESS_KEY_ID, "admin".to_string()), (S3_SECRET_ACCESS_KEY, "password".to_string()), (S3_REGION, "us-east-1".to_string()), @@ -230,10 +230,22 @@ mod tests { assert!(!file_io.exists(&file_path).await.unwrap()); } + fn assert_minio_kms_rejection(err: &iceberg::Error) { + let err_str = err.to_string(); + assert!( + err_str.contains("501") + || err_str.contains("NotImplemented") + || err_str.contains("KMS is not configured") + || err_str.contains("kms") + || err_str.contains("ServerNotInitialized"), + "Expected KMS rejection from KES-less MinIO confirming SSE header was sent, got: {err_str}" + ); + } + #[tokio::test] async fn test_file_io_s3_sse_kms_default() { set_up(); - let endpoint = get_minio_endpoint(); + let endpoint = get_object_store_endpoint(); let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) .with_props(vec![ @@ -252,33 +264,23 @@ mod tests { ); let _ = file_io.delete(&file_path).await; - match file_io + // MinIO without KES does not configure KMS; asserting that MinIO returns 501 / KMS error + // verifies that the SSE-KMS header was attached (an unencrypted PUT would succeed). + let err = file_io .new_output(&file_path) .unwrap() .write(Bytes::from_static(b"kms-encrypted-data")) .await - { - Ok(_) => { - assert!(file_io.exists(&file_path).await.unwrap()); - let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); - assert_eq!(content, Bytes::from_static(b"kms-encrypted-data")); - file_io.delete(&file_path).await.unwrap(); - } - Err(e) - if e.to_string().contains("501") - || e.to_string().contains("NotImplemented") - || e.to_string().contains("KMS is not configured") => - { - // MinIO without KES does not configure KMS; passing 501 verifies header was sent - } - Err(e) => panic!("Unexpected error: {e:?}"), - } + .expect_err( + "MinIO without KES must reject SSE-KMS requests, proving SSE header was sent", + ); + assert_minio_kms_rejection(&err); } #[tokio::test] async fn test_file_io_s3_sse_kms_custom_key() { set_up(); - let endpoint = get_minio_endpoint(); + let endpoint = get_object_store_endpoint(); let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) .with_props(vec![ @@ -302,30 +304,17 @@ mod tests { ); let _ = file_io.delete(&file_path).await; - match file_io + // MinIO without KES does not configure KMS; asserting that MinIO returns 501 / KMS error + // verifies that the SSE-KMS custom key header was attached. + let err = file_io .new_output(&file_path) .unwrap() .write(Bytes::from_static(b"kms-custom-key-encrypted-data")) .await - { - Ok(_) => { - assert!(file_io.exists(&file_path).await.unwrap()); - let content = file_io.new_input(&file_path).unwrap().read().await.unwrap(); - assert_eq!( - content, - Bytes::from_static(b"kms-custom-key-encrypted-data") - ); - file_io.delete(&file_path).await.unwrap(); - } - Err(e) - if e.to_string().contains("501") - || e.to_string().contains("NotImplemented") - || e.to_string().contains("KMS is not configured") => - { - // MinIO without KES does not configure KMS; passing 501 verifies header was sent - } - Err(e) => panic!("Unexpected error: {e:?}"), - } + .expect_err( + "MinIO without KES must reject SSE-KMS requests, proving SSE custom key header was sent", + ); + assert_minio_kms_rejection(&err); } /// Writes 12 MiB (3 × 4 MiB chunks) to exercise real S3 multipart uploads @@ -478,7 +467,7 @@ mod tests { #[tokio::test] async fn test_file_io_s3_sse_s3_aes256() { set_up(); - let endpoint = get_minio_endpoint(); + let endpoint = get_object_store_endpoint(); let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) .with_props(vec![