Skip to content

Commit 5c00d46

Browse files
authored
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.
1 parent e6b2ed4 commit 5c00d46

2 files changed

Lines changed: 168 additions & 29 deletions

File tree

‎crates/catalog/loader/public-api.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ pub mod iceberg_catalog_loader
22
pub struct iceberg_catalog_loader::CatalogLoader<'a>
33
impl iceberg_catalog_loader::CatalogLoader<'_>
44
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>>
5+
pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc<dyn iceberg::encryption::kms::factory::KmsClientFactory>) -> Self
6+
pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self
7+
pub fn iceberg_catalog_loader::CatalogLoader<'_>::with_storage_factory(self, storage_factory: alloc::sync::Arc<dyn iceberg::io::storage::StorageFactory>) -> Self
58
impl<'a> core::convert::From<&'a str> for iceberg_catalog_loader::CatalogLoader<'a>
69
pub fn iceberg_catalog_loader::CatalogLoader<'a>::from(s: &'a str) -> Self
710
pub trait iceberg_catalog_loader::BoxedCatalogBuilder: core::marker::Send

‎crates/catalog/loader/src/lib.rs‎

Lines changed: 165 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -128,22 +128,68 @@ pub fn load(r#type: &str) -> Result<Box<dyn BoxedCatalogBuilder>> {
128128
/// Ergonomic catalog loader builder pattern.
129129
pub struct CatalogLoader<'a> {
130130
catalog_type: &'a str,
131+
storage_factory: Option<Arc<dyn StorageFactory>>,
132+
kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
133+
runtime: Option<Runtime>,
131134
}
132135

133136
impl<'a> From<&'a str> for CatalogLoader<'a> {
134137
fn from(s: &'a str) -> Self {
135-
Self { catalog_type: s }
138+
Self {
139+
catalog_type: s,
140+
storage_factory: None,
141+
kms_client_factory: None,
142+
runtime: None,
143+
}
136144
}
137145
}
138146

139147
impl CatalogLoader<'_> {
148+
/// Sets the storage factory used to build the catalog's `FileIO`; see
149+
/// [`CatalogBuilder::with_storage_factory`].
150+
pub fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
151+
self.storage_factory = Some(storage_factory);
152+
self
153+
}
154+
155+
/// Sets the KMS client factory used to enable table encryption; see
156+
/// [`CatalogBuilder::with_kms_client_factory`].
157+
pub fn with_kms_client_factory(
158+
mut self,
159+
kms_client_factory: Arc<dyn KmsClientFactory>,
160+
) -> Self {
161+
self.kms_client_factory = Some(kms_client_factory);
162+
self
163+
}
164+
165+
/// Sets the runtime the catalog, and the tables it creates, spawn their
166+
/// tasks on; see [`CatalogBuilder::with_runtime`].
167+
pub fn with_runtime(mut self, runtime: Runtime) -> Self {
168+
self.runtime = Some(runtime);
169+
self
170+
}
171+
140172
pub async fn load(
141173
self,
142174
name: String,
143175
props: HashMap<String, String>,
144176
) -> Result<Arc<dyn Catalog>> {
145177
let builder = load(self.catalog_type)?;
146-
builder.load(name, props).await
178+
self.configure(builder).load(name, props).await
179+
}
180+
181+
/// Applies the settings given to this loader to `builder`.
182+
fn configure(self, mut builder: Box<dyn BoxedCatalogBuilder>) -> Box<dyn BoxedCatalogBuilder> {
183+
if let Some(storage_factory) = self.storage_factory {
184+
builder = builder.with_storage_factory(storage_factory);
185+
}
186+
if let Some(kms_client_factory) = self.kms_client_factory {
187+
builder = builder.with_kms_client_factory(kms_client_factory);
188+
}
189+
if let Some(runtime) = self.runtime {
190+
builder = builder.with_runtime(runtime);
191+
}
192+
builder
147193
}
148194
}
149195

