diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index aa36a074e7..7731e207b8 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 @@ -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,20 @@ 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. 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(); + prefixed.extend(cred.config.clone()); + builder = builder.with_prefixed_props(cred.prefix.clone(), prefixed); + } + } + + Ok(builder.build()) } } @@ -1181,14 +1193,15 @@ 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 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() @@ -1219,6 +1232,9 @@ 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 @@ -1245,14 +1261,15 @@ impl SessionCatalog for RestSessionCatalog { } }; - let config = response - .config - .into_iter() - .chain(self.user_config.props.clone()) - .collect(); + 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() @@ -1391,7 +1408,16 @@ impl SessionCatalog for RestSessionCatalog { "Metadata location missing in `register_table` response!", ))?; - let file_io = self.load_file_io(Some(metadata_location), None).await?; + 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(base_config), + response.storage_credentials.as_deref(), + ) + .await?; let mut table_builder = Table::builder() .identifier(table_ident.clone()) @@ -1469,8 +1495,12 @@ 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, 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) + .load_file_io(Some(&response.metadata_location), None, None) .await?; let mut table_builder = Table::builder() @@ -3794,6 +3824,114 @@ 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() + ); + + // 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; + } + + #[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; @@ -4009,6 +4147,7 @@ mod tests { let config_mock = create_config_mock(&mut server).await; + // 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) @@ -4017,6 +4156,7 @@ mod tests { env!("CARGO_MANIFEST_DIR"), "load_table_response.json" )) + .expect(1) .create_async() .await; @@ -4273,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; 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..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<()> @@ -747,6 +748,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..94d7f26a91 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, @@ -59,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, @@ -67,6 +68,37 @@ 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. +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 { @@ -80,12 +112,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 +147,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(MemoryStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -109,6 +159,7 @@ impl FileIO { config: StorageConfig::new(), factory: Arc::new(LocalFsStorageFactory), storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(Vec::new()), } } @@ -131,6 +182,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 +199,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 +227,61 @@ 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()); - } + /// 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". + pub fn config_for(&self, path: &str) -> &StorageConfig { + self.route(path).1 + } - // Build the storage - let storage = self.factory.build(&self.config)?; + /// 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> { + let (cell, config) = self.route(path); + Self::get_or_build(cell, &self.factory, config) + } - // Try to set it (another thread might have set it first) - let _ = self.storage.set(storage.clone()); + /// 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. + /// + /// 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) { + return (&ps.storage, &ps.config); + } + } + (&self.storage, &self.config) + } - // 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 +290,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 +305,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 +319,43 @@ 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; + } + + // 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 { + let key = self + .prefixed + .iter() + .find(|ps| path.starts_with(&ps.prefix)) + .map(|ps| ps.prefix.clone()) + .unwrap_or_default(); + 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() { + if batch.is_empty() { + continue; + } + self.get_storage(&batch[0])? + .delete_stream(stream::iter(batch).boxed()) + .await?; + } + Ok(()) } /// Check file exists. @@ -215,7 +364,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 +373,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 +382,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 +396,8 @@ pub struct FileIOBuilder { factory: Arc, /// Storage configuration config: StorageConfig, + /// Per-location-prefix configs (prefix, config). + prefixed: Vec<(String, StorageConfig)>, } impl FileIOBuilder { @@ -255,6 +406,7 @@ impl FileIOBuilder { Self { factory, config: StorageConfig::new(), + prefixed: Vec::new(), } } @@ -275,6 +427,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 +451,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_key(|item| std::cmp::Reverse(item.prefix.len())); FileIO { config: self.config, factory: self.factory, storage: Arc::new(OnceLock::new()), + prefixed: Arc::new(prefixed), } } } @@ -450,10 +631,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() @@ -646,6 +834,132 @@ 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()) + ); + } + + /// 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(); @@ -677,4 +991,164 @@ 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); + 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_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); + 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()); + } + + #[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()); + } + } } 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 78c2dff05b..d33e3056ca 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -219,6 +219,18 @@ 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.object_cache = Arc::new( + self.object_cache + .as_ref() + .clone() + .with_file_io(file_io.clone()), + ); + self.file_io = file_io; + self + } + /// Returns a TableBuilder to build a table pub fn builder() -> TableBuilder { TableBuilder::new() @@ -412,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!( @@ -497,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 3622471277..cc6c6804a7 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -217,7 +217,14 @@ impl Transaction { self.table = refreshed.clone(); } + // Actions read and write manifests, so they need the credentials the + // 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![]; @@ -232,13 +239,40 @@ impl Transaction { )?; } + let moves_location = existing_updates + .iter() + .any(|update| matches!(update, TableUpdate::SetLocation { .. })); + let table_commit = TableCommit::builder() .ident(self.table.identifier().to_owned()) .updates(existing_updates) .requirements(existing_requirements) .build(); - catalog.update_table(table_commit).await + 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. + 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())) } } @@ -250,8 +284,10 @@ 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; + use crate::io::{FileIO, FileIOBuilder, MemoryStorageFactory}; use crate::memory::tests::new_memory_catalog; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, Literal, Struct, TableMetadata, @@ -259,8 +295,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!( @@ -395,6 +432,228 @@ 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()), + ); + } + + 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")); + } + + /// 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(); + 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. + 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![])) + } + } + + /// 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] + 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 { Ok(moved_table()) })); + + 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,