laskoviymishka commented on code in PR #3165: URL: https://github.com/apache/iceberg-rust/pull/3165#discussion_r4242741596
########## crates/storage/object_store/src/lib.rs: ########## @@ -0,0 +1,1245 @@ +// 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`] 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, +//! 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, BytesMut}; +use dashmap::DashMap; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +#[cfg(feature = "object_store-s3")] +use iceberg::io::{S3Config, S3_DELETE_BATCH_SIZE, S3_MULTIPART_PART_SIZE_BYTES}; +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::{MultipartUpload, ObjectStore, ObjectStoreExt, PutPayload}; +#[cfg(feature = "object_store-s3")] +use s3::{build_s3_store, parse_s3_url}; +use serde::{Deserialize, Serialize}; +use tokio::task::JoinSet; + +/// 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 { + let (kind, msg) = match &e { + object_store::Error::NotFound { path, .. } => { + (ErrorKind::Unexpected, format!("Object not found: {path}")) + } + object_store::Error::AlreadyExists { path, .. } => ( + ErrorKind::Unexpected, + format!("Object already exists: {path}"), + ), + object_store::Error::PermissionDenied { path, .. } => { + (ErrorKind::Unexpected, format!("Permission denied: {path}")) + } + object_store::Error::Unauthenticated { path, .. } => { + (ErrorKind::Unexpected, 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) +} + +/// Default batch size for S3 bulk deletions (matches Iceberg Java specification). +pub(crate) const DEFAULT_DELETE_BATCH_SIZE: usize = 250; +/// Maximum batch size allowed by the AWS S3 DeleteObjects API specification. +pub(crate) const S3_MAX_DELETE_BATCH_SIZE: usize = 1000; +/// Default part size for S3 multipart uploads (32 MiB, matches Iceberg Java specification). +pub(crate) const DEFAULT_PART_SIZE_BYTES: usize = 32 * 1024 * 1024; +/// Minimum part size allowed by the AWS S3 MultipartUpload API specification (5 MiB). +pub(crate) const MIN_PART_SIZE_BYTES: usize = 5 * 1024 * 1024; + +fn parse_delete_batch_size(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_DELETE_BATCH_SIZE) { + match val.parse::<usize>() { + 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 250" + ); + DEFAULT_DELETE_BATCH_SIZE + } + } + } else { + DEFAULT_DELETE_BATCH_SIZE + } +} + +fn default_delete_batch_size() -> usize { + DEFAULT_DELETE_BATCH_SIZE +} + +fn parse_part_size_bytes(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_MULTIPART_PART_SIZE_BYTES) { + match val.parse::<usize>() { + Ok(parsed) if parsed >= MIN_PART_SIZE_BYTES => parsed, + _ => { + tracing::warn!( + val = %val, + "Invalid s3.multipart.part-size-bytes (must be >= 5 MiB); falling back to default 32 MiB" + ); + DEFAULT_PART_SIZE_BYTES + } + } + } else { + DEFAULT_PART_SIZE_BYTES + } +} + +fn default_part_size_bytes() -> usize { + DEFAULT_PART_SIZE_BYTES +} + +/// 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 +/// backed by the `object_store` crate. +#[non_exhaustive] +#[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<Arc<dyn Storage>> { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorageFactory::S3 => { + let s3_config = S3Config::try_from(config)?; + let delete_batch_size = parse_delete_batch_size(config); + let part_size_bytes = parse_part_size_bytes(config); + tracing::info!( + batch_size = delete_batch_size, + part_size_bytes = part_size_bytes, + "Initialized S3 storage with delete batch size {} and multipart part size {} bytes (configure via '{}' and '{}' to adjust)", + delete_batch_size, + part_size_bytes, + S3_DELETE_BATCH_SIZE, + S3_MULTIPART_PART_SIZE_BYTES + ); + Ok(Arc::new(ObjectStoreStorage::S3(S3Storage { + config: Arc::new(s3_config), + delete_batch_size, + part_size_bytes, + store_cache: Arc::new(DashMap::new()), + }))) + } + } + } +} + +type StoreCache = Arc<DashMap<String, Arc<dyn ObjectStore>>>; + +/// `object_store` S3 storage state. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct S3Storage { + config: Arc<S3Config>, + #[serde(default = "default_delete_batch_size")] + delete_batch_size: usize, + #[serde(default = "default_part_size_bytes")] + part_size_bytes: usize, + #[serde(skip, default)] + store_cache: StoreCache, +} + +/// `object_store`-based storage implementation. +/// +/// Stores are cached per bucket to avoid rebuilding the client on every operation. +#[non_exhaustive] +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum ObjectStoreStorage { + /// S3 storage variant. + #[cfg(feature = "object_store-s3")] + S3(S3Storage), +} + +#[derive(Debug)] +struct StoreAndPath { + bucket: String, + store: Arc<dyn ObjectStore>, + path: ObjectStorePath, +} + +/// Helper for batching deletions per bucket. +struct BucketBatch { + store: Arc<dyn ObjectStore>, + locations: Vec<ObjectStorePath>, +} + +impl ObjectStoreStorage { + fn delete_batch_size(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.delete_batch_size.clamp(1, S3_MAX_DELETE_BATCH_SIZE), + } + } + + fn part_size_bytes(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.part_size_bytes.max(MIN_PART_SIZE_BYTES), + } + } + + /// Get or create a cached store and extract the relative `ObjectStorePath`. + fn get_store_and_path(&self, path: &str) -> Result<StoreAndPath> { + match self { + #[cfg(feature = "object_store-s3")] + 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(); + + let object_path = ObjectStorePath::parse(parsed.relative.as_str()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid relative path: {}", parsed.relative), + ) + .with_source(e) + })?; + + Ok(StoreAndPath { + bucket: parsed.bucket, + store, + path: object_path, + }) + } + } + } +} + +#[typetag::serde(name = "ObjectStoreStorage")] +#[async_trait] +impl Storage for ObjectStoreStorage { + async fn exists(&self, path: &str) -> Result<bool> { + 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)), + } + } + + async fn metadata(&self, path: &str) -> Result<FileMetadata> { + let target = self.get_store_and_path(path)?; + let meta = target + .store + .head(&target.path) + .await + .map_err(from_object_store_error)?; + Ok(to_file_metadata(meta)) + } + + async fn read(&self, path: &str) -> Result<Bytes> { + 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<Box<dyn FileRead>> { + let target = self.get_store_and_path(path)?; + Ok(Box::new(ObjectStoreReader { + store: target.store, + path: target.path, + })) + } + + async fn write(&self, path: &str, bs: Bytes) -> Result<()> { + 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<Box<dyn FileWrite>> { + let target = self.get_store_and_path(path)?; + let upload = target + .store + .put_multipart(&target.path) + .await + .map_err(from_object_store_error)?; + Ok(Box::new(ObjectStoreWriter { + upload: Some(upload), + buffer: BytesMut::new(), + tasks: JoinSet::new(), + parts_submitted: 0, + bytes_written: 0, + part_size_bytes: self.part_size_bytes(), + })) + } + + async fn delete(&self, path: &str) -> Result<()> { + 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 target = self.get_store_and_path(path)?; + let locations = target + .store + .list(Some(&target.path)) + .map_ok(|m| m.location) + .boxed(); + target + .store + .delete_stream(locations) + .try_for_each(|_| async { Ok(()) }) + .await + .map_err(from_object_store_error)?; + Ok(()) + } + + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { + 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<String, BucketBatch> = + 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, 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(()) + } + + fn new_input(&self, path: &str) -> Result<InputFile> { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> Result<OutputFile> { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } +} + +/// Reader that implements `FileRead` using `object_store`. +struct ObjectStoreReader { + store: Arc<dyn ObjectStore>, + path: ObjectStorePath, +} + +#[async_trait] +impl FileRead for ObjectStoreReader { + async fn read(&self, range: Range<u64>) -> Result<Bytes> { + if range.is_empty() { + return Ok(Bytes::new()); + } + 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) + } +} + +/// 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<Box<dyn MultipartUpload>>, + buffer: BytesMut, + tasks: JoinSet<object_store::Result<()>>, + parts_submitted: usize, + bytes_written: u64, + part_size_bytes: usize, +} + +impl ObjectStoreWriter { + /// 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<dyn MultipartUpload>, + 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) + } + + /// Drains in-flight tasks down below `MAX_CONCURRENT_PART_UPLOADS`. + async fn drain_capacity( + tasks: &mut JoinSet<object_store::Result<()>>, + upload: &mut Box<dyn MultipartUpload>, + ) -> Result<()> { + while tasks.len() >= MAX_CONCURRENT_PART_UPLOADS { + if let Some(res) = tasks.join_next().await { + match res { + Ok(Err(e)) => { + return Err(Self::abort_and_wrap(upload, e, "part upload failure").await); + } + Err(join_err) => { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after task panic" + ); + } + return Err( + Error::new(ErrorKind::Unexpected, "Part upload task panicked") + .with_source(join_err), + ); + } + Ok(Ok(())) => {} + } + } + } + Ok(()) + } + + /// 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 JoinSet<object_store::Result<()>>, + parts_submitted: &mut usize, + upload: &mut Box<dyn MultipartUpload>, + ) -> Result<()> { + if !buffer.is_empty() { + Self::drain_capacity(tasks, upload).await?; + + let part_data = std::mem::take(buffer).freeze(); + let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); + tasks.spawn(part_fut); + *parts_submitted += 1; + } + Ok(()) + } + + /// Accumulates bytes into `buffer`, flushing parts when the configured part size is reached. + async fn append_bytes( + buffer: &mut BytesMut, + tasks: &mut JoinSet<object_store::Result<()>>, + parts_submitted: &mut usize, + part_size_bytes: usize, + mut bs: Bytes, + upload: &mut Box<dyn MultipartUpload>, + ) -> Result<()> { + while !bs.is_empty() { + let remaining = part_size_bytes.saturating_sub(buffer.len()); + if remaining == 0 { + Self::flush_buffer(buffer, tasks, parts_submitted, upload).await?; + continue; + } + + if bs.len() < remaining { + buffer.extend_from_slice(&bs); + return Ok(()); + } + let chunk = bs.split_to(remaining); + buffer.extend_from_slice(&chunk); + Self::flush_buffer(buffer, tasks, parts_submitted, upload).await?; + } + Ok(()) + } +} + +impl Drop for ObjectStoreWriter { + fn drop(&mut self) { + if let Some(mut upload) = self.upload.take() { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + if let Err(e) = upload.abort().await { + tracing::warn!( + error = %e, + "Failed to abort multipart upload on drop" + ); + } + }); + } else { + tracing::warn!( + "ObjectStoreWriter dropped outside a Tokio runtime; multipart upload abort skipped" + ); + } + } + } +} + +#[async_trait] +impl FileWrite for ObjectStoreWriter { + async fn write(&mut self, bs: Bytes) -> Result<()> { + let upload = self.upload.as_mut().ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "Writer has already been closed", + ) + })?; + self.bytes_written += bs.len() as u64; + if let Err(e) = Self::append_bytes( + &mut self.buffer, + &mut self.tasks, + &mut self.parts_submitted, + self.part_size_bytes, + bs, + upload, + ) + .await + { + let _ = self.upload.take(); + return Err(e); + } + Ok(()) + } + + async fn close(&mut self) -> Result<FileMetadata> { + let upload = self.upload.as_mut().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.parts_submitted == 0 { + if let Err(e) = Self::drain_capacity(&mut self.tasks, upload).await { + let _ = self.upload.take(); + return Err(e); + } + let part_data = std::mem::take(&mut self.buffer).freeze(); + let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); + self.tasks.spawn(part_fut); + self.parts_submitted += 1; + } + + // Await all in-flight part uploads; abort if any part fails. + while let Some(res) = self.tasks.join_next().await { + match res { + Ok(Err(e)) => { + let err = Self::abort_and_wrap(upload, e, "part upload failure").await; Review Comment: Before `upload.abort()` fires in `abort_and_wrap`, I'd `tasks.abort_all()` and drain — right now we return on the first failing part while up to 7 other `put_part` tasks are still live in `self.tasks`, since they're only cancelled when the `JoinSet` drops after we've already returned. So `AbortMultipartUpload` races the in-flight `UploadPart`s, which can land after the abort and leave orphaned parts that bill until a lifecycle rule clears them — the same leak-until-expiry we held on for the panic arm in round 6. Same drain needs to go on all three terminal arms (this one, the `Err(join_err)` panic arm just below, and the `complete()`-failure arm). The abort-exactly-once mock only counts calls, so it's blind to this — a test where a sibling blocks on a `Notify` and one part fails would pin the ordering. ########## crates/storage/object_store/src/lib.rs: ########## @@ -0,0 +1,1245 @@ +// 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`] 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, +//! 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, BytesMut}; +use dashmap::DashMap; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +#[cfg(feature = "object_store-s3")] +use iceberg::io::{S3Config, S3_DELETE_BATCH_SIZE, S3_MULTIPART_PART_SIZE_BYTES}; +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::{MultipartUpload, ObjectStore, ObjectStoreExt, PutPayload}; +#[cfg(feature = "object_store-s3")] +use s3::{build_s3_store, parse_s3_url}; +use serde::{Deserialize, Serialize}; +use tokio::task::JoinSet; + +/// 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 { + let (kind, msg) = match &e { + object_store::Error::NotFound { path, .. } => { + (ErrorKind::Unexpected, format!("Object not found: {path}")) + } + object_store::Error::AlreadyExists { path, .. } => ( + ErrorKind::Unexpected, + format!("Object already exists: {path}"), + ), + object_store::Error::PermissionDenied { path, .. } => { + (ErrorKind::Unexpected, format!("Permission denied: {path}")) + } + object_store::Error::Unauthenticated { path, .. } => { + (ErrorKind::Unexpected, 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) +} + +/// Default batch size for S3 bulk deletions (matches Iceberg Java specification). +pub(crate) const DEFAULT_DELETE_BATCH_SIZE: usize = 250; +/// Maximum batch size allowed by the AWS S3 DeleteObjects API specification. +pub(crate) const S3_MAX_DELETE_BATCH_SIZE: usize = 1000; +/// Default part size for S3 multipart uploads (32 MiB, matches Iceberg Java specification). +pub(crate) const DEFAULT_PART_SIZE_BYTES: usize = 32 * 1024 * 1024; +/// Minimum part size allowed by the AWS S3 MultipartUpload API specification (5 MiB). +pub(crate) const MIN_PART_SIZE_BYTES: usize = 5 * 1024 * 1024; + +fn parse_delete_batch_size(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_DELETE_BATCH_SIZE) { + match val.parse::<usize>() { + 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 250" + ); + DEFAULT_DELETE_BATCH_SIZE + } + } + } else { + DEFAULT_DELETE_BATCH_SIZE + } +} + +fn default_delete_batch_size() -> usize { + DEFAULT_DELETE_BATCH_SIZE +} + +fn parse_part_size_bytes(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_MULTIPART_PART_SIZE_BYTES) { + match val.parse::<usize>() { + Ok(parsed) if parsed >= MIN_PART_SIZE_BYTES => parsed, + _ => { + tracing::warn!( + val = %val, + "Invalid s3.multipart.part-size-bytes (must be >= 5 MiB); falling back to default 32 MiB" + ); + DEFAULT_PART_SIZE_BYTES + } + } + } else { + DEFAULT_PART_SIZE_BYTES + } +} + +fn default_part_size_bytes() -> usize { + DEFAULT_PART_SIZE_BYTES +} + +/// 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 +/// backed by the `object_store` crate. +#[non_exhaustive] +#[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<Arc<dyn Storage>> { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorageFactory::S3 => { + let s3_config = S3Config::try_from(config)?; + let delete_batch_size = parse_delete_batch_size(config); + let part_size_bytes = parse_part_size_bytes(config); + tracing::info!( + batch_size = delete_batch_size, + part_size_bytes = part_size_bytes, + "Initialized S3 storage with delete batch size {} and multipart part size {} bytes (configure via '{}' and '{}' to adjust)", + delete_batch_size, + part_size_bytes, + S3_DELETE_BATCH_SIZE, + S3_MULTIPART_PART_SIZE_BYTES + ); + Ok(Arc::new(ObjectStoreStorage::S3(S3Storage { + config: Arc::new(s3_config), + delete_batch_size, + part_size_bytes, + store_cache: Arc::new(DashMap::new()), + }))) + } + } + } +} + +type StoreCache = Arc<DashMap<String, Arc<dyn ObjectStore>>>; + +/// `object_store` S3 storage state. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct S3Storage { + config: Arc<S3Config>, + #[serde(default = "default_delete_batch_size")] + delete_batch_size: usize, + #[serde(default = "default_part_size_bytes")] + part_size_bytes: usize, + #[serde(skip, default)] + store_cache: StoreCache, +} + +/// `object_store`-based storage implementation. +/// +/// Stores are cached per bucket to avoid rebuilding the client on every operation. +#[non_exhaustive] +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum ObjectStoreStorage { + /// S3 storage variant. + #[cfg(feature = "object_store-s3")] + S3(S3Storage), +} + +#[derive(Debug)] +struct StoreAndPath { + bucket: String, + store: Arc<dyn ObjectStore>, + path: ObjectStorePath, +} + +/// Helper for batching deletions per bucket. +struct BucketBatch { + store: Arc<dyn ObjectStore>, + locations: Vec<ObjectStorePath>, +} + +impl ObjectStoreStorage { + fn delete_batch_size(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.delete_batch_size.clamp(1, S3_MAX_DELETE_BATCH_SIZE), + } + } + + fn part_size_bytes(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.part_size_bytes.max(MIN_PART_SIZE_BYTES), + } + } + + /// Get or create a cached store and extract the relative `ObjectStorePath`. + fn get_store_and_path(&self, path: &str) -> Result<StoreAndPath> { + match self { + #[cfg(feature = "object_store-s3")] + 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(); + + let object_path = ObjectStorePath::parse(parsed.relative.as_str()).map_err(|e| { Review Comment: This is the same `ObjectStorePath::parse` that fixed the double-encode, which is good — but it also rejects empty path segments, so `s3://b/wh//db/tbl/f.parquet` comes back `DataInvalid`. Java's `S3URI` and opendal both accept those keys, and they turn up whenever a trailing-slash warehouse path gets joined with `/`, so tables other engines wrote become unreadable through this backend — and in `delete_stream` one such path aborts the whole stream after earlier chunks are already deleted. I'd at least document this as a known parity gap with an error that names it ("object_store cannot represent empty path segments"), and make `delete_stream` collect-and-continue rather than bail midway. Avoiding `Path::parse` for the empty-segment case (e.g. `Path::from_iter` on pre-split parts) would close it properly. ########## crates/storage/object_store/src/lib.rs: ########## @@ -0,0 +1,1245 @@ +// 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`] 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, +//! 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, BytesMut}; +use dashmap::DashMap; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +#[cfg(feature = "object_store-s3")] +use iceberg::io::{S3Config, S3_DELETE_BATCH_SIZE, S3_MULTIPART_PART_SIZE_BYTES}; +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::{MultipartUpload, ObjectStore, ObjectStoreExt, PutPayload}; +#[cfg(feature = "object_store-s3")] +use s3::{build_s3_store, parse_s3_url}; +use serde::{Deserialize, Serialize}; +use tokio::task::JoinSet; + +/// 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 { + let (kind, msg) = match &e { + object_store::Error::NotFound { path, .. } => { + (ErrorKind::Unexpected, format!("Object not found: {path}")) + } + object_store::Error::AlreadyExists { path, .. } => ( + ErrorKind::Unexpected, + format!("Object already exists: {path}"), + ), + object_store::Error::PermissionDenied { path, .. } => { + (ErrorKind::Unexpected, format!("Permission denied: {path}")) + } + object_store::Error::Unauthenticated { path, .. } => { + (ErrorKind::Unexpected, 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) +} + +/// Default batch size for S3 bulk deletions (matches Iceberg Java specification). +pub(crate) const DEFAULT_DELETE_BATCH_SIZE: usize = 250; +/// Maximum batch size allowed by the AWS S3 DeleteObjects API specification. +pub(crate) const S3_MAX_DELETE_BATCH_SIZE: usize = 1000; +/// Default part size for S3 multipart uploads (32 MiB, matches Iceberg Java specification). +pub(crate) const DEFAULT_PART_SIZE_BYTES: usize = 32 * 1024 * 1024; +/// Minimum part size allowed by the AWS S3 MultipartUpload API specification (5 MiB). +pub(crate) const MIN_PART_SIZE_BYTES: usize = 5 * 1024 * 1024; + +fn parse_delete_batch_size(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_DELETE_BATCH_SIZE) { + match val.parse::<usize>() { + 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 250" + ); + DEFAULT_DELETE_BATCH_SIZE + } + } + } else { + DEFAULT_DELETE_BATCH_SIZE + } +} + +fn default_delete_batch_size() -> usize { + DEFAULT_DELETE_BATCH_SIZE +} + +fn parse_part_size_bytes(config: &StorageConfig) -> usize { + if let Some(val) = config.get(S3_MULTIPART_PART_SIZE_BYTES) { + match val.parse::<usize>() { + Ok(parsed) if parsed >= MIN_PART_SIZE_BYTES => parsed, + _ => { + tracing::warn!( + val = %val, + "Invalid s3.multipart.part-size-bytes (must be >= 5 MiB); falling back to default 32 MiB" + ); + DEFAULT_PART_SIZE_BYTES + } + } + } else { + DEFAULT_PART_SIZE_BYTES + } +} + +fn default_part_size_bytes() -> usize { + DEFAULT_PART_SIZE_BYTES +} + +/// 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 +/// backed by the `object_store` crate. +#[non_exhaustive] +#[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<Arc<dyn Storage>> { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorageFactory::S3 => { + let s3_config = S3Config::try_from(config)?; + let delete_batch_size = parse_delete_batch_size(config); + let part_size_bytes = parse_part_size_bytes(config); + tracing::info!( + batch_size = delete_batch_size, + part_size_bytes = part_size_bytes, + "Initialized S3 storage with delete batch size {} and multipart part size {} bytes (configure via '{}' and '{}' to adjust)", + delete_batch_size, + part_size_bytes, + S3_DELETE_BATCH_SIZE, + S3_MULTIPART_PART_SIZE_BYTES + ); + Ok(Arc::new(ObjectStoreStorage::S3(S3Storage { + config: Arc::new(s3_config), + delete_batch_size, + part_size_bytes, + store_cache: Arc::new(DashMap::new()), + }))) + } + } + } +} + +type StoreCache = Arc<DashMap<String, Arc<dyn ObjectStore>>>; + +/// `object_store` S3 storage state. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct S3Storage { + config: Arc<S3Config>, + #[serde(default = "default_delete_batch_size")] + delete_batch_size: usize, + #[serde(default = "default_part_size_bytes")] + part_size_bytes: usize, + #[serde(skip, default)] + store_cache: StoreCache, +} + +/// `object_store`-based storage implementation. +/// +/// Stores are cached per bucket to avoid rebuilding the client on every operation. +#[non_exhaustive] +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum ObjectStoreStorage { + /// S3 storage variant. + #[cfg(feature = "object_store-s3")] + S3(S3Storage), +} + +#[derive(Debug)] +struct StoreAndPath { + bucket: String, + store: Arc<dyn ObjectStore>, + path: ObjectStorePath, +} + +/// Helper for batching deletions per bucket. +struct BucketBatch { + store: Arc<dyn ObjectStore>, + locations: Vec<ObjectStorePath>, +} + +impl ObjectStoreStorage { + fn delete_batch_size(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.delete_batch_size.clamp(1, S3_MAX_DELETE_BATCH_SIZE), + } + } + + fn part_size_bytes(&self) -> usize { + match self { + #[cfg(feature = "object_store-s3")] + ObjectStoreStorage::S3(s3) => s3.part_size_bytes.max(MIN_PART_SIZE_BYTES), + } + } + + /// Get or create a cached store and extract the relative `ObjectStorePath`. + fn get_store_and_path(&self, path: &str) -> Result<StoreAndPath> { + match self { + #[cfg(feature = "object_store-s3")] + 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(); + + let object_path = ObjectStorePath::parse(parsed.relative.as_str()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid relative path: {}", parsed.relative), + ) + .with_source(e) + })?; + + Ok(StoreAndPath { + bucket: parsed.bucket, + store, + path: object_path, + }) + } + } + } +} + +#[typetag::serde(name = "ObjectStoreStorage")] +#[async_trait] +impl Storage for ObjectStoreStorage { + async fn exists(&self, path: &str) -> Result<bool> { + 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)), + } + } + + async fn metadata(&self, path: &str) -> Result<FileMetadata> { + let target = self.get_store_and_path(path)?; + let meta = target + .store + .head(&target.path) + .await + .map_err(from_object_store_error)?; + Ok(to_file_metadata(meta)) + } + + async fn read(&self, path: &str) -> Result<Bytes> { + 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<Box<dyn FileRead>> { + let target = self.get_store_and_path(path)?; + Ok(Box::new(ObjectStoreReader { + store: target.store, + path: target.path, + })) + } + + async fn write(&self, path: &str, bs: Bytes) -> Result<()> { + 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<Box<dyn FileWrite>> { + let target = self.get_store_and_path(path)?; + let upload = target + .store + .put_multipart(&target.path) + .await + .map_err(from_object_store_error)?; + Ok(Box::new(ObjectStoreWriter { + upload: Some(upload), + buffer: BytesMut::new(), + tasks: JoinSet::new(), + parts_submitted: 0, + bytes_written: 0, + part_size_bytes: self.part_size_bytes(), + })) + } + + async fn delete(&self, path: &str) -> Result<()> { + 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 target = self.get_store_and_path(path)?; + let locations = target + .store + .list(Some(&target.path)) + .map_ok(|m| m.location) + .boxed(); + target + .store + .delete_stream(locations) + .try_for_each(|_| async { Ok(()) }) + .await + .map_err(from_object_store_error)?; + Ok(()) + } + + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { + 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<String, BucketBatch> = + 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, 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(()) + } + + fn new_input(&self, path: &str) -> Result<InputFile> { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> Result<OutputFile> { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } +} + +/// Reader that implements `FileRead` using `object_store`. +struct ObjectStoreReader { + store: Arc<dyn ObjectStore>, + path: ObjectStorePath, +} + +#[async_trait] +impl FileRead for ObjectStoreReader { + async fn read(&self, range: Range<u64>) -> Result<Bytes> { + if range.is_empty() { + return Ok(Bytes::new()); + } + 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) + } +} + +/// 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<Box<dyn MultipartUpload>>, + buffer: BytesMut, + tasks: JoinSet<object_store::Result<()>>, + parts_submitted: usize, + bytes_written: u64, + part_size_bytes: usize, +} + +impl ObjectStoreWriter { + /// 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<dyn MultipartUpload>, + 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) + } + + /// Drains in-flight tasks down below `MAX_CONCURRENT_PART_UPLOADS`. + async fn drain_capacity( + tasks: &mut JoinSet<object_store::Result<()>>, + upload: &mut Box<dyn MultipartUpload>, + ) -> Result<()> { + while tasks.len() >= MAX_CONCURRENT_PART_UPLOADS { + if let Some(res) = tasks.join_next().await { + match res { + Ok(Err(e)) => { + return Err(Self::abort_and_wrap(upload, e, "part upload failure").await); + } + Err(join_err) => { + if let Err(abort_err) = upload.abort().await { + tracing::warn!( + error = %abort_err, + "Failed to abort multipart upload after task panic" + ); + } + return Err( + Error::new(ErrorKind::Unexpected, "Part upload task panicked") + .with_source(join_err), + ); + } + Ok(Ok(())) => {} + } + } + } + Ok(()) + } + + /// 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 JoinSet<object_store::Result<()>>, + parts_submitted: &mut usize, + upload: &mut Box<dyn MultipartUpload>, + ) -> Result<()> { + if !buffer.is_empty() { + Self::drain_capacity(tasks, upload).await?; + + let part_data = std::mem::take(buffer).freeze(); + let part_fut = upload.put_part(PutPayload::from_bytes(part_data)); + tasks.spawn(part_fut); Review Comment: `tasks.spawn` panics when there's no Tokio runtime, and this is library code reachable from `write()`/`close()` under any executor (same for the `close()` spawn). `Drop` already guards exactly this with `Handle::try_current()`, so the two paths disagree. Returning an error — or routing through the same `try_current` check — keeps it consistent and off the panic path. ########## crates/storage/object_store/tests/file_io_s3_test.rs: ########## @@ -0,0 +1,581 @@ +// 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, TryStreamExt}; + use iceberg::io::{ + FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, + S3_SECRET_ACCESS_KEY, S3_SSE_KEY, S3_SSE_TYPE, + }; + use iceberg_storage_object_store::ObjectStoreStorageFactory; + use iceberg_test_utils::{get_object_store_endpoint, normalize_test_name_with_parts, set_up}; + use object_store::aws::AmazonS3Builder; + use object_store::path::Path as ObjectStorePath; + use object_store::ObjectStore; + + async fn get_file_io() -> FileIO { + set_up(); + + let endpoint = get_object_store_endpoint(); + + 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()), + ]) + .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; + let file_path = format!( + "s3://bucket1/{}", + normalize_test_name_with_parts!("test_file_io_s3_exists") + ); + + let _ = file_io.delete(&file_path).await; + assert!(!file_io.exists(&file_path).await.unwrap()); + + let output_file = file_io.new_output(&file_path).unwrap(); + output_file + .write(Bytes::from_static(b"test_exists")) + .await + .unwrap(); + assert!(file_io.exists(&file_path).await.unwrap()); + + file_io.delete(&file_path).await.unwrap(); + assert!(!file_io.exists(&file_path).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<String> = (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(); + } + + #[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_percent_encoded_bucket() { + let file_io = get_file_io().await; + let file_path = format!( + "s3://bucket%31/{}", + normalize_test_name_with_parts!("test_file_io_s3_percent_encoded_bucket") + ); + let canonical_path = format!( + "s3://bucket1/{}", + 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_encoded_partition_path() { + set_up(); + let endpoint = get_object_store_endpoint(); + + let file_io = FileIOBuilder::new(Arc::new(ObjectStoreStorageFactory::S3)) + .with_props(vec![ + (S3_ENDPOINT, endpoint.clone()), + (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(); + + let prefix = normalize_test_name_with_parts!("test_file_io_s3_encoded_partition_path"); + let file_path = format!("s3://bucket1/{prefix}/dt=a%2Fb/data.parquet"); + + let _ = file_io.delete(&file_path).await; + file_io + .new_output(&file_path) + .unwrap() + .write(Bytes::from_static(b"partition-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"partition-data")); + + // Direct S3 client listing assertion: verify the raw key in S3 is literally `dt=a%2Fb/data.parquet` and NOT double-encoded (`dt=a%252Fb/data.parquet`) + let raw_store = AmazonS3Builder::new() + .with_endpoint(&endpoint) + .with_allow_http(true) + .with_access_key_id("admin") + .with_secret_access_key("password") + .with_region("us-east-1") + .with_bucket_name("bucket1") + .with_virtual_hosted_style_request(false) + .build() + .unwrap(); + + let raw_prefix = ObjectStorePath::parse(format!("{prefix}/dt=a%2Fb")).unwrap(); + let listed_keys: Vec<String> = raw_store + .list(Some(&raw_prefix)) + .map_ok(|meta| meta.location.to_string()) + .try_collect() + .await + .unwrap(); + + assert_eq!( + listed_keys, + vec![format!("{prefix}/dt=a%2Fb/data.parquet")], + "Expected raw S3 key with single percent encoding, found double-encoded or wrong key!" + ); + + file_io.delete(&file_path).await.unwrap(); + 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("ServerNotInitialized") + || err_str.contains("A header you provided implies functionality that is not implemented"), + "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_object_store_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; + // 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 + .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_object_store_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; + // 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 + .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 + /// past the 10 MiB `WriteMultipart` buffer threshold, then reads back and Review Comment: This comment still cites the 10 MiB `WriteMultipart` threshold, but the code uses `part_size_bytes` and the default is 32 MiB, so writing 12 MiB produces exactly one part at `close()` — this never exercises a real multi-part upload. Setting `s3.multipart.part-size-bytes=5MiB` and writing several parts would actually run part numbering/ETag assembly across ≥2 parts, the 5 MiB-min rule, and the 8-part throttle against MinIO. Worth fixing the stale comment in the same pass. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