@@ -152,7 +198,7 @@ mod tests {
152198
use std::collections::HashMap;
153199
use std::sync::{Arc, Mutex};
154200

155-
use iceberg::encryption::kms::KmsClientFactory;
201+
use iceberg::encryption::kms::{KmsClientFactory, MemoryKmsClientFactory};
156202
use iceberg::io::{LocalFsStorageFactory, StorageFactory};
157203
use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder};
158204
use iceberg::{CatalogBuilder, Result, Runtime};
@@ -302,25 +348,59 @@ mod tests {
302348
assert!(catalog.is_ok());
303349
}
304350

305-
/// A memory catalog builder that records the runtime it is given.
351+
#[tokio::test]
352+
async fn test_catalog_loader_with_storage_factory_loads_sql_catalog() {
353+
use iceberg_catalog_sql::{SQL_CATALOG_PROP_URI, SQL_CATALOG_PROP_WAREHOUSE};
354+
355+
let uri = format!("sqlite:{}", temp_path());
356+
sqlx::Sqlite::create_database(&uri).await.unwrap();
357+
358+
// A SQL catalog can't be loaded without a storage factory.
359+
CatalogLoader::from("sql")
360+
.with_storage_factory(Arc::new(LocalFsStorageFactory))
361+
.load(
362+
"sql".to_string(),
363+
HashMap::from([
364+
(SQL_CATALOG_PROP_URI.to_string(), uri),
365+
(
366+
SQL_CATALOG_PROP_WAREHOUSE.to_string(),
367+
"s3://warehouse".to_string(),
368+
),
369+
]),
370+
)
371+
.await
372+
.unwrap();
373+
}
374+
375+
/// The settings a [`RecordingBuilder`] was given.
376+
#[derive(Debug, Default)]
377+
struct Recorded {
378+
storage_factory: Option<Arc<dyn StorageFactory>>,
379+
kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
380+
runtime: Option<Runtime>,
381+
}
382+
383+
/// A memory catalog builder that records the settings it is given.
306384
#[derive(Debug, Default)]
307-
struct RuntimeRecordingBuilder {
308-
runtime: Arc<Mutex<Option<Runtime>>>,
385+
struct RecordingBuilder {
386+
recorded: Arc<Mutex<Recorded>>,
309387
}
310388

311-
impl CatalogBuilder for RuntimeRecordingBuilder {
389+
impl CatalogBuilder for RecordingBuilder {
312390
type C = MemoryCatalog;
313391

314-
fn with_storage_factory(self, _storage_factory: Arc<dyn StorageFactory>) -> Self {
392+
fn with_storage_factory(self, storage_factory: Arc<dyn StorageFactory>) -> Self {
393+
self.recorded.lock().unwrap().storage_factory = Some(storage_factory);
315394
self
316395
}
317396

318-
fn with_kms_client_factory(self, _kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
397+
fn with_kms_client_factory(self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
398+
self.recorded.lock().unwrap().kms_client_factory = Some(kms_client_factory);
319399
self
320400
}
321401

322402
fn with_runtime(self, runtime: Runtime) -> Self {
323-
*self.runtime.lock().unwrap() = Some(runtime);
403+
self.recorded.lock().unwrap().runtime = Some(runtime);
324404
self
325405
}
326406

@@ -333,18 +413,45 @@ mod tests {
333413
}
334414
}
335415

336-
#[test]
337-
fn test_with_runtime_reaches_the_catalog_builder() {
338-
let tokio_runtime = tokio::runtime::Builder::new_multi_thread()
416+
/// A boxed [`RecordingBuilder`], and the settings it records.
417+
fn recording_builder() -> (Box<dyn BoxedCatalogBuilder>, Arc<Mutex<Recorded>>) {
418+
let recorded = Arc::new(Mutex::new(Recorded::default()));
419+
let builder = Box::new(RecordingBuilder {
420+
recorded: recorded.clone(),
421+
});
422+
(builder, recorded)
423+
}
424+
425+
/// A tokio runtime whose worker thread is named `loader-test-runtime`.
426+
fn named_tokio_runtime() -> tokio::runtime::Runtime {
427+
tokio::runtime::Builder::new_multi_thread()
339428
.worker_threads(1)
340429
.thread_name("loader-test-runtime")
341430
.enable_all()
342431
.build()
432+
.unwrap()
433+
}
434+
435+
/// Asserts that tasks spawned on `runtime` run on a thread of
436+
/// [`named_tokio_runtime`].
437+
async fn assert_spawns_on_named_runtime(runtime: &Runtime) {
438+
let thread = runtime
439+
.io()
440+
.spawn(async { std::thread::current().name().map(str::to_string) })
441+
.await
343442
.unwrap();
344-
let recorded = Arc::new(Mutex::new(None));
345-
let builder: Box<dyn BoxedCatalogBuilder> = Box::new(RuntimeRecordingBuilder {
346-
runtime: recorded.clone(),
347-
});
443+
assert!(
444+
thread
445+
.as_deref()
446+
.is_some_and(|n| n.starts_with("loader-test-runtime")),
447+
"got: {thread:?}"
448+
);
449+
}
450+
451+
#[test]
452+
fn test_with_runtime_reaches_the_catalog_builder() {
453+
let tokio_runtime = named_tokio_runtime();
454+
let (builder, recorded) = recording_builder();
348455

349456
tokio_runtime.block_on(async {
350457
builder
@@ -358,21 +465,50 @@ mod tests {
358465

359466
// The builder got the runtime passed to the boxed builder: its
360467
// tasks run on that runtime's threads.
361-
let runtime = recorded.lock().unwrap().take().unwrap();
362-
let thread = runtime
363-
.io()
364-
.spawn(async { std::thread::current().name().map(str::to_string) })
365-
.await
366-
.unwrap();
367-
assert!(
368-
thread
369-
.as_deref()
370-
.is_some_and(|n| n.starts_with("loader-test-runtime")),
371-
"got: {thread:?}"
372-
);
468+
let runtime = recorded.lock().unwrap().runtime.take().unwrap();
469+
assert_spawns_on_named_runtime(&runtime).await;
373470
});
374471
}
375472

473+
#[test]
474+
fn test_catalog_loader_settings_reach_the_catalog_builder() {
475+
let tokio_runtime = named_tokio_runtime();
476+
let storage_factory: Arc<dyn StorageFactory> = Arc::new(LocalFsStorageFactory);
477+
let kms_client_factory: Arc<dyn KmsClientFactory> =
478+
Arc::new(MemoryKmsClientFactory::default());
479+
let (builder, recorded) = recording_builder();
480+
481+
CatalogLoader::from("memory")
482+
.with_storage_factory(storage_factory.clone())
483+
.with_kms_client_factory(kms_client_factory.clone())
484+
.with_runtime(Runtime::new(&tokio_runtime))
485+
.configure(builder);
486+
487+
let recorded = std::mem::take(&mut *recorded.lock().unwrap());
488+
assert!(Arc::ptr_eq(
489+
&recorded.storage_factory.unwrap(),
490+
&storage_factory
491+
));
492+
assert!(Arc::ptr_eq(
493+
&recorded.kms_client_factory.unwrap(),
494+
&kms_client_factory
495+
));
496+
tokio_runtime.block_on(assert_spawns_on_named_runtime(&recorded.runtime.unwrap()));
497+
}
498+
499+
#[test]
500+
fn test_catalog_loader_without_settings_leaves_the_builder_default() {
501+
let (builder, recorded) = recording_builder();
502+
503+
CatalogLoader::from("memory").configure(builder);
504+
505+
assert!(matches!(*recorded.lock().unwrap(), Recorded {
506+
storage_factory: None,
507+
kms_client_factory: None,
508+
runtime: None,
509+
}));
510+
}
511+
376512
#[tokio::test]
377513
async fn test_error_message_includes_supported_types() {
378514
let err = match load("does-not-exist") {

0 commit comments

Comments
 (0)