This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3408-b154603a710d730d2f43f8af4db257cc137237ef in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit e9a21181360200c60a4e08538bc385586d300397 Author: NoahKusaba <[email protected]> AuthorDate: Tue Oct 13 09:23:14 2026 +0000 feat(catalog-loader): give CatalogLoader the BoxedCatalogBuilder settings (#3408) #3286 added with_runtime to BoxedCatalogBuilder, but CatalogLoader, the one-line entry point, still had no way to set a runtime, so a catalog loaded through it always fell back to Runtime::current(). It had no way to set a storage or KMS client factory either, so it couldn't load a SQL catalog at all: SqlCatalog requires a storage factory. Add with_storage_factory, with_kms_client_factory and with_runtime to CatalogLoader, and apply them to the builder in load(). Raised in review of #3286. --- crates/catalog/loader/public-api.txt | 3 + crates/catalog/loader/src/lib.rs | 194 +++++++++++++++++++++++++++++------ 2 files changed, 168 insertions(+), 29 deletions(-) diff --git a/crates/catalog/loader/public-api.txt b/crates/catalog/loader/public-api.txt index 4e3c28861..cf4fb8d73 100644 --- a/crates/catalog/loader/public-api.txt +++ b/crates/catalog/loader/public-api.txt @@ -2,6 +2,9 @@ pub mod iceberg_catalog_loader pub struct iceberg_catalog_loader::CatalogLoader<'a> impl iceberg_catalog_loader::CatalogLoader<'_> pub async fn iceberg_catalog_loader::CatalogLoader<'_>::load(self, name: alloc::string::String, props: std::collections::hash::map::HashMap<alloc::string::String, alloc::string::String>) -> iceberg::error::Result<alloc::sync::Arc<dyn iceberg::catalog::Catalog>> +pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc<dyn iceberg::encryption::kms::factory::KmsClientFactory>) -> Self +pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self +pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_storage_factory(self, storage_factory: alloc::sync::Arc<dyn iceberg::io::storage::StorageFactory>) -> Self impl<'a> core::convert::From<&'a str> for iceberg_catalog_loader::CatalogLoader<'a> pub fn iceberg_catalog_loader::CatalogLoader<'a>::from(s: &'a str) -> Self pub trait iceberg_catalog_loader::BoxedCatalogBuilder: core::marker::Send diff --git a/crates/catalog/loader/src/lib.rs b/crates/catalog/loader/src/lib.rs index 3993de676..0ddb29fc8 100644 --- a/crates/catalog/loader/src/lib.rs +++ b/crates/catalog/loader/src/lib.rs @@ -128,22 +128,68 @@ pub fn load(r#type: &str) -> Result<Box<dyn BoxedCatalogBuilder>> { /// Ergonomic catalog loader builder pattern. pub struct CatalogLoader<'a> { catalog_type: &'a str, + storage_factory: Option<Arc<dyn StorageFactory>>, + kms_client_factory: Option<Arc<dyn KmsClientFactory>>, + runtime: Option<Runtime>, } impl<'a> From<&'a str> for CatalogLoader<'a> { fn from(s: &'a str) -> Self { - Self { catalog_type: s } + Self { + catalog_type: s, + storage_factory: None, + kms_client_factory: None, + runtime: None, + } } } impl CatalogLoader<'_> { + /// Sets the storage factory used to build the catalog's `FileIO`; see + /// [`CatalogBuilder::with_storage_factory`]. + pub fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self { + self.storage_factory = Some(storage_factory); + self + } + + /// Sets the KMS client factory used to enable table encryption; see + /// [`CatalogBuilder::with_kms_client_factory`]. + pub fn with_kms_client_factory( + mut self, + kms_client_factory: Arc<dyn KmsClientFactory>, + ) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + + /// Sets the runtime the catalog, and the tables it creates, spawn their + /// tasks on; see [`CatalogBuilder::with_runtime`]. + pub fn with_runtime(mut self, runtime: Runtime) -> Self { + self.runtime = Some(runtime); + self + } + pub async fn load( self, name: String, props: HashMap<String, String>, ) -> Result<Arc<dyn Catalog>> { let builder = load(self.catalog_type)?; - builder.load(name, props).await + self.configure(builder).load(name, props).await + } + + /// Applies the settings given to this loader to `builder`. + fn configure(self, mut builder: Box<dyn BoxedCatalogBuilder>) -> Box<dyn BoxedCatalogBuilder> { + if let Some(storage_factory) = self.storage_factory { + builder = builder.with_storage_factory(storage_factory); + } + if let Some(kms_client_factory) = self.kms_client_factory { + builder = builder.with_kms_client_factory(kms_client_factory); + } + if let Some(runtime) = self.runtime { + builder = builder.with_runtime(runtime); + } + builder } } @@ -152,7 +198,7 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, Mutex}; - use iceberg::encryption::kms::KmsClientFactory; + use iceberg::encryption::kms::{KmsClientFactory, MemoryKmsClientFactory}; use iceberg::io::{LocalFsStorageFactory, StorageFactory}; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder}; use iceberg::{CatalogBuilder, Result, Runtime}; @@ -302,25 +348,59 @@ mod tests { assert!(catalog.is_ok()); } - /// A memory catalog builder that records the runtime it is given. + #[tokio::test] + async fn test_catalog_loader_with_storage_factory_loads_sql_catalog() { + use iceberg_catalog_sql::{SQL_CATALOG_PROP_URI, SQL_CATALOG_PROP_WAREHOUSE}; + + let uri = format!("sqlite:{}", temp_path()); + sqlx::Sqlite::create_database(&uri).await.unwrap(); + + // A SQL catalog can't be loaded without a storage factory. + CatalogLoader::from("sql") + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "sql".to_string(), + HashMap::from([ + (SQL_CATALOG_PROP_URI.to_string(), uri), + ( + SQL_CATALOG_PROP_WAREHOUSE.to_string(), + "s3://warehouse".to_string(), + ), + ]), + ) + .await + .unwrap(); + } + + /// The settings a [`RecordingBuilder`] was given. + #[derive(Debug, Default)] + struct Recorded { + storage_factory: Option<Arc<dyn StorageFactory>>, + kms_client_factory: Option<Arc<dyn KmsClientFactory>>, + runtime: Option<Runtime>, + } + + /// A memory catalog builder that records the settings it is given. #[derive(Debug, Default)] - struct RuntimeRecordingBuilder { - runtime: Arc<Mutex<Option<Runtime>>>, + struct RecordingBuilder { + recorded: Arc<Mutex<Recorded>>, } - impl CatalogBuilder for RuntimeRecordingBuilder { + impl CatalogBuilder for RecordingBuilder { type C = MemoryCatalog; - fn with_storage_factory(self, _storage_factory: Arc<dyn StorageFactory>) -> Self { + fn with_storage_factory(self, storage_factory: Arc<dyn StorageFactory>) -> Self { + self.recorded.lock().unwrap().storage_factory = Some(storage_factory); self } - fn with_kms_client_factory(self, _kms_client_factory: Arc<dyn KmsClientFactory>) -> Self { + fn with_kms_client_factory(self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self { + self.recorded.lock().unwrap().kms_client_factory = Some(kms_client_factory); self } fn with_runtime(self, runtime: Runtime) -> Self { - *self.runtime.lock().unwrap() = Some(runtime); + self.recorded.lock().unwrap().runtime = Some(runtime); self } @@ -333,18 +413,45 @@ mod tests { } } - #[test] - fn test_with_runtime_reaches_the_catalog_builder() { - let tokio_runtime = tokio::runtime::Builder::new_multi_thread() + /// A boxed [`RecordingBuilder`], and the settings it records. + fn recording_builder() -> (Box<dyn BoxedCatalogBuilder>, Arc<Mutex<Recorded>>) { + let recorded = Arc::new(Mutex::new(Recorded::default())); + let builder = Box::new(RecordingBuilder { + recorded: recorded.clone(), + }); + (builder, recorded) + } + + /// A tokio runtime whose worker thread is named `loader-test-runtime`. + fn named_tokio_runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .thread_name("loader-test-runtime") .enable_all() .build() + .unwrap() + } + + /// Asserts that tasks spawned on `runtime` run on a thread of + /// [`named_tokio_runtime`]. + async fn assert_spawns_on_named_runtime(runtime: &Runtime) { + let thread = runtime + .io() + .spawn(async { std::thread::current().name().map(str::to_string) }) + .await .unwrap(); - let recorded = Arc::new(Mutex::new(None)); - let builder: Box<dyn BoxedCatalogBuilder> = Box::new(RuntimeRecordingBuilder { - runtime: recorded.clone(), - }); + assert!( + thread + .as_deref() + .is_some_and(|n| n.starts_with("loader-test-runtime")), + "got: {thread:?}" + ); + } + + #[test] + fn test_with_runtime_reaches_the_catalog_builder() { + let tokio_runtime = named_tokio_runtime(); + let (builder, recorded) = recording_builder(); tokio_runtime.block_on(async { builder @@ -358,21 +465,50 @@ mod tests { // The builder got the runtime passed to the boxed builder: its // tasks run on that runtime's threads. - let runtime = recorded.lock().unwrap().take().unwrap(); - let thread = runtime - .io() - .spawn(async { std::thread::current().name().map(str::to_string) }) - .await - .unwrap(); - assert!( - thread - .as_deref() - .is_some_and(|n| n.starts_with("loader-test-runtime")), - "got: {thread:?}" - ); + let runtime = recorded.lock().unwrap().runtime.take().unwrap(); + assert_spawns_on_named_runtime(&runtime).await; }); } + #[test] + fn test_catalog_loader_settings_reach_the_catalog_builder() { + let tokio_runtime = named_tokio_runtime(); + let storage_factory: Arc<dyn StorageFactory> = Arc::new(LocalFsStorageFactory); + let kms_client_factory: Arc<dyn KmsClientFactory> = + Arc::new(MemoryKmsClientFactory::default()); + let (builder, recorded) = recording_builder(); + + CatalogLoader::from("memory") + .with_storage_factory(storage_factory.clone()) + .with_kms_client_factory(kms_client_factory.clone()) + .with_runtime(Runtime::new(&tokio_runtime)) + .configure(builder); + + let recorded = std::mem::take(&mut *recorded.lock().unwrap()); + assert!(Arc::ptr_eq( + &recorded.storage_factory.unwrap(), + &storage_factory + )); + assert!(Arc::ptr_eq( + &recorded.kms_client_factory.unwrap(), + &kms_client_factory + )); + tokio_runtime.block_on(assert_spawns_on_named_runtime(&recorded.runtime.unwrap())); + } + + #[test] + fn test_catalog_loader_without_settings_leaves_the_builder_default() { + let (builder, recorded) = recording_builder(); + + CatalogLoader::from("memory").configure(builder); + + assert!(matches!(*recorded.lock().unwrap(), Recorded { + storage_factory: None, + kms_client_factory: None, + runtime: None, + })); + } + #[tokio::test] async fn test_error_message_includes_supported_types() { let err = match load("does-not-exist") {
