From 38bc779e2999946c366b4154fd873c1f8a7b9735 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Mon, 15 Jun 2026 18:28:21 +0800 Subject: [PATCH 1/4] feat(rest): use vended storage credentials from LoadTable response --- crates/catalog/rest/src/catalog.rs | 43 +++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 13 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 8642b32d22..97e27171cb 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -45,7 +45,7 @@ use crate::endpoint::{Endpoint, V1_NAMESPACE_EXISTS, V1_TABLE_EXISTS}; use crate::types::{ CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult, - NamespaceResponse, RegisterTableRequest, RenameTableRequest, + NamespaceResponse, RegisterTableRequest, RenameTableRequest, StorageCredential, }; /// REST catalog URI @@ -803,6 +803,7 @@ impl Catalog for RestCatalog { let request = context .client .request(Method::POST, context.config.tables_endpoint(namespace)) + .header("X-Iceberg-Access-Delegation", "vended-credentials") .json(&CreateTableRequest { name: creation.name, location: creation.location, @@ -846,11 +847,7 @@ impl Catalog for RestCatalog { "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)) @@ -883,6 +880,8 @@ impl Catalog for RestCatalog { let request = context .client .request(Method::GET, context.config.table_endpoint(table_ident)) + // Opt in to vended storage credentials. + .header("X-Iceberg-Access-Delegation", "vended-credentials") .build()?; let http_response = context.client.query_catalog(request).await?; @@ -906,11 +905,7 @@ impl Catalog for RestCatalog { } }; - 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)) @@ -1122,9 +1117,12 @@ impl Catalog for RestCatalog { } }; + // 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(commit.identifier()) + .await? + .file_io() + .clone(); let mut table_builder = Table::builder() .identifier(commit.identifier().clone()) @@ -1139,6 +1137,23 @@ impl Catalog for RestCatalog { } } +/// 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; @@ -2888,6 +2903,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) @@ -2896,6 +2912,7 @@ mod tests { env!("CARGO_MANIFEST_DIR"), "load_table_response.json" )) + .expect(2) .create_async() .await; From 5221745ab0627248000041e0fdbb73f62469f9b5 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 16 Jun 2026 00:29:12 +0800 Subject: [PATCH 2/4] feat(rest): per-prefix vended credentials with opt-in delegation header --- crates/catalog/rest/src/catalog.rs | 155 ++++++++++---- .../load_table_response_with_credentials.json | 78 ++++++++ crates/iceberg/public-api.txt | 1 + crates/iceberg/src/io/file_io.rs | 189 +++++++++++++++--- crates/iceberg/src/table.rs | 6 + crates/iceberg/src/transaction/mod.rs | 4 +- 6 files changed, 375 insertions(+), 58 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 97e27171cb..fdf2b6eb7a 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -509,6 +509,7 @@ impl RestCatalog { &self, metadata_location: Option<&str>, extra_config: Option>, + storage_credentials: Option<&[StorageCredential]>, ) -> Result { let mut props = self.context().await?.config.props.clone(); if let Some(config) = extra_config { @@ -541,9 +542,19 @@ impl RestCatalog { ) })?; - 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()) } /// Invalidate the current token without generating a new one. On the next request, the client @@ -803,7 +814,6 @@ impl Catalog for RestCatalog { let request = context .client .request(Method::POST, context.config.tables_endpoint(namespace)) - .header("X-Iceberg-Access-Delegation", "vended-credentials") .json(&CreateTableRequest { name: creation.name, location: creation.location, @@ -847,10 +857,15 @@ impl Catalog for RestCatalog { "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() @@ -877,11 +892,12 @@ impl Catalog for RestCatalog { async fn load_table(&self, table_ident: &TableIdent) -> Result { let context = self.context().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 = context .client .request(Method::GET, context.config.table_endpoint(table_ident)) - // Opt in to vended storage credentials. - .header("X-Iceberg-Access-Delegation", "vended-credentials") .build()?; let http_response = context.client.query_catalog(request).await?; @@ -905,10 +921,15 @@ impl Catalog for RestCatalog { } }; - 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() @@ -1043,7 +1064,9 @@ impl Catalog for RestCatalog { "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()) @@ -1117,12 +1140,11 @@ impl Catalog for RestCatalog { } }; - // 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(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()) @@ -1137,23 +1159,6 @@ impl Catalog for RestCatalog { } } -/// 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; @@ -2686,6 +2691,88 @@ 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( + RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + ); + + 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( + RestCatalogConfig::builder().uri(server.url()).build(), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + ); + + 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; @@ -2903,7 +2990,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) @@ -2912,7 +2999,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 610675c83d..0c989dbfd7 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -739,6 +739,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 cd0a4434c4..549133bfe7 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>, } impl FileIO { @@ -78,6 +90,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(MemoryStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -89,6 +102,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(LocalFsStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -97,24 +111,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. @@ -123,7 +144,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. @@ -138,7 +159,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. @@ -150,7 +173,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. @@ -159,7 +204,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. @@ -168,7 +213,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. @@ -177,7 +222,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()) } } @@ -191,6 +236,8 @@ pub struct FileIOBuilder { factory: Arc, /// Storage configuration config: StorageConfig, + /// Per-location-prefix configs (prefix, config). + prefixed: Vec<(String, StorageConfig)>, } impl FileIOBuilder { @@ -199,6 +246,7 @@ impl FileIOBuilder { Self { factory, config: StorageConfig::new(), + prefixed: Vec::new(), } } @@ -219,6 +267,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 @@ -226,10 +291,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), } } } @@ -544,4 +621,70 @@ mod tests { assert_eq!(file_io.config().get("key1"), Some(&"value1".to_string())); assert_eq!(file_io.config().get("key2"), Some(&"value2".to_string())); } + + #[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 31feade038..2885e33644 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -220,6 +220,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 7b8e02da36..cfde2aac40 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -246,7 +246,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 d4c1b362e0a4d2277a86853c0635aad97de1167b Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 17 Jun 2026 17:05:34 +0800 Subject: [PATCH 3/4] 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 549133bfe7..61dbcee248 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -178,7 +178,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 { @@ -188,12 +190,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(()) } @@ -635,6 +649,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); @@ -687,4 +734,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 da370e8248163923388e65e3f2fc9fc6cd26867f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Fri, 10 Jul 2026 04:56:21 -0400 Subject: [PATCH 4/4] fix(rest): route vended credentials by prefix and keep them in object cache --- crates/catalog/rest/src/catalog.rs | 5 ++- 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, 116 insertions(+), 5 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index fdf2b6eb7a..37125bb630 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -545,7 +545,8 @@ impl RestCatalog { 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(); @@ -2723,6 +2724,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table = catalog @@ -2762,6 +2764,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); catalog diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs index 61dbcee248..4480f44e6c 100644 --- a/crates/iceberg/src/io/file_io.rs +++ b/crates/iceberg/src/io/file_io.rs @@ -115,6 +115,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); @@ -315,7 +317,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 9d7815569b..29af0169f6 100644 --- a/crates/iceberg/src/io/object_cache.rs +++ b/crates/iceberg/src/io/object_cache.rs @@ -95,6 +95,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 2350aab6dd..d75f31ae1e 100644 --- a/crates/iceberg/src/io/storage/config/mod.rs +++ b/crates/iceberg/src/io/storage/config/mod.rs @@ -51,12 +51,23 @@ use serde::{Deserialize, Serialize}; /// 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 { @@ -151,6 +162,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 2885e33644..a273a5c819 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -222,6 +222,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 } @@ -416,7 +422,9 @@ mod tests { use super::*; 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!( @@ -501,6 +509,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 cfde2aac40..dbdb0dfef8 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -240,6 +240,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) @@ -247,8 +253,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())) + } } }