From ad85c537712ca082c3d07c8c59f6cd190848db53 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Mon, 15 Jun 2026 18:28:21 +0800 Subject: [PATCH 01/13] feat(rest): use vended storage credentials from LoadTable response --- crates/catalog/rest/src/catalog.rs | 47 ++++++++++++++++++++---------- 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index aa36a074e7..bc164186bc 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -49,7 +49,7 @@ use crate::response::HttpResponse; use crate::types::{ CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult, - NamespaceResponse, RegisterTableRequest, RenameTableRequest, + NamespaceResponse, RegisterTableRequest, RenameTableRequest, StorageCredential, }; /// REST catalog URI @@ -1141,6 +1141,7 @@ impl SessionCatalog for RestSessionCatalog { client .http_client .request(Method::POST, client.config.tables_endpoint(namespace)) + .header("X-Iceberg-Access-Delegation", "vended-credentials") .json(&CreateTableRequest { name: creation.name, location: creation.location, @@ -1181,11 +1182,7 @@ impl SessionCatalog for RestSessionCatalog { "Metadata location missing in `create_table` response!", ))?; - let config = response - .config - .into_iter() - .chain(self.user_config.props.clone()) - .collect(); + let config = table_file_io_config(&response, &self.user_config.props); let file_io = self .load_file_io(Some(metadata_location), Some(config)) @@ -1222,7 +1219,9 @@ impl SessionCatalog for RestSessionCatalog { let request = HttpRequest::build( client .http_client - .request(Method::GET, client.config.table_endpoint(table_ident)), + .request(Method::GET, client.config.table_endpoint(table_ident)) + // Opt in to vended storage credentials. + .header("X-Iceberg-Access-Delegation", "vended-credentials"), )?; let http_response = client.query_catalog(request).await?; @@ -1245,11 +1244,7 @@ impl SessionCatalog for RestSessionCatalog { } }; - let config = response - .config - .into_iter() - .chain(self.user_config.props.clone()) - .collect(); + let config = table_file_io_config(&response, &self.user_config.props); let file_io = self .load_file_io(response.metadata_location.as_deref(), Some(config)) @@ -1407,7 +1402,7 @@ impl SessionCatalog for RestSessionCatalog { async fn update_table( &self, - _context: &SessionContext, + context: &SessionContext, mut commit: TableCommit, ) -> Result { let client = self.client().await?; @@ -1469,9 +1464,12 @@ impl SessionCatalog for RestSessionCatalog { } }; + // Reload for a credentialed FileIO (commit response carries no credentials). let file_io = self - .load_file_io(Some(&response.metadata_location), None) - .await?; + .load_table(context, commit.identifier()) + .await? + .file_io() + .clone(); let mut table_builder = Table::builder() .identifier(commit.identifier().clone()) @@ -1660,6 +1658,23 @@ impl RestSessionCatalogBuilder { } } +/// FileIO props: server `config`, then vended `storage_credentials` (longest prefix wins), then user props. +fn table_file_io_config( + response: &LoadTableResult, + user_props: &HashMap, +) -> HashMap { + let mut config: HashMap = response.config.clone(); + if let Some(creds) = response.storage_credentials.as_ref() { + let mut sorted: Vec<&StorageCredential> = creds.iter().collect(); + sorted.sort_by_key(|c| c.prefix.len()); + for cred in sorted { + config.extend(cred.config.clone()); + } + } + config.extend(user_props.clone()); + config +} + #[cfg(test)] mod tests { use std::fs::File; @@ -4009,6 +4024,7 @@ mod tests { let config_mock = create_config_mock(&mut server).await; + // GET hit twice: commit refresh + post-commit reload. let load_table_mock = server .mock("GET", "/v1/namespaces/ns1/tables/test1") .with_status(200) @@ -4017,6 +4033,7 @@ mod tests { env!("CARGO_MANIFEST_DIR"), "load_table_response.json" )) + .expect(2) .create_async() .await; From 4c694119a36532b6e2e0e6e6993a7f4fa64212ea Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 16 Jun 2026 00:29:12 +0800 Subject: [PATCH 02/13] feat(rest): per-prefix vended credentials with opt-in delegation header --- crates/catalog/rest/src/catalog.rs | 165 ++++++++++--- .../load_table_response_with_credentials.json | 78 ++++++ crates/iceberg/public-api.txt | 1 + crates/iceberg/src/io/file_io.rs | 231 ++++++++++++++++-- crates/iceberg/src/table.rs | 6 + crates/iceberg/src/transaction/mod.rs | 4 +- 6 files changed, 424 insertions(+), 61 deletions(-) create mode 100644 crates/catalog/rest/testdata/load_table_response_with_credentials.json diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index bc164186bc..ebce25d01f 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -847,6 +847,7 @@ impl RestSessionCatalog { &self, metadata_location: Option<&str>, extra_config: Option>, + storage_credentials: Option<&[StorageCredential]>, ) -> Result { let mut props = self.client().await?.config.props.clone(); if let Some(config) = extra_config { @@ -879,9 +880,19 @@ impl RestSessionCatalog { ) })?; - let file_io = FileIOBuilder::new(factory).with_props(props).build(); + let mut builder = FileIOBuilder::new(factory).with_props(props.clone()); - Ok(file_io) + // Vended credentials are scoped per location prefix: give each its own + // storage so reads/writes use the matching credentials. + if let Some(creds) = storage_credentials { + for cred in creds { + let mut prefixed = props.clone(); + prefixed.extend(cred.config.clone()); + builder = builder.with_prefixed_props(cred.prefix.clone(), prefixed); + } + } + + Ok(builder.build()) } } @@ -1141,7 +1152,6 @@ impl SessionCatalog for RestSessionCatalog { client .http_client .request(Method::POST, client.config.tables_endpoint(namespace)) - .header("X-Iceberg-Access-Delegation", "vended-credentials") .json(&CreateTableRequest { name: creation.name, location: creation.location, @@ -1182,10 +1192,15 @@ impl SessionCatalog for RestSessionCatalog { "Metadata location missing in `create_table` response!", ))?; - let config = table_file_io_config(&response, &self.user_config.props); + let mut base_config = response.config.clone(); + base_config.extend(self.user_config.props.clone()); let file_io = self - .load_file_io(Some(metadata_location), Some(config)) + .load_file_io( + Some(metadata_location), + Some(base_config), + response.storage_credentials.as_deref(), + ) .await?; let mut table_builder = Table::builder() @@ -1216,12 +1231,13 @@ impl SessionCatalog for RestSessionCatalog { ) -> Result
{ let client = self.client().await?; + // Vended credentials are opt-in via a `header.X-Iceberg-Access-Delegation` + // catalog property (applied to every request like the Iceberg Java client); + // any returned `storage_credentials` are wired into the FileIO below. let request = HttpRequest::build( client .http_client - .request(Method::GET, client.config.table_endpoint(table_ident)) - // Opt in to vended storage credentials. - .header("X-Iceberg-Access-Delegation", "vended-credentials"), + .request(Method::GET, client.config.table_endpoint(table_ident)), )?; let http_response = client.query_catalog(request).await?; @@ -1244,10 +1260,15 @@ impl SessionCatalog for RestSessionCatalog { } }; - let config = table_file_io_config(&response, &self.user_config.props); + let mut base_config = response.config.clone(); + base_config.extend(self.user_config.props.clone()); let file_io = self - .load_file_io(response.metadata_location.as_deref(), Some(config)) + .load_file_io( + response.metadata_location.as_deref(), + Some(base_config), + response.storage_credentials.as_deref(), + ) .await?; let mut table_builder = Table::builder() @@ -1386,7 +1407,9 @@ impl SessionCatalog for RestSessionCatalog { "Metadata location missing in `register_table` response!", ))?; - let file_io = self.load_file_io(Some(metadata_location), None).await?; + let file_io = self + .load_file_io(Some(metadata_location), None, None) + .await?; let mut table_builder = Table::builder() .identifier(table_ident.clone()) @@ -1402,7 +1425,7 @@ impl SessionCatalog for RestSessionCatalog { async fn update_table( &self, - context: &SessionContext, + _context: &SessionContext, mut commit: TableCommit, ) -> Result
{ let client = self.client().await?; @@ -1464,12 +1487,11 @@ impl SessionCatalog for RestSessionCatalog { } }; - // Reload for a credentialed FileIO (commit response carries no credentials). + // The commit response carries no credentials, so build a plain FileIO; + // the transaction layer reuses the credentialed one it already holds. let file_io = self - .load_table(context, commit.identifier()) - .await? - .file_io() - .clone(); + .load_file_io(Some(&response.metadata_location), None, None) + .await?; let mut table_builder = Table::builder() .identifier(commit.identifier().clone()) @@ -1658,23 +1680,6 @@ impl RestSessionCatalogBuilder { } } -/// FileIO props: server `config`, then vended `storage_credentials` (longest prefix wins), then user props. -fn table_file_io_config( - response: &LoadTableResult, - user_props: &HashMap, -) -> HashMap { - let mut config: HashMap = response.config.clone(); - if let Some(creds) = response.storage_credentials.as_ref() { - let mut sorted: Vec<&StorageCredential> = creds.iter().collect(); - sorted.sort_by_key(|c| c.prefix.len()); - for cred in sorted { - config.extend(cred.config.clone()); - } - } - config.extend(user_props.clone()); - config -} - #[cfg(test)] mod tests { use std::fs::File; @@ -3809,6 +3814,94 @@ mod tests { rename_table_mock.assert_async().await; } + #[tokio::test] + async fn test_load_table_uses_vended_credentials() { + let mut server = Server::new_async().await; + + let config_mock = create_config_mock(&mut server).await; + + // Vended credentials are opt-in via a `header.*` catalog property (like the + // Java client). With it configured, the header is sent and the response's + // `storage-credentials` are accepted (the FileIO builds). + let load_table_mock = server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .match_header("x-iceberg-access-delegation", "vended-credentials") + .with_status(200) + .with_body_from_file(format!( + "{}/testdata/{}", + env!("CARGO_MANIFEST_DIR"), + "load_table_response_with_credentials.json" + )) + .create_async() + .await; + + let props = HashMap::from([( + "header.X-Iceberg-Access-Delegation".to_string(), + "vended-credentials".to_string(), + )]); + let catalog = RestCatalog::new( + SessionContext::empty(), + RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(), + None, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + let table = catalog + .load_table(&TableIdent::from_strs(["ns1", "test1"]).unwrap()) + .await + .unwrap(); + + assert_eq!( + "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json", + table.metadata_location().unwrap() + ); + + config_mock.assert_async().await; + load_table_mock.assert_async().await; + } + + #[tokio::test] + async fn test_load_table_omits_delegation_header_by_default() { + let mut server = Server::new_async().await; + + let config_mock = create_config_mock(&mut server).await; + + // No delegation header is hardcoded: without a `header.*` prop, none is sent. + let load_table_mock = server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .match_header("x-iceberg-access-delegation", mockito::Matcher::Missing) + .with_status(200) + .with_body_from_file(format!( + "{}/testdata/{}", + env!("CARGO_MANIFEST_DIR"), + "load_table_response.json" + )) + .create_async() + .await; + + let catalog = RestCatalog::new( + SessionContext::empty(), + RestCatalogConfig::builder().uri(server.url()).build(), + None, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + catalog + .load_table(&TableIdent::from_strs(["ns1", "test1"]).unwrap()) + .await + .unwrap(); + + config_mock.assert_async().await; + load_table_mock.assert_async().await; + } + #[tokio::test] async fn test_create_table() { let mut server = Server::new_async().await; @@ -4024,7 +4117,7 @@ mod tests { let config_mock = create_config_mock(&mut server).await; - // GET hit twice: commit refresh + post-commit reload. + // GET hit once: the transaction refreshes the table before committing. let load_table_mock = server .mock("GET", "/v1/namespaces/ns1/tables/test1") .with_status(200) @@ -4033,7 +4126,7 @@ mod tests { env!("CARGO_MANIFEST_DIR"), "load_table_response.json" )) - .expect(2) + .expect(1) .create_async() .await; diff --git a/crates/catalog/rest/testdata/load_table_response_with_credentials.json b/crates/catalog/rest/testdata/load_table_response_with_credentials.json new file mode 100644 index 0000000000..c0a540fbac --- /dev/null +++ b/crates/catalog/rest/testdata/load_table_response_with_credentials.json @@ -0,0 +1,78 @@ +{ + "metadata-location": "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json", + "metadata": { + "format-version": 1, + "table-uuid": "b55d9dda-6561-423a-8bfc-787980ce421f", + "location": "s3://warehouse/database/table", + "last-updated-ms": 1646787054459, + "last-column-id": 2, + "schema": { + "type": "struct", + "schema-id": 0, + "fields": [ + {"id": 1, "name": "id", "required": false, "type": "int"}, + {"id": 2, "name": "data", "required": false, "type": "string"} + ] + }, + "current-schema-id": 0, + "schemas": [ + { + "type": "struct", + "schema-id": 0, + "fields": [ + {"id": 1, "name": "id", "required": false, "type": "int"}, + {"id": 2, "name": "data", "required": false, "type": "string"} + ] + } + ], + "partition-spec": [], + "default-spec-id": 0, + "partition-specs": [{"spec-id": 0, "fields": []}], + "last-partition-id": 999, + "default-sort-order-id": 0, + "sort-orders": [{"order-id": 0, "fields": []}], + "properties": {"owner": "bryan", "write.metadata.compression-codec": "gzip"}, + "current-snapshot-id": 3497810964824022504, + "refs": {"main": {"snapshot-id": 3497810964824022504, "type": "branch"}}, + "snapshots": [ + { + "snapshot-id": 3497810964824022504, + "timestamp-ms": 1646787054459, + "summary": { + "operation": "append", + "spark.app.id": "local-1646787004168", + "added-data-files": "1", + "added-records": "1", + "added-files-size": "697", + "changed-partition-count": "1", + "total-records": "1", + "total-files-size": "697", + "total-data-files": "1", + "total-delete-files": "0", + "total-position-deletes": "0", + "total-equality-deletes": "0" + }, + "manifest-list": "s3://warehouse/database/table/metadata/snap-3497810964824022504-1-c4f68204-666b-4e50-a9df-b10c34bf6b82.avro", + "schema-id": 0 + } + ], + "snapshot-log": [{"timestamp-ms": 1646787054459, "snapshot-id": 3497810964824022504}], + "metadata-log": [ + { + "timestamp-ms": 1646787031514, + "metadata-file": "s3://warehouse/database/table/metadata/00000-88484a1c-00e5-4a07-a787-c0e7aeffa805.gz.metadata.json" + } + ] + }, + "config": {"client.factory": "io.tabular.iceberg.catalog.TabularAwsClientFactory", "region": "us-west-2"}, + "storage-credentials": [ + { + "prefix": "s3://warehouse/database/table", + "config": { + "s3.access-key-id": "vended-key-id", + "s3.secret-access-key": "vended-secret", + "s3.session-token": "vended-token" + } + } + ] +} diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index e2e78073e9..ef924b50a6 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -747,6 +747,7 @@ impl iceberg::io::FileIOBuilder pub fn iceberg::io::FileIOBuilder::build(self) -> iceberg::io::FileIO pub fn iceberg::io::FileIOBuilder::config(&self) -> &iceberg::io::StorageConfig pub fn iceberg::io::FileIOBuilder::new(factory: alloc::sync::Arc) -> Self +pub fn iceberg::io::FileIOBuilder::with_prefixed_props(self, prefix: impl core::convert::Into, props: impl core::iter::traits::collect::IntoIterator) -> Self pub fn iceberg::io::FileIOBuilder::with_prop(self, key: impl alloc::string::ToString, value: impl alloc::string::ToString) -> Self pub fn iceberg::io::FileIOBuilder::with_props(self, args: impl core::iter::traits::collect::IntoIterator) -> Self impl core::clone::Clone for iceberg::io::FileIOBuilder diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 42eec4db56..10112f8458 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -15,11 +15,12 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::ops::Range; use std::sync::{Arc, OnceLock}; use bytes::Bytes; -use futures::{Stream, StreamExt}; +use futures::{Stream, StreamExt, stream}; use super::storage::{ LocalFsStorageFactory, MemoryStorageFactory, Storage, StorageConfig, StorageFactory, @@ -67,6 +68,17 @@ pub struct FileIO { factory: Arc, /// Cached storage instance (lazily initialized) storage: Arc>>, + /// Per-prefix storages (longest prefix first) for tables that vend distinct + /// credentials per location prefix. Paths matching none use `storage` above. + prefixed: Arc>, +} + +/// A storage scoped to a location `prefix`, lazily built from its own config. +#[derive(Debug)] +struct PrefixedStorage { + prefix: String, + config: StorageConfig, + storage: OnceLock>, } mod _serde { @@ -80,12 +92,29 @@ mod _serde { pub(super) struct SerializableFileIO<'a> { pub(super) config: &'a StorageConfig, pub(super) factory: &'a Arc, + /// Per-prefix credentials travel too: a worker that deserializes this + /// FileIO has to read the same data the catalog vended them for. + pub(super) prefixed: Vec>, + } + + #[derive(Serialize)] + pub(super) struct SerializablePrefixed<'a> { + pub(super) prefix: &'a str, + pub(super) config: &'a StorageConfig, } #[derive(Deserialize)] pub(super) struct DeserializedFileIO { pub(super) config: StorageConfig, pub(super) factory: Arc, + #[serde(default)] + pub(super) prefixed: Vec, + } + + #[derive(Deserialize)] + pub(super) struct DeserializedPrefixed { + pub(super) prefix: String, + pub(super) config: StorageConfig, } } @@ -98,6 +127,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(MemoryStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -109,6 +139,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(LocalFsStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -131,6 +162,14 @@ impl FileIO { Ok(serde_json::to_vec(&_serde::SerializableFileIO { config: &self.config, factory: &self.factory, + prefixed: self + .prefixed + .iter() + .map(|p| _serde::SerializablePrefixed { + prefix: &p.prefix, + config: &p.config, + }) + .collect(), })?) } @@ -140,11 +179,26 @@ impl FileIO { /// implementation so it is registered with `typetag`. Backend-specific requirements are /// documented by each storage factory implementation. pub fn deserialize_all(bytes: &[u8]) -> Result { - let _serde::DeserializedFileIO { config, factory } = serde_json::from_slice(bytes)?; + let _serde::DeserializedFileIO { + config, + factory, + prefixed, + } = serde_json::from_slice(bytes)?; + // Order is preserved from `serialize_all`, so this is still longest + // prefix first. + let prefixed = prefixed + .into_iter() + .map(|p| PrefixedStorage { + prefix: p.prefix, + config: p.config, + storage: OnceLock::new(), + }) + .collect(); Ok(Self { config, factory, storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(prefixed), }) } @@ -153,24 +207,31 @@ impl FileIO { &self.config } - /// Get or create the storage instance. - /// - /// The factory is invoked on first access and the result is cached - /// for all subsequent operations. - fn get_storage(&self) -> Result> { - // Check if already initialized - if let Some(storage) = self.storage.get() { - return Ok(storage.clone()); + /// Get or create the storage for `path`, routing to the longest-matching + /// prefix storage if any, else the default. Built once, then cached. + fn get_storage(&self, path: &str) -> Result> { + // `prefixed` is sorted longest-first, so the first match is most specific. + for ps in self.prefixed.iter() { + if path.starts_with(&ps.prefix) { + return Self::get_or_build(&ps.storage, &self.factory, &ps.config); + } } + Self::get_or_build(&self.storage, &self.factory, &self.config) + } - // Build the storage - let storage = self.factory.build(&self.config)?; - - // Try to set it (another thread might have set it first) - let _ = self.storage.set(storage.clone()); - - // Return whatever is in the cell (either ours or another thread's) - Ok(self.storage.get().unwrap().clone()) + /// Get a cached storage from `cell`, building it from `config` on first use. + fn get_or_build( + cell: &OnceLock>, + factory: &Arc, + config: &StorageConfig, + ) -> Result> { + if let Some(storage) = cell.get() { + return Ok(storage.clone()); + } + let storage = factory.build(config)?; + // Another thread might have set it first; keep whatever ends up in the cell. + let _ = cell.set(storage); + Ok(cell.get().unwrap().clone()) } /// Deletes file. @@ -179,7 +240,7 @@ impl FileIO { /// /// * path: It should be *absolute* path starting with scheme string used to construct [`FileIO`]. pub async fn delete(&self, path: impl AsRef) -> Result<()> { - self.get_storage()?.delete(path.as_ref()).await + self.get_storage(path.as_ref())?.delete(path.as_ref()).await } /// Remove the path and all nested dirs and files recursively. @@ -194,7 +255,9 @@ impl FileIO { /// - If the path is a empty directory, this function will remove the directory itself. /// - If the path is a non-empty directory, this function will remove the directory and all nested files and directories. pub async fn delete_prefix(&self, path: impl AsRef) -> Result<()> { - self.get_storage()?.delete_prefix(path.as_ref()).await + self.get_storage(path.as_ref())? + .delete_prefix(path.as_ref()) + .await } /// Delete multiple files from a stream of paths. @@ -206,7 +269,29 @@ impl FileIO { &self, paths: impl Stream + Send + 'static, ) -> Result<()> { - self.get_storage()?.delete_stream(paths.boxed()).await + // No per-prefix storages: delete the whole batch on the default storage. + if self.prefixed.is_empty() { + return self.get_storage("")?.delete_stream(paths.boxed()).await; + } + + // Otherwise group paths by routed storage, then batch-delete each group. + let mut groups: HashMap> = HashMap::new(); + let mut paths = paths.boxed(); + while let Some(path) = paths.next().await { + let key = self + .prefixed + .iter() + .find(|ps| path.starts_with(&ps.prefix)) + .map(|ps| ps.prefix.clone()) + .unwrap_or_default(); + groups.entry(key).or_default().push(path); + } + + for batch in groups.into_values() { + let storage = self.get_storage(&batch[0])?; + storage.delete_stream(stream::iter(batch).boxed()).await?; + } + Ok(()) } /// Check file exists. @@ -215,7 +300,7 @@ impl FileIO { /// /// * path: It should be *absolute* path starting with scheme string used to construct [`FileIO`]. pub async fn exists(&self, path: impl AsRef) -> Result { - self.get_storage()?.exists(path.as_ref()).await + self.get_storage(path.as_ref())?.exists(path.as_ref()).await } /// Creates input file. @@ -224,7 +309,7 @@ impl FileIO { /// /// * path: It should be *absolute* path starting with scheme string used to construct [`FileIO`]. pub fn new_input(&self, path: impl AsRef) -> Result { - self.get_storage()?.new_input(path.as_ref()) + self.get_storage(path.as_ref())?.new_input(path.as_ref()) } /// Creates output file. @@ -233,7 +318,7 @@ impl FileIO { /// /// * path: It should be *absolute* path starting with scheme string used to construct [`FileIO`]. pub fn new_output(&self, path: impl AsRef) -> Result { - self.get_storage()?.new_output(path.as_ref()) + self.get_storage(path.as_ref())?.new_output(path.as_ref()) } } @@ -247,6 +332,8 @@ pub struct FileIOBuilder { factory: Arc, /// Storage configuration config: StorageConfig, + /// Per-location-prefix configs (prefix, config). + prefixed: Vec<(String, StorageConfig)>, } impl FileIOBuilder { @@ -255,6 +342,7 @@ impl FileIOBuilder { Self { factory, config: StorageConfig::new(), + prefixed: Vec::new(), } } @@ -275,6 +363,23 @@ impl FileIOBuilder { self } + /// Add a per-prefix storage config. Paths starting with `prefix` (longest + /// match wins) use these props instead of the default config. + pub fn with_prefixed_props( + mut self, + prefix: impl Into, + props: impl IntoIterator, + ) -> Self { + let config = StorageConfig::from_props( + props + .into_iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ); + self.prefixed.push((prefix.into(), config)); + self + } + /// Get the storage configuration. pub fn config(&self) -> &StorageConfig { &self.config @@ -282,10 +387,22 @@ impl FileIOBuilder { /// Builds [`FileIO`]. pub fn build(self) -> FileIO { + let mut prefixed: Vec = self + .prefixed + .into_iter() + .map(|(prefix, config)| PrefixedStorage { + prefix, + config, + storage: OnceLock::new(), + }) + .collect(); + // Longest prefix first so routing picks the most specific match. + prefixed.sort_by(|a, b| b.prefix.len().cmp(&a.prefix.len())); FileIO { config: self.config, factory: self.factory, storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(prefixed), } } } @@ -677,4 +794,70 @@ mod tests { ); assert!(deserialized.storage.get().is_some()); } + + #[tokio::test] + async fn test_prefixed_props_sorted_by_descending_prefix_length() { + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prefixed_props("memory://a/", [("k", "short")]) + .with_prefixed_props("memory://a/longer/", [("k", "long")]) + .build(); + + // Longest prefix first so the most specific match wins at routing time. + let prefixes: Vec<&str> = file_io.prefixed.iter().map(|p| p.prefix.as_str()).collect(); + assert_eq!(prefixes, vec!["memory://a/longer/", "memory://a/"]); + } + + #[tokio::test] + async fn test_get_storage_routes_by_prefix() { + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prop("scope", "default") + .with_prefixed_props("memory://creds/", [("scope", "prefixed")]) + .build(); + + let default_a = file_io.get_storage("memory://other/x").unwrap(); + let default_b = file_io.get_storage("memory://other/y").unwrap(); + let prefixed_a = file_io.get_storage("memory://creds/x").unwrap(); + let prefixed_b = file_io.get_storage("memory://creds/y").unwrap(); + + // Repeated routing to the same bucket returns the memoized storage... + assert!(Arc::ptr_eq(&default_a, &default_b)); + assert!(Arc::ptr_eq(&prefixed_a, &prefixed_b)); + // ...and a prefix-matching path resolves to a distinct storage from the default. + assert!(!Arc::ptr_eq(&default_a, &prefixed_a)); + } + + #[tokio::test] + async fn test_delete_stream_routes_by_prefix() { + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prefixed_props("memory:/creds/", [("k", "v")]) + .build(); + + // One file under each routing bucket (default vs prefixed storage). + let default_path = "memory:/other/a.txt"; + let prefixed_path = "memory:/creds/b.txt"; + for path in [default_path, prefixed_path] { + file_io + .new_output(path) + .unwrap() + .write("x".into()) + .await + .unwrap(); + assert!(file_io.exists(path).await.unwrap()); + } + + // delete_stream must route each path to the storage that holds it. + file_io + .delete_stream(futures::stream::iter(vec![ + default_path.to_string(), + prefixed_path.to_string(), + ])) + .await + .unwrap(); + + assert!(!file_io.exists(default_path).await.unwrap()); + assert!(!file_io.exists(prefixed_path).await.unwrap()); + } } diff --git a/crates/iceberg/src/table.rs b/crates/iceberg/src/table.rs index 78c2dff05b..63b5cac355 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -219,6 +219,12 @@ impl Table { self } + /// Sets the [`Table`] `FileIO` and returns an updated instance. + pub(crate) fn with_file_io(mut self, file_io: FileIO) -> Self { + self.file_io = file_io; + self + } + /// Returns a TableBuilder to build a table pub fn builder() -> TableBuilder { TableBuilder::new() diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index 3622471277..887275d241 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -238,7 +238,9 @@ impl Transaction { .requirements(existing_requirements) .build(); - catalog.update_table(table_commit).await + let committed = catalog.update_table(table_commit).await?; + // Reuse the credentialed FileIO we already hold; the commit response has none. + Ok(committed.with_file_io(self.table.file_io().clone())) } } From 77d5c7d3c60656cfa2c14fab8f4537ecb7a299e4 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 17 Jun 2026 17:05:34 +0800 Subject: [PATCH 03/13] fix(io): bound delete_stream memory by flushing per-prefix batches --- crates/iceberg/src/io/file_io.rs | 86 ++++++++++++++++++++++++++++++-- 1 file changed, 82 insertions(+), 4 deletions(-) diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 10112f8458..27baff4dcc 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -274,7 +274,9 @@ impl FileIO { return self.get_storage("")?.delete_stream(paths.boxed()).await; } - // Otherwise group paths by routed storage, then batch-delete each group. + // Route by prefix, flushing bounded batches as we iterate so memory stays + // bounded on large streams (like Java's `S3FileIO.deleteFiles`). + const DELETE_BATCH_SIZE: usize = 1000; let mut groups: HashMap> = HashMap::new(); let mut paths = paths.boxed(); while let Some(path) = paths.next().await { @@ -284,12 +286,24 @@ impl FileIO { .find(|ps| path.starts_with(&ps.prefix)) .map(|ps| ps.prefix.clone()) .unwrap_or_default(); - groups.entry(key).or_default().push(path); + let buf = groups.entry(key).or_default(); + buf.push(path); + if buf.len() >= DELETE_BATCH_SIZE { + let full = std::mem::take(buf); + self.get_storage(&full[0])? + .delete_stream(stream::iter(full).boxed()) + .await?; + } } + // Flush remainders. for batch in groups.into_values() { - let storage = self.get_storage(&batch[0])?; - storage.delete_stream(stream::iter(batch).boxed()).await?; + if batch.is_empty() { + continue; + } + self.get_storage(&batch[0])? + .delete_stream(stream::iter(batch).boxed()) + .await?; } Ok(()) } @@ -808,6 +822,39 @@ mod tests { assert_eq!(prefixes, vec!["memory://a/longer/", "memory://a/"]); } + #[tokio::test] + async fn test_prefixed_config_carries_credential_values() { + // Prefix config gets the vended credentials; default config keeps only base props. + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prop("s3.region", "us-east-1") + .with_prefixed_props("s3://bucket/table", [ + ("s3.region", "us-east-1"), + ("s3.access-key-id", "vended-key"), + ("s3.secret-access-key", "vended-secret"), + ]) + .build(); + + // Default: base props, no credentials. + assert_eq!( + file_io.config().get("s3.region"), + Some(&"us-east-1".to_string()) + ); + assert_eq!(file_io.config().get("s3.access-key-id"), None); + + // Prefix: base props + vended credentials. + let prefixed = &file_io.prefixed[0].config; + assert_eq!(prefixed.get("s3.region"), Some(&"us-east-1".to_string())); + assert_eq!( + prefixed.get("s3.access-key-id"), + Some(&"vended-key".to_string()) + ); + assert_eq!( + prefixed.get("s3.secret-access-key"), + Some(&"vended-secret".to_string()) + ); + } + #[tokio::test] async fn test_get_storage_routes_by_prefix() { let factory = Arc::new(MemoryStorageFactory); @@ -860,4 +907,35 @@ mod tests { assert!(!file_io.exists(default_path).await.unwrap()); assert!(!file_io.exists(prefixed_path).await.unwrap()); } + + #[tokio::test] + async fn test_delete_stream_flushes_across_batches() { + // More than the flush threshold (1000): exercises mid-stream flush + remainder. + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prefixed_props("memory:/creds/", [("k", "v")]) + .build(); + + let n = 1050; + let mut paths = Vec::with_capacity(n); + for i in 0..n { + let p = format!("memory:/creds/f{i}.txt"); + file_io + .new_output(&p) + .unwrap() + .write("x".into()) + .await + .unwrap(); + paths.push(p); + } + + file_io + .delete_stream(futures::stream::iter(paths.clone())) + .await + .unwrap(); + + for p in &paths { + assert!(!file_io.exists(p).await.unwrap()); + } + } } From 1b2dd744d531900a01d49d14b32be266faaff66f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Fri, 10 Jul 2026 04:56:21 -0400 Subject: [PATCH 04/13] fix(rest): route vended credentials by prefix and keep them in object cache --- crates/catalog/rest/src/catalog.rs | 3 +- crates/iceberg/src/io/file_io.rs | 4 +- crates/iceberg/src/io/object_cache.rs | 11 ++++++ crates/iceberg/src/io/storage/config/mod.rs | 41 +++++++++++++++++++- crates/iceberg/src/table.rs | 42 +++++++++++++++++++++ crates/iceberg/src/transaction/mod.rs | 18 ++++++++- 6 files changed, 114 insertions(+), 5 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index ebce25d01f..9b264fd998 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -883,7 +883,8 @@ impl RestSessionCatalog { let mut builder = FileIOBuilder::new(factory).with_props(props.clone()); // Vended credentials are scoped per location prefix: give each its own - // storage so reads/writes use the matching credentials. + // storage. Paths under no vended prefix fall back to the default `props` + // above, which carry no credentials. if let Some(creds) = storage_credentials { for cred in creds { let mut prefixed = props.clone(); diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 27baff4dcc..b49ba2ea3e 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -211,6 +211,8 @@ impl FileIO { /// prefix storage if any, else the default. Built once, then cached. fn get_storage(&self, path: &str) -> Result> { // `prefixed` is sorted longest-first, so the first match is most specific. + // Selection is by longest matching string prefix, per the Iceberg REST + // spec's storage-credentials semantics (and Java's `S3FileIO`). for ps in self.prefixed.iter() { if path.starts_with(&ps.prefix) { return Self::get_or_build(&ps.storage, &self.factory, &ps.config); @@ -411,7 +413,7 @@ impl FileIOBuilder { }) .collect(); // Longest prefix first so routing picks the most specific match. - prefixed.sort_by(|a, b| b.prefix.len().cmp(&a.prefix.len())); + prefixed.sort_by_key(|item| std::cmp::Reverse(item.prefix.len())); FileIO { config: self.config, factory: self.factory, diff --git a/crates/iceberg/src/io/object_cache.rs b/crates/iceberg/src/io/object_cache.rs index dfc6c124c9..e87c361f44 100644 --- a/crates/iceberg/src/io/object_cache.rs +++ b/crates/iceberg/src/io/object_cache.rs @@ -104,6 +104,17 @@ impl ObjectCache { } } + /// Returns a cache that uses `file_io` for future cache misses. + pub(crate) fn with_file_io(mut self, file_io: FileIO) -> Self { + self.file_io = file_io; + self + } + + #[cfg(test)] + pub(crate) fn file_io(&self) -> &FileIO { + &self.file_io + } + /// Retrieves an Arc [`Manifest`] from the cache /// or retrieves one from FileIO and parses it if not present pub(crate) async fn get_manifest(&self, manifest_file: &ManifestFile) -> Result> { diff --git a/crates/iceberg/src/io/storage/config/mod.rs b/crates/iceberg/src/io/storage/config/mod.rs index d8d356de16..c7ebed7526 100644 --- a/crates/iceberg/src/io/storage/config/mod.rs +++ b/crates/iceberg/src/io/storage/config/mod.rs @@ -50,12 +50,23 @@ use serde::{Deserialize, Serialize}; /// This struct contains only configuration properties without specifying /// which storage backend to use. The storage type is determined by the /// explicit factory selection. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)] +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct StorageConfig { /// Configuration properties for the storage backend props: HashMap, } +impl std::fmt::Debug for StorageConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Property values may hold vended credentials (e.g. `s3.secret-access-key`, + // `s3.session-token`). Debug is reachable through the `FileIO`/`Table` + // derives, so print only the keys and never the secret values. + f.debug_struct("StorageConfig") + .field("keys", &self.props.keys().collect::>()) + .finish_non_exhaustive() + } +} + impl StorageConfig { /// Create a new empty StorageConfig. pub fn new() -> Self { @@ -150,6 +161,34 @@ mod tests { assert!(config.props().is_empty()); } + #[test] + fn test_debug_redacts_credential_values() { + let config = StorageConfig::from_props(HashMap::from([ + ("s3.access-key-id".to_string(), "vended-key".to_string()), + ( + "s3.secret-access-key".to_string(), + "super-secret".to_string(), + ), + ("s3.session-token".to_string(), "vended-token".to_string()), + ])); + + let rendered = format!("{config:?}"); + // Secret values must never appear in Debug output (reachable via FileIO/Table). + assert!( + !rendered.contains("super-secret"), + "leaked secret: {rendered}" + ); + assert!( + !rendered.contains("vended-token"), + "leaked token: {rendered}" + ); + // Keys stay visible so routing/config is still diagnosable. + assert!( + rendered.contains("s3.secret-access-key"), + "keys hidden: {rendered}" + ); + } + #[test] fn test_storage_config_get() { let config = StorageConfig::new().with_prop("region", "us-east-1"); diff --git a/crates/iceberg/src/table.rs b/crates/iceberg/src/table.rs index 63b5cac355..d33e3056ca 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -221,6 +221,12 @@ impl Table { /// Sets the [`Table`] `FileIO` and returns an updated instance. pub(crate) fn with_file_io(mut self, file_io: FileIO) -> Self { + self.object_cache = Arc::new( + self.object_cache + .as_ref() + .clone() + .with_file_io(file_io.clone()), + ); self.file_io = file_io; self } @@ -418,7 +424,9 @@ mod tests { use crate::ErrorKind; use crate::encryption::SensitiveBytes; use crate::encryption::kms::MemoryKeyManagementClient; + use crate::io::{FileIOBuilder, MemoryStorageFactory}; use crate::spec::TableProperties; + use crate::test_utils::test_runtime; fn load_test_metadata(filename: &str) -> TableMetadata { let path = format!( @@ -503,6 +511,40 @@ mod tests { assert_eq!(table.identifier.name(), "table"); } + #[test] + fn test_with_file_io_updates_object_cache_file_io() { + let metadata = load_test_metadata("TableMetadataV2ValidMinimal.json"); + let original_file_io = FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prop("marker", "original") + .build(); + let replacement_file_io = FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prop("marker", "replacement") + .build(); + + let table = Table::builder() + .metadata(metadata) + .identifier(TableIdent::from_strs(["ns", "table"]).unwrap()) + .file_io(original_file_io) + .runtime(test_runtime()) + .build() + .unwrap() + .with_file_io(replacement_file_io); + + assert_eq!( + table.file_io().config().get("marker").map(String::as_str), + Some("replacement") + ); + assert_eq!( + table + .object_cache() + .file_io() + .config() + .get("marker") + .map(String::as_str), + Some("replacement") + ); + } + fn make_kms() -> Arc { let kms = MemoryKeyManagementClient::new(); kms.add_master_key("master-1").unwrap(); diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index 887275d241..39e5c99164 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -232,6 +232,12 @@ impl Transaction { )?; } + // A location change moves metadata/data to a new prefix that the refresh + // load's vended credentials do not cover, so it needs a post-commit reload. + let location_changed = existing_updates + .iter() + .any(|update| matches!(update, TableUpdate::SetLocation { .. })); + let table_commit = TableCommit::builder() .ident(self.table.identifier().to_owned()) .updates(existing_updates) @@ -239,8 +245,16 @@ impl Transaction { .build(); let committed = catalog.update_table(table_commit).await?; - // Reuse the credentialed FileIO we already hold; the commit response has none. - Ok(committed.with_file_io(self.table.file_io().clone())) + if location_changed { + // The new location has its own vended credentials; the reused FileIO is + // scoped to the old prefix, so reload the table to pick them up. + catalog.load_table(committed.identifier()).await + } else { + // The commit response carries no credentials. Reuse the FileIO from the + // refresh load above (not `self.table`, which is left untouched when the + // metadata is unchanged) so freshly vended credentials are not dropped. + Ok(committed.with_file_io(refreshed.file_io().clone())) + } } } From 23302df15c343cbc6889e020299c46705725498c Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Fri, 14 Aug 2026 03:16:48 -0400 Subject: [PATCH 05/13] test(rest): cover vended credentials reaching FileIO and prefix routing --- crates/catalog/rest/src/catalog.rs | 20 ++++++++++ crates/iceberg/public-api.txt | 1 + crates/iceberg/src/io/file_io.rs | 54 ++++++++++++++++++++++++--- crates/iceberg/src/transaction/mod.rs | 35 ++++++++++++++++- 4 files changed, 104 insertions(+), 6 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 9b264fd998..9650e77f00 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -3862,6 +3862,26 @@ mod tests { table.metadata_location().unwrap() ); + // The point of the feature: the vended credentials reach the FileIO, + // scoped to the prefix the server sent them for. + let file_io = table.file_io(); + let vended = file_io.config_for("s3://warehouse/database/table/data/f.parquet"); + assert_eq!( + vended.get("s3.access-key-id"), + Some(&"vended-key-id".to_string()) + ); + assert_eq!( + vended.get("s3.session-token"), + Some(&"vended-token".to_string()) + ); + // A path outside the prefix keeps the credential-free default. + assert_eq!( + file_io + .config_for("s3://warehouse/other/f.parquet") + .get("s3.access-key-id"), + None + ); + config_mock.assert_async().await; load_table_mock.assert_async().await; } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index ef924b50a6..54421fc94a 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -728,6 +728,7 @@ pub fn iceberg::io::AzdlsConfig::deserialize<__D>(__deserializer: __D) -> core:: pub struct iceberg::io::FileIO impl iceberg::io::FileIO pub fn iceberg::io::FileIO::config(&self) -> &iceberg::io::StorageConfig +pub fn iceberg::io::FileIO::config_for(&self, path: &str) -> &iceberg::io::StorageConfig pub async fn iceberg::io::FileIO::delete(&self, path: impl core::convert::AsRef) -> iceberg::Result<()> pub async fn iceberg::io::FileIO::delete_prefix(&self, path: impl core::convert::AsRef) -> iceberg::Result<()> pub async fn iceberg::io::FileIO::delete_stream(&self, paths: impl futures_core::stream::Stream + core::marker::Send + 'static) -> iceberg::Result<()> diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index b49ba2ea3e..991be4967b 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -207,18 +207,32 @@ impl FileIO { &self.config } + /// The configuration `path` routes to: the longest-matching prefix's if + /// any, else the default. Vended credentials live on the prefix configs, + /// so this answers "which credentials apply to this path". + pub fn config_for(&self, path: &str) -> &StorageConfig { + self.route(path).1 + } + /// Get or create the storage for `path`, routing to the longest-matching /// prefix storage if any, else the default. Built once, then cached. fn get_storage(&self, path: &str) -> Result> { - // `prefixed` is sorted longest-first, so the first match is most specific. - // Selection is by longest matching string prefix, per the Iceberg REST - // spec's storage-credentials semantics (and Java's `S3FileIO`). + let (cell, config) = self.route(path); + Self::get_or_build(cell, &self.factory, config) + } + + /// The storage cell and configuration serving `path`. + /// + /// `prefixed` is sorted longest-first, so the first match is the most + /// specific one, per the Iceberg REST spec's storage-credentials semantics + /// (and Java's `S3FileIO`). + fn route(&self, path: &str) -> (&OnceLock>, &StorageConfig) { for ps in self.prefixed.iter() { if path.starts_with(&ps.prefix) { - return Self::get_or_build(&ps.storage, &self.factory, &ps.config); + return (&ps.storage, &ps.config); } } - Self::get_or_build(&self.storage, &self.factory, &self.config) + (&self.storage, &self.config) } /// Get a cached storage from `cell`, building it from `config` on first use. @@ -811,6 +825,36 @@ mod tests { assert!(deserialized.storage.get().is_some()); } + #[tokio::test] + async fn test_routes_a_path_to_its_longest_matching_prefix() { + // Overlapping prefixes: the more specific credentials must win, so + // sorting alone isn't enough — the lookup has to respect that order. + // Each prefix owns a storage cell, so identity tells them apart. + let factory = Arc::new(MemoryStorageFactory); + let file_io = FileIOBuilder::new(factory) + .with_prefixed_props("memory://bucket/data", [("k", "short")]) + .with_prefixed_props("memory://bucket/data/warehouse", [("k", "long")]) + .build(); + + let long = file_io + .get_storage("memory://bucket/data/warehouse/t/f") + .unwrap(); + let short = file_io.get_storage("memory://bucket/data/other/f").unwrap(); + let default = file_io.get_storage("memory://elsewhere/f").unwrap(); + + // Routing the shortest prefix first would collapse these two. + assert!(!Arc::ptr_eq(&long, &short)); + assert!(!Arc::ptr_eq(&short, &default)); + assert!(!Arc::ptr_eq(&long, &default)); + // The same prefix keeps serving the same storage. + assert!(Arc::ptr_eq( + &long, + &file_io + .get_storage("memory://bucket/data/warehouse/other") + .unwrap() + )); + } + #[tokio::test] async fn test_prefixed_props_sorted_by_descending_prefix_length() { let factory = Arc::new(MemoryStorageFactory); diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index 39e5c99164..e6da22f261 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -267,7 +267,7 @@ mod tests { use std::sync::atomic::{AtomicU32, Ordering}; use crate::catalog::MockCatalog; - use crate::io::FileIO; + use crate::io::{FileIO, FileIOBuilder, MemoryStorageFactory}; use crate::memory::tests::new_memory_catalog; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, Literal, Struct, TableMetadata, @@ -411,6 +411,39 @@ mod tests { .unwrap() } + #[tokio::test] + async fn test_commit_keeps_the_refreshed_file_io() { + // The commit response carries no vended credentials, so the FileIO from + // the refresh load has to survive; taking the committed table's own + // would silently drop them. + let refreshed = make_v2_table().with_file_io( + FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prefixed_props("memory://warehouse", [("s3.access-key-id", "vended")]) + .build(), + ); + let mut mock_catalog = MockCatalog::new(); + mock_catalog.expect_load_table().returning_st(move |_| { + let refreshed = refreshed.clone(); + Box::pin(async move { Ok(refreshed) }) + }); + mock_catalog + .expect_update_table() + .returning_st(|_| Box::pin(async move { Ok(make_v2_table()) })); + + let table = create_test_transaction(&make_v2_table()) + .commit(&mock_catalog) + .await + .unwrap(); + + assert_eq!( + table + .file_io() + .config_for("memory://warehouse/t/f") + .get("s3.access-key-id"), + Some(&"vended".to_string()), + ); + } + /// Helper function to set up a mock catalog with retryable errors fn setup_mock_catalog_with_retryable_errors( success_after_attempts: Option, From eaa5f1b26c8f8e41fbb7765c1f9fb32f7edae051 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Fri, 14 Aug 2026 10:15:45 -0400 Subject: [PATCH 06/13] docs(io): note that prefix matching is raw-string, as in Java --- crates/iceberg/src/io/file_io.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 991be4967b..8a27aade24 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -224,8 +224,10 @@ impl FileIO { /// The storage cell and configuration serving `path`. /// /// `prefixed` is sorted longest-first, so the first match is the most - /// specific one, per the Iceberg REST spec's storage-credentials semantics - /// (and Java's `S3FileIO`). + /// specific one, per the Iceberg REST spec's storage-credentials semantics. + /// + /// Matching is on the raw string, so `s3://bucket/data` also serves + /// `s3://bucket/database/` — as in Java's `S3FileIO.clientForStoragePath`. fn route(&self, path: &str) -> (&OnceLock>, &StorageConfig) { for ps in self.prefixed.iter() { if path.starts_with(&ps.prefix) { From 7decc433d183d2aa56ef85bc05fa9c4a1ce8d71b Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 2 Sep 2026 01:46:15 -0400 Subject: [PATCH 07/13] fix(rest): keep vended credentials on register_table --- crates/catalog/rest/src/catalog.rs | 65 ++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 3 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 9650e77f00..1a69fc0904 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -1408,8 +1408,15 @@ impl SessionCatalog for RestSessionCatalog { "Metadata location missing in `register_table` response!", ))?; + let mut base_config = response.config.clone(); + base_config.extend(self.user_config.props.clone()); + let file_io = self - .load_file_io(Some(metadata_location), None, None) + .load_file_io( + Some(metadata_location), + Some(base_config), + response.storage_credentials.as_deref(), + ) .await?; let mut table_builder = Table::builder() @@ -1488,8 +1495,10 @@ impl SessionCatalog for RestSessionCatalog { } }; - // The commit response carries no credentials, so build a plain FileIO; - // the transaction layer reuses the credentialed one it already holds. + // The commit response carries no credentials, so this FileIO has only the + // catalog-level config. `Transaction::do_commit` swaps in the credentialed + // one from its pre-commit load; only a direct `update_table` caller sees + // this plain one. let file_io = self .load_file_io(Some(&response.metadata_location), None, None) .await?; @@ -4404,6 +4413,56 @@ mod tests { register_table_mock.assert_async().await; } + /// The register response is a `LoadTableResult`, so the vended credentials + /// it carries have to reach the table's FileIO like `load_table`'s do. + #[tokio::test] + async fn test_register_table_uses_vended_credentials() { + let mut server = Server::new_async().await; + let config_mock = create_config_mock(&mut server).await; + let register_table_mock = server + .mock("POST", "/v1/namespaces/ns1/register") + .with_status(200) + .with_body_from_file(format!( + "{}/testdata/{}", + env!("CARGO_MANIFEST_DIR"), + "load_table_response_with_credentials.json" + )) + .create_async() + .await; + + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); + let table_ident = TableIdent::from_strs(["ns1", "test1"]).unwrap(); + let table = catalog + .register_table( + &SessionContext::empty(), + &table_ident, + "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json".to_string(), + ) + .await + .unwrap(); + + let vended = table + .file_io() + .config_for("s3://warehouse/database/table/data/f.parquet"); + assert_eq!( + vended.get("s3.access-key-id").map(String::as_str), + Some("vended-key-id") + ); + // Scoped to the prefix the server sent them for. + let outside = table + .file_io() + .config_for("s3://other-bucket/data/f.parquet"); + assert_eq!(outside.get("s3.access-key-id"), None); + // The response's table config reaches the FileIO too, as in load_table. + assert_eq!( + table.file_io().config().get("region").map(String::as_str), + Some("us-west-2") + ); + + config_mock.assert_async().await; + register_table_mock.assert_async().await; + } + #[tokio::test] async fn test_register_table_404() { let mut server = Server::new_async().await; From 41f74e1013674eb840b36b2dd4872b449c2e3d99 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 9 Sep 2026 00:56:49 -0400 Subject: [PATCH 08/13] test(io): prefixed credentials survive FileIO serialization --- crates/iceberg/src/io/file_io.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 8a27aade24..807da449c4 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -795,6 +795,31 @@ mod tests { assert!(deserialized.storage.get().is_some()); } + /// Per-prefix vended credentials must survive serialization, or a worker + /// that receives this FileIO cannot read the data they were vended for. + #[tokio::test] + async fn test_prefixed_credentials_survive_serialization_roundtrip() { + let file_io = FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prop("s3.access-key-id", "default-key") + .with_prefixed_props("memory://warehouse/t", [("s3.access-key-id", "vended-key")]) + .build(); + + let deserialized = FileIO::deserialize_all(&file_io.serialize_all().unwrap()).unwrap(); + + assert_eq!( + deserialized + .config_for("memory://warehouse/t/data/f.parquet") + .get("s3.access-key-id"), + Some(&"vended-key".to_string()) + ); + assert_eq!( + deserialized + .config_for("memory://elsewhere/f.parquet") + .get("s3.access-key-id"), + Some(&"default-key".to_string()) + ); + } + #[tokio::test] async fn test_local_fs_file_io_serialization_roundtrip() { let tmp_dir = TempDir::new().unwrap(); From c7fcc2a38dbfb769df9bda1e3779c5e473ef066f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Thu, 10 Sep 2026 02:11:44 -0400 Subject: [PATCH 09/13] docs(rest): note that only Transaction can reach update_table --- crates/catalog/rest/src/catalog.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 1a69fc0904..7731e207b8 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -1497,8 +1497,8 @@ impl SessionCatalog for RestSessionCatalog { // The commit response carries no credentials, so this FileIO has only the // catalog-level config. `Transaction::do_commit` swaps in the credentialed - // one from its pre-commit load; only a direct `update_table` caller sees - // this plain one. + // one from its pre-commit load, and it is the only caller there can be: + // `TableCommit` is buildable inside the `iceberg` crate alone. let file_io = self .load_file_io(Some(&response.metadata_location), None, None) .await?; From 2e27d9f4bcd3f3a3206c5420c5fe221ec6160a19 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Sun, 20 Sep 2026 15:06:03 -0400 Subject: [PATCH 10/13] fix(io): keep cached storages out of FileIO's Debug --- crates/iceberg/src/io/file_io.rs | 136 ++++++++++++++++++++++++++++++- 1 file changed, 132 insertions(+), 4 deletions(-) diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 807da449c4..6dbb6d92e9 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -60,7 +60,7 @@ use crate::Result; /// .with_prop("key", "value") /// .build(); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct FileIO { /// Storage configuration containing properties config: StorageConfig, @@ -74,13 +74,33 @@ pub struct FileIO { } /// A storage scoped to a location `prefix`, lazily built from its own config. -#[derive(Debug)] struct PrefixedStorage { prefix: String, config: StorageConfig, storage: OnceLock>, } +// A backend's Debug may print the raw credential map that `StorageConfig` +// redacts, so neither cached storage is shown. +impl std::fmt::Debug for FileIO { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FileIO") + .field("config", &self.config) + .field("factory", &self.factory) + .field("prefixed", &self.prefixed) + .finish_non_exhaustive() + } +} + +impl std::fmt::Debug for PrefixedStorage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrefixedStorage") + .field("prefix", &self.prefix) + .field("config", &self.config) + .finish_non_exhaustive() + } +} + mod _serde { use std::sync::Arc; @@ -599,10 +619,17 @@ mod tests { use bytes::Bytes; use futures::AsyncReadExt; use futures::io::AllowStdIo; + use futures::stream::BoxStream; + use serde::{Deserialize, Serialize}; use tempfile::TempDir; - use super::{FileIO, FileIOBuilder}; - use crate::io::{LocalFsStorageFactory, MemoryStorageFactory}; + use super::{ + FileIO, FileIOBuilder, InputFile, OutputFile, Storage, StorageConfig, StorageFactory, + }; + use crate::Result; + use crate::io::{ + FileMetadata, FileRead, FileWrite, LocalFsStorageFactory, MemoryStorageFactory, + }; fn create_local_file_io() -> FileIO { FileIO::new_with_fs() @@ -820,6 +847,107 @@ mod tests { ); } + /// Stands in for a backend whose Debug prints its raw props. + #[derive(Serialize, Deserialize)] + struct LeakyStorage(String); + + impl std::fmt::Debug for LeakyStorage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "LeakyStorage {{ secret: {} }}", self.0) + } + } + + #[typetag::serde] + #[async_trait::async_trait] + impl Storage for LeakyStorage { + async fn exists(&self, _: &str) -> Result { + Ok(false) + } + async fn metadata(&self, _: &str) -> Result { + unimplemented!() + } + async fn read(&self, _: &str) -> Result { + unimplemented!() + } + async fn reader(&self, _: &str) -> Result> { + unimplemented!() + } + async fn write(&self, _: &str, _: Bytes) -> Result<()> { + unimplemented!() + } + async fn writer(&self, _: &str) -> Result> { + unimplemented!() + } + async fn delete(&self, _: &str) -> Result<()> { + unimplemented!() + } + async fn delete_prefix(&self, _: &str) -> Result<()> { + unimplemented!() + } + async fn delete_stream(&self, _: BoxStream<'static, String>) -> Result<()> { + unimplemented!() + } + fn new_input(&self, _: &str) -> Result { + unimplemented!() + } + fn new_output(&self, _: &str) -> Result { + unimplemented!() + } + } + + #[derive(Debug, Serialize, Deserialize)] + struct LeakyFactory; + + #[typetag::serde] + impl StorageFactory for LeakyFactory { + fn build(&self, config: &StorageConfig) -> Result> { + let secret = config + .get("s3.secret-access-key") + .cloned() + .unwrap_or_default(); + Ok(Arc::new(LeakyStorage(secret))) + } + } + + /// Once a storage is initialized its own Debug is reachable through + /// `FileIO`; the vended secret must not come out that way. + #[tokio::test] + async fn test_debug_omits_initialized_storages() { + let file_io = FileIOBuilder::new(Arc::new(LeakyFactory)) + .with_prop("s3.secret-access-key", "DEFAULT-SECRET") + .with_prefixed_props("memory://warehouse/t", [( + "s3.secret-access-key", + "VENDED-SECRET", + )]) + .build(); + // Initialize both the default and the prefixed storage. + file_io.exists("memory://elsewhere/f").await.unwrap(); + file_io.exists("memory://warehouse/t/f").await.unwrap(); + + let debug = format!("{file_io:?}"); + assert!(!debug.contains("VENDED-SECRET"), "{debug}"); + assert!(!debug.contains("DEFAULT-SECRET"), "{debug}"); + // Still informative: the prefix and the redacted config keys show. + assert!(debug.contains("memory://warehouse/t"), "{debug}"); + assert!(debug.contains("s3.secret-access-key"), "{debug}"); + } + + /// Two prefixes where one nests in the other: the longest-first order + /// that routing depends on has to survive the roundtrip. + #[tokio::test] + async fn test_overlapping_prefixes_survive_serialization_roundtrip() { + let file_io = FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prefixed_props("memory://w/t", [("s3.access-key-id", "outer")]) + .with_prefixed_props("memory://w/t/nested", [("s3.access-key-id", "inner")]) + .build(); + + let deserialized = FileIO::deserialize_all(&file_io.serialize_all().unwrap()).unwrap(); + + let key = |p: &str| deserialized.config_for(p).get("s3.access-key-id").cloned(); + assert_eq!(key("memory://w/t/nested/f"), Some("inner".to_string())); + assert_eq!(key("memory://w/t/other/f"), Some("outer".to_string())); + } + #[tokio::test] async fn test_local_fs_file_io_serialization_roundtrip() { let tmp_dir = TempDir::new().unwrap(); From 193e77fadaae4be65cd26977152e21da707c404f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Sun, 20 Sep 2026 15:06:09 -0400 Subject: [PATCH 11/13] fix(transaction): act with the refreshed FileIO and reload on any location change --- crates/iceberg/src/transaction/mod.rs | 109 ++++++++++++++++++++++++-- 1 file changed, 104 insertions(+), 5 deletions(-) diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index e6da22f261..e5f7235bff 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -217,7 +217,9 @@ impl Transaction { self.table = refreshed.clone(); } - let mut current_table = self.table.clone(); + // Actions read and write manifests, so they need the credentials the + // refresh load vended, whether or not the base was stale. + let mut current_table = self.table.clone().with_file_io(refreshed.file_io().clone()); let mut existing_updates: Vec = vec![]; let mut existing_requirements: Vec = vec![]; @@ -232,9 +234,7 @@ impl Transaction { )?; } - // A location change moves metadata/data to a new prefix that the refresh - // load's vended credentials do not cover, so it needs a post-commit reload. - let location_changed = existing_updates + let moves_location = existing_updates .iter() .any(|update| matches!(update, TableUpdate::SetLocation { .. })); @@ -245,6 +245,12 @@ impl Transaction { .build(); let committed = catalog.update_table(table_commit).await?; + // A location change moves metadata/data to a new prefix that the refresh + // load's vended credentials do not cover, so it needs a post-commit + // reload. Another writer may have moved the table meanwhile: a + // property-only commit carries no requirement that would catch that. + let location_changed = + moves_location || committed.metadata().location() != refreshed.metadata().location(); if location_changed { // The new location has its own vended credentials; the reused FileIO is // scoped to the old prefix, so reload the table to pick them up. @@ -266,6 +272,8 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; + use async_trait::async_trait; + use crate::catalog::MockCatalog; use crate::io::{FileIO, FileIOBuilder, MemoryStorageFactory}; use crate::memory::tests::new_memory_catalog; @@ -275,8 +283,9 @@ mod tests { }; use crate::table::Table; use crate::test_utils::{make_encrypted_table, test_runtime}; + use crate::transaction::action::{ActionCommit, TransactionAction}; use crate::transaction::{ApplyTransactionAction, Transaction}; - use crate::{Catalog, Error, ErrorKind, TableCreation, TableIdent}; + use crate::{Catalog, Error, ErrorKind, Result, TableCreation, TableIdent}; pub fn make_v1_table() -> Table { let file = File::open(format!( @@ -444,6 +453,96 @@ mod tests { ); } + fn table_with_marker(marker: &str) -> Table { + make_v2_table().with_file_io( + FileIOBuilder::new(Arc::new(MemoryStorageFactory)) + .with_prop("marker", marker) + .build(), + ) + } + + fn marker_of(table: &Table) -> String { + table + .file_io() + .config() + .get("marker") + .cloned() + .unwrap_or_default() + } + + /// Records which FileIO the action was handed. + struct RecordingAction(Arc>>); + + #[async_trait] + impl TransactionAction for RecordingAction { + async fn commit(self: Arc, table: &Table) -> Result { + *self.0.lock().unwrap() = Some(marker_of(table)); + Ok(ActionCommit::new(vec![], vec![])) + } + } + + /// The refresh load may carry rotated credentials while the metadata is + /// unchanged; actions still have to run with those, not the stale ones. + #[tokio::test] + async fn test_actions_run_with_the_refreshed_file_io() { + let seen = Arc::new(std::sync::Mutex::new(None)); + let mut mock_catalog = MockCatalog::new(); + mock_catalog + .expect_load_table() + .returning_st(|_| Box::pin(async { Ok(table_with_marker("refreshed")) })); + mock_catalog + .expect_update_table() + .returning_st(|_| Box::pin(async { Ok(make_v2_table()) })); + + let tx = RecordingAction(seen.clone()) + .apply(Transaction::new(&table_with_marker("stale"))) + .unwrap(); + tx.commit(&mock_catalog).await.unwrap(); + + assert_eq!(seen.lock().unwrap().as_deref(), Some("refreshed")); + } + + /// Another writer moved the table between our refresh and our commit; the + /// committed metadata says so even though we sent no `SetLocation`. + #[tokio::test] + async fn test_commit_reloads_when_someone_else_moved_the_table() { + let loads = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut mock_catalog = MockCatalog::new(); + let counter = loads.clone(); + mock_catalog.expect_load_table().returning_st(move |_| { + let n = counter.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { Ok(table_with_marker(if n == 0 { "first" } else { "second" })) }) + }); + mock_catalog.expect_update_table().returning_st(|_| { + Box::pin(async { + let base = make_v2_table(); + let moved = base + .metadata() + .clone() + .into_builder(None) + .set_location("s3://moved-by-someone-else/table".to_string()) + .build() + .unwrap() + .metadata; + Ok(Table::builder() + .identifier(base.identifier().clone()) + .metadata(moved) + .file_io(base.file_io().clone()) + .runtime(test_runtime()) + .build() + .unwrap()) + }) + }); + + let table = create_test_transaction(&make_v2_table()) + .commit(&mock_catalog) + .await + .unwrap(); + + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(marker_of(&table), "second"); + } + /// Helper function to set up a mock catalog with retryable errors fn setup_mock_catalog_with_retryable_errors( success_after_attempts: Option, From e421454d1b1bc5f700163b624eec4e70d5533c4f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Sun, 20 Sep 2026 19:33:15 -0400 Subject: [PATCH 12/13] fix(transaction): keep the table's initialized storage when a reload changes nothing --- crates/iceberg/src/io/file_io.rs | 12 ++++++ crates/iceberg/src/transaction/mod.rs | 62 ++++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 6dbb6d92e9..94d7f26a91 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -227,6 +227,18 @@ impl FileIO { &self.config } + /// Whether `other` would route and authenticate identically. Used to keep + /// an initialized backend when a reload brought back the same settings. + pub(crate) fn same_routing_as(&self, other: &FileIO) -> bool { + self.config == other.config + && self.prefixed.len() == other.prefixed.len() + && self + .prefixed + .iter() + .zip(other.prefixed.iter()) + .all(|(a, b)| a.prefix == b.prefix && a.config == b.config) + } + /// The configuration `path` routes to: the longest-matching prefix's if /// any, else the default. Vended credentials live on the prefix configs, /// so this answers "which credentials apply to this path". diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index e5f7235bff..c27cb6bbc3 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -218,8 +218,13 @@ impl Transaction { } // Actions read and write manifests, so they need the credentials the - // refresh load vended, whether or not the base was stale. - let mut current_table = self.table.clone().with_file_io(refreshed.file_io().clone()); + // refresh load vended, whether or not the base was stale. Identical + // settings keep the original: a rebuilt FileIO would drop a backend + // the table has already initialized. + let mut current_table = self.table.clone(); + if !current_table.file_io().same_routing_as(refreshed.file_io()) { + current_table = current_table.with_file_io(refreshed.file_io().clone()); + } let mut existing_updates: Vec = vec![]; let mut existing_requirements: Vec = vec![]; @@ -502,6 +507,59 @@ mod tests { assert_eq!(seen.lock().unwrap().as_deref(), Some("refreshed")); } + /// A reload that changed nothing must not cost the table its initialized + /// backend: `MemoryStorageFactory` builds a fresh, empty store each time. + #[tokio::test] + async fn test_unchanged_settings_keep_the_initialized_storage() { + let table = table_with_marker("same"); + table + .file_io() + .new_output("memory://warehouse/manifest") + .unwrap() + .write("written-before-commit".into()) + .await + .unwrap(); + + let seen = Arc::new(std::sync::Mutex::new(None)); + let action_io = seen.clone(); + let mut mock_catalog = MockCatalog::new(); + // A reload rebuilds the FileIO, as a REST catalog does. + mock_catalog + .expect_load_table() + .returning_st(|_| Box::pin(async { Ok(table_with_marker("same")) })); + mock_catalog + .expect_update_table() + .returning_st(|_| Box::pin(async { Ok(make_v2_table()) })); + + let tx = ReadingAction(action_io) + .apply(Transaction::new(&table)) + .unwrap(); + tx.commit(&mock_catalog).await.unwrap(); + + assert_eq!( + seen.lock().unwrap().as_deref(), + Some("written-before-commit") + ); + } + + /// Reads a file the table wrote before the transaction started. + struct ReadingAction(Arc>>); + + #[async_trait] + impl TransactionAction for ReadingAction { + async fn commit(self: Arc, table: &Table) -> Result { + let read = table + .file_io() + .new_input("memory://warehouse/manifest")? + .read() + .await + .ok() + .map(|b| String::from_utf8_lossy(&b).to_string()); + *self.0.lock().unwrap() = read; + Ok(ActionCommit::new(vec![], vec![])) + } + } + /// Another writer moved the table between our refresh and our commit; the /// committed metadata says so even though we sent no `SetLocation`. #[tokio::test] From 0334dafcb78e5993c448a29f90397300b0769119 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Mon, 28 Sep 2026 00:14:35 -0400 Subject: [PATCH 13/13] fix(transaction): return the actions' FileIO and never replay a durable commit --- crates/iceberg/src/transaction/mod.rs | 107 +++++++++++++++++++------- 1 file changed, 80 insertions(+), 27 deletions(-) diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index c27cb6bbc3..cc6c6804a7 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -259,13 +259,20 @@ impl Transaction { if location_changed { // The new location has its own vended credentials; the reused FileIO is // scoped to the old prefix, so reload the table to pick them up. - catalog.load_table(committed.identifier()).await - } else { - // The commit response carries no credentials. Reuse the FileIO from the - // refresh load above (not `self.table`, which is left untouched when the - // metadata is unchanged) so freshly vended credentials are not dropped. - Ok(committed.with_file_io(refreshed.file_io().clone())) + match catalog.load_table(committed.identifier()).await { + Ok(reloaded) => return Ok(reloaded), + // The commit is already durable: returning this error would let the + // retry loop replay it. Keep the old credentials; a reload fixes them. + Err(e) => tracing::warn!( + "committed {} but could not reload it for the new location: {e}", + committed.identifier() + ), + } } + // The commit response carries no credentials, so keep the FileIO the + // actions ran with: the refreshed one, or the original when its settings + // were unchanged and it holds an initialized backend. + Ok(committed.with_file_io(current_table.file_io().clone())) } } @@ -534,12 +541,21 @@ mod tests { let tx = ReadingAction(action_io) .apply(Transaction::new(&table)) .unwrap(); - tx.commit(&mock_catalog).await.unwrap(); + let committed = tx.commit(&mock_catalog).await.unwrap(); assert_eq!( seen.lock().unwrap().as_deref(), Some("written-before-commit") ); + // The returned table keeps that backend too, not the reload's empty one. + let after = committed + .file_io() + .new_input("memory://warehouse/manifest") + .unwrap() + .read() + .await + .unwrap(); + assert_eq!(&after[..], b"written-before-commit"); } /// Reads a file the table wrote before the transaction started. @@ -560,6 +576,60 @@ mod tests { } } + /// What `update_table` returns once the table root has moved. + fn moved_table() -> Table { + let base = make_v2_table(); + let moved = base + .metadata() + .clone() + .into_builder(None) + .set_location("s3://moved-by-someone-else/table".to_string()) + .build() + .unwrap() + .metadata; + Table::builder() + .identifier(base.identifier().clone()) + .metadata(moved) + .file_io(base.file_io().clone()) + .runtime(test_runtime()) + .build() + .unwrap() + } + + /// The commit is durable before the reload runs, so a failed reload must + /// not send the transaction round the retry loop to be committed again. + #[tokio::test] + async fn test_a_failed_reload_does_not_replay_the_commit() { + let loads = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut mock_catalog = MockCatalog::new(); + let counter = loads.clone(); + mock_catalog.expect_load_table().returning_st(move |_| { + let n = counter.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + if n == 1 { + // The post-commit reload. + Err(Error::new(ErrorKind::Unexpected, "503").with_retryable(true)) + } else { + Ok(make_v2_table()) + } + }) + }); + mock_catalog + .expect_update_table() + .times(1) + .returning_st(|_| Box::pin(async { Ok(moved_table()) })); + + let table = create_test_transaction(&make_v2_table()) + .commit(&mock_catalog) + .await + .unwrap(); + + assert_eq!( + table.metadata().location(), + "s3://moved-by-someone-else/table" + ); + } + /// Another writer moved the table between our refresh and our commit; the /// committed metadata says so even though we sent no `SetLocation`. #[tokio::test] @@ -571,26 +641,9 @@ mod tests { let n = counter.fetch_add(1, Ordering::SeqCst); Box::pin(async move { Ok(table_with_marker(if n == 0 { "first" } else { "second" })) }) }); - mock_catalog.expect_update_table().returning_st(|_| { - Box::pin(async { - let base = make_v2_table(); - let moved = base - .metadata() - .clone() - .into_builder(None) - .set_location("s3://moved-by-someone-else/table".to_string()) - .build() - .unwrap() - .metadata; - Ok(Table::builder() - .identifier(base.identifier().clone()) - .metadata(moved) - .file_io(base.file_io().clone()) - .runtime(test_runtime()) - .build() - .unwrap()) - }) - }); + mock_catalog + .expect_update_table() + .returning_st(|_| Box::pin(async { Ok(moved_table()) })); let table = create_test_transaction(&make_v2_table()) .commit(&mock_catalog)