diff --git a/crates/catalog/glue/public-api.txt b/crates/catalog/glue/public-api.txt index 8b7a38da50..902daf635f 100644 --- a/crates/catalog/glue/public-api.txt +++ b/crates/catalog/glue/public-api.txt @@ -29,6 +29,7 @@ pub fn iceberg_catalog_glue::GlueCatalogBuilder::fmt(&self, f: &mut core::fmt::F impl iceberg::catalog::CatalogBuilder for iceberg_catalog_glue::GlueCatalogBuilder pub type iceberg_catalog_glue::GlueCatalogBuilder::C = iceberg_catalog_glue::GlueCatalog pub fn iceberg_catalog_glue::GlueCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg_catalog_glue::GlueCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_glue::GlueCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_glue::GlueCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub const iceberg_catalog_glue::AWS_ACCESS_KEY_ID: &str diff --git a/crates/catalog/glue/src/catalog.rs b/crates/catalog/glue/src/catalog.rs index 0a69797ef9..b2c93fdfbe 100644 --- a/crates/catalog/glue/src/catalog.rs +++ b/crates/catalog/glue/src/catalog.rs @@ -25,6 +25,7 @@ use async_trait::async_trait; use aws_sdk_glue::operation::create_table::CreateTableError; use aws_sdk_glue::operation::update_table::UpdateTableError; use aws_sdk_glue::types::TableInput; +use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; use iceberg::io::{ FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_REGION, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, StorageFactory, @@ -58,6 +59,7 @@ pub const GLUE_CATALOG_PROP_WAREHOUSE: &str = "warehouse"; pub struct GlueCatalogBuilder { config: GlueCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -72,6 +74,7 @@ impl Default for GlueCatalogBuilder { props: HashMap::new(), }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -85,6 +88,11 @@ impl CatalogBuilder for GlueCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -140,7 +148,11 @@ impl CatalogBuilder for GlueCatalogBuilder { Some(rt) => rt, None => Runtime::try_current()?, }; - GlueCatalog::new(self.config, self.storage_factory, runtime).await + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + GlueCatalog::new(self.config, self.storage_factory, runtime, kms_client).await } } } @@ -163,6 +175,7 @@ pub struct GlueCatalog { client: GlueClient, file_io: FileIO, runtime: Runtime, + kms_client: Option>, } impl Debug for GlueCatalog { @@ -179,6 +192,7 @@ impl GlueCatalog { config: GlueCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Result { let sdk_config = create_sdk_config(&config.props, config.uri.as_ref()).await; let mut file_io_props = config.props.clone(); @@ -228,6 +242,7 @@ impl GlueCatalog { client: GlueClient(client), file_io, runtime, + kms_client, }) } /// Get the catalogs `FileIO` @@ -278,7 +293,7 @@ impl GlueCatalog { let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?; - let table = Table::builder() + let mut builder = Table::builder() .file_io(self.file_io()) .metadata_location(metadata_location) .metadata(metadata) @@ -286,8 +301,11 @@ impl GlueCatalog { NamespaceIdent::new(db_name), table_name.to_owned(), )) - .runtime(self.runtime.clone()) - .build()?; + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + let table = builder.build()?; Ok((table, version_id)) } @@ -638,13 +656,16 @@ impl Catalog for GlueCatalog { builder.send().await.map_err(from_aws_sdk_error)?; - Table::builder() + let mut builder = Table::builder() .file_io(self.file_io()) .metadata_location(metadata_location_str) .metadata(metadata) .identifier(TableIdent::new(NamespaceIdent::new(db_name), table_name)) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } /// Loads a table from the Glue Catalog and constructs a `Table` object @@ -868,13 +889,16 @@ impl Catalog for GlueCatalog { .with_source(anyhow!("aws sdk error: {error:?}")) })?; - Ok(Table::builder() + let mut builder = Table::builder() .identifier(table_ident.clone()) .metadata_location(metadata_location) .metadata(metadata) .file_io(self.file_io()) - .runtime(self.runtime.clone()) - .build()?) + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + Ok(builder.build()?) } async fn update_table(&self, commit: TableCommit) -> Result { diff --git a/crates/catalog/hms/public-api.txt b/crates/catalog/hms/public-api.txt index 744819c437..9ea5f553b3 100644 --- a/crates/catalog/hms/public-api.txt +++ b/crates/catalog/hms/public-api.txt @@ -35,6 +35,7 @@ pub fn iceberg_catalog_hms::HmsCatalogBuilder::fmt(&self, f: &mut core::fmt::For impl iceberg::catalog::CatalogBuilder for iceberg_catalog_hms::HmsCatalogBuilder pub type iceberg_catalog_hms::HmsCatalogBuilder::C = iceberg_catalog_hms::HmsCatalog pub fn iceberg_catalog_hms::HmsCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg_catalog_hms::HmsCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_hms::HmsCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_hms::HmsCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub const iceberg_catalog_hms::HMS_CATALOG_PROP_THRIFT_TRANSPORT: &str diff --git a/crates/catalog/hms/src/catalog.rs b/crates/catalog/hms/src/catalog.rs index 0f15890c77..d54b9acf65 100644 --- a/crates/catalog/hms/src/catalog.rs +++ b/crates/catalog/hms/src/catalog.rs @@ -26,6 +26,7 @@ use hive_metastore::{ ThriftHiveMetastoreClient, ThriftHiveMetastoreClientBuilder, ThriftHiveMetastoreGetDatabaseException, ThriftHiveMetastoreGetTableException, }; +use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; use iceberg::io::{FileIO, FileIOBuilder, StorageFactory}; use iceberg::spec::{TableMetadata, TableMetadataBuilder}; use iceberg::table::Table; @@ -56,6 +57,7 @@ pub const HMS_CATALOG_PROP_WAREHOUSE: &str = "warehouse"; pub struct HmsCatalogBuilder { config: HmsCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -70,6 +72,7 @@ impl Default for HmsCatalogBuilder { props: HashMap::new(), }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -83,6 +86,11 @@ impl CatalogBuilder for HmsCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -123,7 +131,12 @@ impl CatalogBuilder for HmsCatalogBuilder { }) .collect(); - let result = (|| -> Result { + async move { + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + if self.config.name.is_none() { return Err(Error::new( ErrorKind::DataInvalid, @@ -146,10 +159,8 @@ impl CatalogBuilder for HmsCatalogBuilder { Some(rt) => rt, None => Runtime::try_current()?, }; - HmsCatalog::new(self.config, self.storage_factory, runtime) - })(); - - std::future::ready(result) + HmsCatalog::new(self.config, self.storage_factory, runtime, kms_client) + } } } @@ -182,6 +193,7 @@ pub struct HmsCatalog { client: HmsClient, file_io: FileIO, runtime: Runtime, + kms_client: Option>, } impl Debug for HmsCatalog { @@ -198,6 +210,7 @@ impl HmsCatalog { config: HmsCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Result { let address = config .address @@ -238,6 +251,7 @@ impl HmsCatalog { client: HmsClient(client), file_io, runtime, + kms_client, }) } /// Get the catalogs `FileIO` @@ -539,13 +553,16 @@ impl Catalog for HmsCatalog { .await .map_err(from_thrift_error)?; - Table::builder() + let mut builder = Table::builder() .file_io(self.file_io()) .metadata_location(metadata_location_str) .metadata(metadata) .identifier(TableIdent::new(NamespaceIdent::new(db_name), table_name)) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } /// Loads a table from the Hive Metastore and constructs a `Table` object @@ -575,7 +592,7 @@ impl Catalog for HmsCatalog { let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?; - Table::builder() + let mut builder = Table::builder() .file_io(self.file_io()) .metadata_location(metadata_location) .metadata(metadata) @@ -583,8 +600,11 @@ impl Catalog for HmsCatalog { NamespaceIdent::new(db_name), table.name.clone(), )) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } /// Asynchronously drops a table from the database. diff --git a/crates/catalog/loader/public-api.txt b/crates/catalog/loader/public-api.txt index c9e50207f4..1efaf831df 100644 --- a/crates/catalog/loader/public-api.txt +++ b/crates/catalog/loader/public-api.txt @@ -6,9 +6,11 @@ impl<'a> core::convert::From<&'a str> for iceberg_catalog_loader::CatalogLoader< pub fn iceberg_catalog_loader::CatalogLoader<'a>::from(s: &'a str) -> Self pub trait iceberg_catalog_loader::BoxedCatalogBuilder: core::marker::Send pub fn iceberg_catalog_loader::BoxedCatalogBuilder::load<'async_trait>(self: alloc::boxed::Box, name: alloc::string::String, props: std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait +pub fn iceberg_catalog_loader::BoxedCatalogBuilder::with_kms_client_factory(self: alloc::boxed::Box, kms_client_factory: alloc::sync::Arc) -> alloc::boxed::Box pub fn iceberg_catalog_loader::BoxedCatalogBuilder::with_storage_factory(self: alloc::boxed::Box, storage_factory: alloc::sync::Arc) -> alloc::boxed::Box impl iceberg_catalog_loader::BoxedCatalogBuilder for T pub fn T::load<'async_trait>(self: alloc::boxed::Box, name: alloc::string::String, props: std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait +pub fn T::with_kms_client_factory(self: alloc::boxed::Box, kms_client_factory: alloc::sync::Arc) -> alloc::boxed::Box pub fn T::with_storage_factory(self: alloc::boxed::Box, storage_factory: alloc::sync::Arc) -> alloc::boxed::Box pub fn iceberg_catalog_loader::load(type: &str) -> iceberg::error::Result> pub fn iceberg_catalog_loader::supported_types() -> alloc::vec::Vec<&'static str> diff --git a/crates/catalog/loader/src/lib.rs b/crates/catalog/loader/src/lib.rs index f681337736..d45decc40e 100644 --- a/crates/catalog/loader/src/lib.rs +++ b/crates/catalog/loader/src/lib.rs @@ -19,6 +19,7 @@ use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; +use iceberg::encryption::kms::KmsClientFactory; use iceberg::io::StorageFactory; use iceberg::{Catalog, CatalogBuilder, Error, ErrorKind, Result}; use iceberg_catalog_glue::GlueCatalogBuilder; @@ -51,6 +52,11 @@ pub trait BoxedCatalogBuilder: Send { storage_factory: Arc, ) -> Box; + fn with_kms_client_factory( + self: Box, + kms_client_factory: Arc, + ) -> Box; + async fn load( self: Box, name: String, @@ -67,6 +73,16 @@ impl BoxedCatalogBuilder for T { Box::new(CatalogBuilder::with_storage_factory(*self, storage_factory)) } + fn with_kms_client_factory( + self: Box, + kms_client_factory: Arc, + ) -> Box { + Box::new(CatalogBuilder::with_kms_client_factory( + *self, + kms_client_factory, + )) + } + async fn load( self: Box, name: String, diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 027df29b24..db7f906017 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -215,6 +215,7 @@ pub fn iceberg_catalog_rest::RestCatalogBuilder::fmt(&self, f: &mut core::fmt::F impl iceberg::catalog::CatalogBuilder for iceberg_catalog_rest::RestCatalogBuilder pub type iceberg_catalog_rest::RestCatalogBuilder::C = iceberg_catalog_rest::RestCatalog pub fn iceberg_catalog_rest::RestCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg_catalog_rest::RestCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub struct iceberg_catalog_rest::StorageCredential diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 0626ce5061..da05fc7e3b 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -23,6 +23,7 @@ use std::str::FromStr; use std::sync::Arc; use async_trait::async_trait; +use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; use iceberg::io::{FileIO, FileIOBuilder, StorageFactory}; use iceberg::table::Table; use iceberg::{ @@ -62,6 +63,7 @@ const PATH_V1: &str = "v1"; pub struct RestCatalogBuilder { config: RestCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -76,6 +78,7 @@ impl Default for RestCatalogBuilder { client: None, }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -89,6 +92,11 @@ impl CatalogBuilder for RestCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -118,7 +126,7 @@ impl CatalogBuilder for RestCatalogBuilder { .filter(|(k, _)| k != REST_CATALOG_PROP_URI && k != REST_CATALOG_PROP_WAREHOUSE) .collect(); - let result = { + async move { if self.config.name.is_none() { Err(Error::new( ErrorKind::DataInvalid, @@ -131,11 +139,18 @@ impl CatalogBuilder for RestCatalogBuilder { )) } else { let runtime = self.runtime.unwrap_or_else(Runtime::current); - Ok(RestCatalog::new(self.config, self.storage_factory, runtime)) + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + Ok(RestCatalog::new( + self.config, + self.storage_factory, + runtime, + kms_client, + )) } - }; - - std::future::ready(result) + } } } @@ -360,6 +375,8 @@ pub struct RestCatalog { /// Storage factory for creating FileIO instances. storage_factory: Option>, runtime: Runtime, + /// Optional KMS client for encrypted tables. + kms_client: Option>, } impl RestCatalog { @@ -368,12 +385,14 @@ impl RestCatalog { config: RestCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Self { Self { user_config: config, ctx: OnceCell::new(), storage_factory, runtime, + kms_client, } } @@ -801,11 +820,14 @@ impl Catalog for RestCatalog { .load_file_io(Some(metadata_location), Some(config)) .await?; - let table_builder = Table::builder() + let mut table_builder = Table::builder() .identifier(table_ident.clone()) .file_io(file_io) .metadata(response.metadata) .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + table_builder = table_builder.kms_client(kms_client); + } if let Some(metadata_location) = response.metadata_location { table_builder.metadata_location(metadata_location).build() @@ -858,11 +880,14 @@ impl Catalog for RestCatalog { .load_file_io(response.metadata_location.as_deref(), Some(config)) .await?; - let table_builder = Table::builder() + let mut table_builder = Table::builder() .identifier(table_ident.clone()) .file_io(file_io) .metadata(response.metadata) .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + table_builder = table_builder.kms_client(kms_client); + } if let Some(metadata_location) = response.metadata_location { table_builder.metadata_location(metadata_location).build() @@ -993,13 +1018,16 @@ impl Catalog for RestCatalog { let file_io = self.load_file_io(Some(metadata_location), None).await?; - Table::builder() + let mut table_builder = Table::builder() .identifier(table_ident.clone()) .file_io(file_io) .metadata(response.metadata) .metadata_location(metadata_location.clone()) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + table_builder = table_builder.kms_client(kms_client); + } + table_builder.build() } async fn update_table(&self, mut commit: TableCommit) -> Result
{ @@ -1066,13 +1094,16 @@ impl Catalog for RestCatalog { .load_file_io(Some(&response.metadata_location), None) .await?; - Table::builder() + let mut table_builder = Table::builder() .identifier(commit.identifier().clone()) .file_io(file_io) .metadata(response.metadata) .metadata_location(response.metadata_location) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + table_builder = table_builder.kms_client(kms_client); + } + table_builder.build() } } @@ -1119,6 +1150,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); assert_eq!( @@ -1194,6 +1226,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1242,6 +1275,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1267,6 +1301,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1299,6 +1334,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1331,6 +1367,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1363,6 +1400,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1477,6 +1515,7 @@ mod tests { .build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1525,6 +1564,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let _namespaces = catalog.list_namespaces(None).await.unwrap(); @@ -1556,6 +1596,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let namespaces = catalog.list_namespaces(None).await.unwrap(); @@ -1608,6 +1649,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let namespaces = catalog.list_namespaces(None).await.unwrap(); @@ -1708,6 +1750,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let namespaces = catalog.list_namespaces(None).await.unwrap(); @@ -1762,6 +1805,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let namespaces = catalog @@ -1806,6 +1850,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let namespaces = catalog @@ -1840,6 +1885,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); assert!( @@ -1869,6 +1915,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); catalog @@ -1910,6 +1957,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let tables = catalog @@ -1979,6 +2027,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let tables = catalog @@ -2111,6 +2160,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let tables = catalog @@ -2156,6 +2206,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); catalog @@ -2186,6 +2237,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); assert!( @@ -2218,6 +2270,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); catalog @@ -2253,6 +2306,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table = catalog @@ -2371,6 +2425,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table = catalog @@ -2408,6 +2463,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table_creation = TableCreation::builder() @@ -2558,6 +2614,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table_creation = TableCreation::builder() @@ -2628,6 +2685,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table1 = { @@ -2773,6 +2831,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table1 = { @@ -2839,6 +2898,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table_ident = TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()); @@ -2891,6 +2951,7 @@ mod tests { RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), + None, ); let table_ident = diff --git a/crates/catalog/s3tables/public-api.txt b/crates/catalog/s3tables/public-api.txt index b6b704e3dd..6d1b089e12 100644 --- a/crates/catalog/s3tables/public-api.txt +++ b/crates/catalog/s3tables/public-api.txt @@ -30,6 +30,7 @@ pub fn iceberg_catalog_s3tables::S3TablesCatalogBuilder::fmt(&self, f: &mut core impl iceberg::catalog::CatalogBuilder for iceberg_catalog_s3tables::S3TablesCatalogBuilder pub type iceberg_catalog_s3tables::S3TablesCatalogBuilder::C = iceberg_catalog_s3tables::S3TablesCatalog pub fn iceberg_catalog_s3tables::S3TablesCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg_catalog_s3tables::S3TablesCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_s3tables::S3TablesCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_s3tables::S3TablesCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub const iceberg_catalog_s3tables::S3TABLES_CATALOG_PROP_ENDPOINT_URL: &str diff --git a/crates/catalog/s3tables/src/catalog.rs b/crates/catalog/s3tables/src/catalog.rs index e7f4174261..3bd75b4833 100644 --- a/crates/catalog/s3tables/src/catalog.rs +++ b/crates/catalog/s3tables/src/catalog.rs @@ -27,6 +27,7 @@ use aws_sdk_s3tables::operation::get_table::{GetTableError, GetTableOutput}; use aws_sdk_s3tables::operation::list_tables::ListTablesOutput; use aws_sdk_s3tables::operation::update_table_metadata_location::UpdateTableMetadataLocationError; use aws_sdk_s3tables::types::OpenTableFormat; +use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; use iceberg::io::{FileIO, FileIOBuilder, StorageFactory}; use iceberg::spec::{TableMetadata, TableMetadataBuilder}; use iceberg::table::Table; @@ -71,6 +72,7 @@ struct S3TablesCatalogConfig { pub struct S3TablesCatalogBuilder { config: S3TablesCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -86,6 +88,7 @@ impl Default for S3TablesCatalogBuilder { props: HashMap::new(), }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -134,6 +137,11 @@ impl CatalogBuilder for S3TablesCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -183,7 +191,11 @@ impl CatalogBuilder for S3TablesCatalogBuilder { Some(rt) => rt, None => Runtime::try_current()?, }; - S3TablesCatalog::new(self.config, self.storage_factory, runtime).await + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + S3TablesCatalog::new(self.config, self.storage_factory, runtime, kms_client).await } } } @@ -196,6 +208,7 @@ pub struct S3TablesCatalog { s3tables_client: aws_sdk_s3tables::Client, file_io: FileIO, runtime: Runtime, + kms_client: Option>, } impl S3TablesCatalog { @@ -204,6 +217,7 @@ impl S3TablesCatalog { config: S3TablesCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Result { let s3tables_client = if let Some(client) = config.client.clone() { client @@ -227,6 +241,7 @@ impl S3TablesCatalog { s3tables_client, file_io, runtime, + kms_client, }) } @@ -270,13 +285,16 @@ impl S3TablesCatalog { })?; let metadata = TableMetadata::read_from(&self.file_io, metadata_location).await?; - let table = Table::builder() + let mut builder = Table::builder() .identifier(table_ident.clone()) .metadata(metadata) .metadata_location(metadata_location) .file_io(self.file_io.clone()) - .runtime(self.runtime.clone()) - .build()?; + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + let table = builder.build()?; Ok((table, resp.version_token)) } } @@ -569,13 +587,16 @@ impl Catalog for S3TablesCatalog { .await .map_err(from_aws_sdk_error)?; - let table = Table::builder() + let mut builder = Table::builder() .identifier(table_ident) .metadata_location(metadata_location_str) .metadata(metadata) .file_io(self.file_io.clone()) - .runtime(self.runtime.clone()) - .build()?; + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + let table = builder.build()?; Ok(table) } @@ -759,7 +780,7 @@ mod tests { }; Ok(Some( - S3TablesCatalog::new(config, None, Runtime::current()).await?, + S3TablesCatalog::new(config, None, Runtime::current(), None).await?, )) } diff --git a/crates/catalog/sql/public-api.txt b/crates/catalog/sql/public-api.txt index 9402ccde4e..7f3c638942 100644 --- a/crates/catalog/sql/public-api.txt +++ b/crates/catalog/sql/public-api.txt @@ -48,6 +48,7 @@ pub fn iceberg_catalog_sql::SqlCatalogBuilder::fmt(&self, f: &mut core::fmt::For impl iceberg::catalog::CatalogBuilder for iceberg_catalog_sql::SqlCatalogBuilder pub type iceberg_catalog_sql::SqlCatalogBuilder::C = iceberg_catalog_sql::SqlCatalog pub fn iceberg_catalog_sql::SqlCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg_catalog_sql::SqlCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_sql::SqlCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_sql::SqlCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub const iceberg_catalog_sql::SQL_CATALOG_PROP_BIND_STYLE: &str diff --git a/crates/catalog/sql/src/catalog.rs b/crates/catalog/sql/src/catalog.rs index 02e32c5f4a..605d2b6219 100644 --- a/crates/catalog/sql/src/catalog.rs +++ b/crates/catalog/sql/src/catalog.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; +use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory}; use iceberg::io::{FileIO, FileIOBuilder, StorageFactory}; use iceberg::spec::{TableMetadata, TableMetadataBuilder}; use iceberg::table::Table; @@ -67,6 +68,7 @@ static TEST_BEFORE_ACQUIRE: bool = true; // Default the health-check of each con pub struct SqlCatalogBuilder { config: SqlCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -81,6 +83,7 @@ impl Default for SqlCatalogBuilder { props: HashMap::new(), }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -145,6 +148,11 @@ impl CatalogBuilder for SqlCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -201,7 +209,11 @@ impl CatalogBuilder for SqlCatalogBuilder { Some(rt) => rt, None => Runtime::try_current()?, }; - SqlCatalog::new(self.config, self.storage_factory, runtime).await + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + SqlCatalog::new(self.config, self.storage_factory, runtime, kms_client).await } } } @@ -233,6 +245,7 @@ pub struct SqlCatalog { fileio: FileIO, sql_bind_style: SqlBindStyle, runtime: Runtime, + kms_client: Option>, } #[derive(Debug, PartialEq, strum::EnumString, strum::Display)] @@ -250,6 +263,7 @@ impl SqlCatalog { config: SqlCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Result { let factory = storage_factory.ok_or_else(|| { Error::new( @@ -321,6 +335,7 @@ impl SqlCatalog { fileio, sql_bind_style: config.sql_bind_style, runtime, + kms_client, }) } @@ -818,13 +833,16 @@ impl Catalog for SqlCatalog { let metadata = TableMetadata::read_from(&self.fileio, &tbl_metadata_location).await?; - Ok(Table::builder() + let mut builder = Table::builder() .file_io(self.fileio.clone()) .identifier(identifier.clone()) .metadata_location(tbl_metadata_location) .metadata(metadata) - .runtime(self.runtime.clone()) - .build()?) + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + Ok(builder.build()?) } async fn create_table( @@ -889,13 +907,16 @@ impl Catalog for SqlCatalog { VALUES (?, ?, ?, ?, ?) "), vec![Some(&self.name), Some(&namespace.join(".")), Some(&tbl_name.clone()), Some(&tbl_metadata_location_str), Some(CATALOG_FIELD_TABLE_RECORD_TYPE)], None).await?; - Ok(Table::builder() + let mut builder = Table::builder() .file_io(self.fileio.clone()) .metadata_location(tbl_metadata_location_str) .identifier(tbl_ident) .metadata(tbl_metadata) - .runtime(self.runtime.clone()) - .build()?) + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + Ok(builder.build()?) } async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { @@ -961,13 +982,16 @@ impl Catalog for SqlCatalog { VALUES (?, ?, ?, ?, ?) "), vec![Some(&self.name), Some(&namespace.join(".")), Some(&tbl_name), Some(&metadata_location), Some(CATALOG_FIELD_TABLE_RECORD_TYPE)], None).await?; - Ok(Table::builder() + let mut builder = Table::builder() .identifier(table_ident.clone()) .metadata_location(metadata_location) .metadata(metadata) .file_io(self.fileio.clone()) - .runtime(self.runtime.clone()) - .build()?) + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + Ok(builder.build()?) } /// Updates an existing table within the SQL catalog. diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 41d4d7691b..c2013294e3 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -195,6 +195,20 @@ pub fn iceberg::encryption::kms::MemoryKeyManagementClient::generate_key<'life0, pub fn iceberg::encryption::kms::MemoryKeyManagementClient::supports_key_generation(&self) -> bool pub fn iceberg::encryption::kms::MemoryKeyManagementClient::unwrap_key<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, wrapped_key: &'life1 [u8], wrapping_key_id: &'life2 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub fn iceberg::encryption::kms::MemoryKeyManagementClient::wrap_key<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, key: &'life1 [u8], wrapping_key_id: &'life2 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub struct iceberg::encryption::kms::MemoryKmsClientFactory +impl iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::add_master_key(&self, key_id: impl core::convert::Into) -> iceberg::Result<()> +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::add_master_key_bytes(&self, key_id: impl core::convert::Into, key_bytes: iceberg::encryption::SensitiveBytes) -> iceberg::Result<()> +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::new() -> Self +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::with_master_key_size(master_key_size: iceberg::encryption::AesKeySize) -> Self +impl core::clone::Clone for iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::clone(&self) -> iceberg::encryption::kms::MemoryKmsClientFactory +impl core::default::Default for iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::default() -> iceberg::encryption::kms::MemoryKmsClientFactory +impl core::fmt::Debug for iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg::encryption::kms::KmsClientFactory for iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::create_kms_client<'life0, 'life1, 'async_trait>(&'life0 self, _properties: &'life1 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub trait iceberg::encryption::kms::KeyManagementClient: core::marker::Send + core::marker::Sync + core::fmt::Debug pub fn iceberg::encryption::kms::KeyManagementClient::generate_key<'life0, 'life1, 'async_trait>(&'life0 self, wrapping_key_id: &'life1 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg::encryption::kms::KeyManagementClient::supports_key_generation(&self) -> bool @@ -210,6 +224,10 @@ pub fn T::generate_key<'life0, 'life1, 'async_trait>(&'life0 self, wrapping_key_ pub fn T::supports_key_generation(&self) -> bool pub fn T::unwrap_key<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, wrapped_key: &'life1 [u8], wrapping_key_id: &'life2 str) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub fn T::wrap_key<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, key: &'life1 [u8], wrapping_key_id: &'life2 str) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub trait iceberg::encryption::kms::KmsClientFactory: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg::encryption::kms::KmsClientFactory::create_kms_client<'life0, 'life1, 'async_trait>(&'life0 self, properties: &'life1 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +impl iceberg::encryption::kms::KmsClientFactory for iceberg::encryption::kms::MemoryKmsClientFactory +pub fn iceberg::encryption::kms::MemoryKmsClientFactory::create_kms_client<'life0, 'life1, 'async_trait>(&'life0 self, _properties: &'life1 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub enum iceberg::encryption::AesKeySize pub iceberg::encryption::AesKeySize::Bits128 = 128 pub iceberg::encryption::AesKeySize::Bits192 = 192 @@ -1075,6 +1093,7 @@ pub fn iceberg::memory::MemoryCatalogBuilder::fmt(&self, f: &mut core::fmt::Form impl iceberg::CatalogBuilder for iceberg::memory::MemoryCatalogBuilder pub type iceberg::memory::MemoryCatalogBuilder::C = iceberg::memory::MemoryCatalog pub fn iceberg::memory::MemoryCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg::memory::MemoryCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_runtime(self, runtime: iceberg::Runtime) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub const iceberg::memory::MEMORY_CATALOG_WAREHOUSE: &str @@ -3685,11 +3704,13 @@ pub fn iceberg::memory::MemoryCatalog::update_table<'life0, 'async_trait>(&'life pub trait iceberg::CatalogBuilder: core::default::Default + core::fmt::Debug + core::marker::Send + core::marker::Sync pub type iceberg::CatalogBuilder::C: iceberg::Catalog pub fn iceberg::CatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg::CatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg::CatalogBuilder::with_runtime(self, runtime: iceberg::Runtime) -> Self pub fn iceberg::CatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self impl iceberg::CatalogBuilder for iceberg::memory::MemoryCatalogBuilder pub type iceberg::memory::MemoryCatalogBuilder::C = iceberg::memory::MemoryCatalog pub fn iceberg::memory::MemoryCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg::memory::MemoryCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_runtime(self, runtime: iceberg::Runtime) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub async fn iceberg::drop_table_data(table_info: &iceberg::table::Table) -> iceberg::Result<()> diff --git a/crates/iceberg/src/catalog/memory/catalog.rs b/crates/iceberg/src/catalog/memory/catalog.rs index 67f8ab8dd1..e004ae9dfd 100644 --- a/crates/iceberg/src/catalog/memory/catalog.rs +++ b/crates/iceberg/src/catalog/memory/catalog.rs @@ -26,6 +26,7 @@ use futures::lock::{Mutex, MutexGuard}; use itertools::Itertools; use super::namespace_state::NamespaceState; +use crate::encryption::kms::{KeyManagementClient, KmsClientFactory}; use crate::io::{FileIO, FileIOBuilder, MemoryStorageFactory, StorageFactory}; use crate::runtime::Runtime; use crate::spec::{TableMetadata, TableMetadataBuilder}; @@ -46,6 +47,7 @@ const LOCATION: &str = "location"; pub struct MemoryCatalogBuilder { config: MemoryCatalogConfig, storage_factory: Option>, + kms_client_factory: Option>, runtime: Option, } @@ -58,6 +60,7 @@ impl Default for MemoryCatalogBuilder { props: HashMap::new(), }, storage_factory: None, + kms_client_factory: None, runtime: None, } } @@ -71,6 +74,11 @@ impl CatalogBuilder for MemoryCatalogBuilder { self } + fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } + fn with_runtime(mut self, runtime: Runtime) -> Self { self.runtime = Some(runtime); self @@ -96,7 +104,7 @@ impl CatalogBuilder for MemoryCatalogBuilder { .filter(|(k, _)| k != MEMORY_CATALOG_WAREHOUSE) .collect(); - let result = { + async move { if self.config.name.is_none() { Err(Error::new( ErrorKind::DataInvalid, @@ -109,11 +117,13 @@ impl CatalogBuilder for MemoryCatalogBuilder { )) } else { let runtime = self.runtime.unwrap_or_else(Runtime::current); - MemoryCatalog::new(self.config, self.storage_factory, runtime) + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + MemoryCatalog::new(self.config, self.storage_factory, runtime, kms_client) } - }; - - std::future::ready(result) + } } } @@ -131,6 +141,7 @@ pub struct MemoryCatalog { file_io: FileIO, warehouse_location: String, runtime: Runtime, + kms_client: Option>, } impl MemoryCatalog { @@ -139,6 +150,7 @@ impl MemoryCatalog { config: MemoryCatalogConfig, storage_factory: Option>, runtime: Runtime, + kms_client: Option>, ) -> Result { // Use provided factory or default to MemoryStorageFactory let factory = storage_factory.unwrap_or_else(|| Arc::new(MemoryStorageFactory)); @@ -148,6 +160,7 @@ impl MemoryCatalog { file_io: FileIOBuilder::new(factory).with_props(config.props).build(), warehouse_location: config.warehouse, runtime, + kms_client, }) } @@ -160,13 +173,16 @@ impl MemoryCatalog { let metadata_location = root_namespace_state.get_existing_table_location(table_ident)?; let metadata = TableMetadata::read_from(&self.file_io, metadata_location).await?; - Table::builder() + let mut builder = Table::builder() .identifier(table_ident.clone()) .metadata(metadata) .metadata_location(metadata_location.to_string()) .file_io(self.file_io.clone()) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } } @@ -315,13 +331,16 @@ impl Catalog for MemoryCatalog { root_namespace_state.insert_new_table(&table_ident, metadata_location.to_string())?; - Table::builder() + let mut builder = Table::builder() .file_io(self.file_io.clone()) .metadata_location(metadata_location.to_string()) .metadata(metadata) .identifier(table_ident) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } /// Load table from the catalog. @@ -382,13 +401,16 @@ impl Catalog for MemoryCatalog { let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?; - Table::builder() + let mut builder = Table::builder() .file_io(self.file_io.clone()) .metadata_location(metadata_location) .metadata(metadata) .identifier(table_ident.clone()) - .runtime(self.runtime.clone()) - .build() + .runtime(self.runtime.clone()); + if let Some(kms_client) = self.kms_client.clone() { + builder = builder.kms_client(kms_client); + } + builder.build() } /// Update a table in the catalog. @@ -428,7 +450,8 @@ pub(crate) mod tests { use tempfile::TempDir; use super::*; - use crate::io::FileIO; + use crate::encryption::kms::MemoryKmsClientFactory; + use crate::io::{FileIO, LocalFsStorageFactory}; use crate::spec::{NestedField, PartitionSpec, PrimitiveType, Schema, SortOrder, Type}; use crate::test_utils::test_runtime; use crate::transaction::{ApplyTransactionAction, Transaction}; @@ -1934,6 +1957,91 @@ pub(crate) mod tests { assert_eq!(err.kind(), ErrorKind::TableNotFound); } + /// Master key bytes used to generate the encrypted testdata fixtures. + /// See `testdata/manifests_lists/README.md`. + const FIXTURE_MASTER_KEY_ID: &str = "master-1"; + const FIXTURE_MASTER_KEY_BYTES: [u8; 16] = [ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, + 0x0f, + ]; + + /// Builds a `MemoryKmsClientFactory` seeded with the fixture master key. + fn fixture_kms_factory() -> MemoryKmsClientFactory { + use crate::encryption::SensitiveBytes; + + let factory = MemoryKmsClientFactory::new(); + factory + .add_master_key_bytes( + FIXTURE_MASTER_KEY_ID, + SensitiveBytes::new(FIXTURE_MASTER_KEY_BYTES), + ) + .unwrap(); + factory + } + + /// Loads the encrypted V3 metadata fixture and patches its snapshot's + /// manifest-list to point at the on-disk encrypted testdata file. + fn load_encrypted_fixture_metadata() -> TableMetadata { + let manifest_dir = env!("CARGO_MANIFEST_DIR"); + let metadata_json = std::fs::read_to_string(format!( + "{manifest_dir}/testdata/table_metadata/TableMetadataV3ValidEncryption.json" + )) + .unwrap(); + let mut metadata: TableMetadata = serde_json::from_str(&metadata_json).unwrap(); + + let manifest_list_path = + format!("{manifest_dir}/testdata/manifests_lists/manifest-list-v3-encrypted.avro"); + let snapshot = metadata.snapshots.get_mut(&1).unwrap(); + let mut patched = snapshot.as_ref().clone(); + patched.manifest_list = manifest_list_path; + *snapshot = Arc::new(patched); + + metadata + } + + #[tokio::test] + async fn catalog_kms_factory_client_reaches_table_encryption_manager() { + let warehouse = temp_path(); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .with_kms_client_factory(Arc::new(fixture_kms_factory())) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await + .unwrap(); + + let namespace_ident = NamespaceIdent::new("enc_ns".into()); + create_namespace(&catalog, &namespace_ident).await; + + let metadata = load_encrypted_fixture_metadata(); + let metadata_dir = TempDir::new().unwrap(); + let metadata_location = + format!("{}/v1.metadata.json", metadata_dir.path().to_str().unwrap()); + std::fs::write(&metadata_location, serde_json::to_vec(&metadata).unwrap()).unwrap(); + + let table_ident = TableIdent::new(namespace_ident, "enc".to_string()); + catalog + .register_table(&table_ident, metadata_location) + .await + .unwrap(); + + let table = catalog.load_table(&table_ident).await.unwrap(); + assert!( + table.encryption_manager().is_some(), + "factory-built KMS client should have reached the table's EncryptionManager" + ); + + let snapshot_ref = table.metadata().current_snapshot().unwrap(); + let manifest_list = table + .object_cache() + .get_manifest_list(snapshot_ref, &table.metadata_ref()) + .await + .unwrap(); + assert_eq!(manifest_list.entries().len(), 0); + } + fn build_table(ident: TableIdent) -> Table { let file_io = FileIO::new_with_fs(); diff --git a/crates/iceberg/src/catalog/mod.rs b/crates/iceberg/src/catalog/mod.rs index 3ab0f2886b..d2dc1ac1a0 100644 --- a/crates/iceberg/src/catalog/mod.rs +++ b/crates/iceberg/src/catalog/mod.rs @@ -39,6 +39,7 @@ use serde_derive::{Deserialize, Serialize}; use typed_builder::TypedBuilder; use uuid::Uuid; +use crate::encryption::kms::KmsClientFactory; use crate::io::StorageFactory; use crate::runtime::Runtime; use crate::spec::{ @@ -152,6 +153,28 @@ pub trait CatalogBuilder: Default + Debug + Send + Sync { /// ``` fn with_storage_factory(self, storage_factory: Arc) -> Self; + /// Set a [`KmsClientFactory`] to enable table encryption. + /// + /// When provided, the catalog calls the factory once during + /// [`load`](Self::load) with the catalog properties to create a shared + /// [`KeyManagementClient`](crate::encryption::KeyManagementClient). + /// That client is then passed to each table's `TableBuilder` so tables + /// with `encryption.key-id` set can construct an `EncryptionManager`. + /// + /// # Example + /// + /// ```rust,ignore + /// use iceberg::CatalogBuilder; + /// use iceberg::encryption::kms::KmsClientFactory; + /// use std::sync::Arc; + /// + /// let catalog = MyCatalogBuilder::default() + /// .with_kms_client_factory(Arc::new(MyKmsClientFactory)) + /// .load("my_catalog", props) + /// .await?; + /// ``` + fn with_kms_client_factory(self, kms_client_factory: Arc) -> Self; + /// Set a custom tokio Runtime to use for spawning async tasks. /// /// When a Runtime is provided, the catalog will propagate it to all tables diff --git a/crates/iceberg/src/encryption/kms/factory.rs b/crates/iceberg/src/encryption/kms/factory.rs new file mode 100644 index 0000000000..b957b5a5f6 --- /dev/null +++ b/crates/iceberg/src/encryption/kms/factory.rs @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Factory trait for creating [`KeyManagementClient`] instances. + +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; + +use super::KeyManagementClient; +use crate::Result; + +/// Factory for creating a [`KeyManagementClient`] from catalog properties. +/// +/// Replaces Java's reflection-based `encryption.kms-impl` + `initialize(properties)` +/// pattern. Users provide an implementation of this trait to the catalog builder via +/// [`CatalogBuilder::with_kms_client_factory`](crate::CatalogBuilder::with_kms_client_factory). +/// +/// The catalog calls [`create_kms_client`](Self::create_kms_client) **once** during +/// catalog initialization with the catalog's properties. The resulting client is +/// shared across all tables in the catalog and passed to each table's +/// [`EncryptionManager`](crate::encryption::EncryptionManager) via +/// `TableBuilder::kms_client(...)`. +#[async_trait] +pub trait KmsClientFactory: Debug + Send + Sync { + /// Create a [`KeyManagementClient`] from catalog properties. + /// + /// Called once during catalog initialization. Properties may include + /// KMS endpoint, region, credentials, or any backend-specific + /// configuration needed to construct the client. + async fn create_kms_client( + &self, + properties: &HashMap, + ) -> Result>; +} diff --git a/crates/iceberg/src/encryption/kms/memory.rs b/crates/iceberg/src/encryption/kms/memory.rs index 65319831dd..4c95b8421a 100644 --- a/crates/iceberg/src/encryption/kms/memory.rs +++ b/crates/iceberg/src/encryption/kms/memory.rs @@ -27,6 +27,7 @@ use std::sync::{Arc, RwLock}; use async_trait::async_trait; use super::KeyManagementClient; +use super::factory::KmsClientFactory; use crate::encryption::{AesGcmCipher, AesKeySize, SecureKey, SensitiveBytes}; use crate::error::lock_error; use crate::{Error, ErrorKind, Result}; @@ -142,6 +143,71 @@ impl MemoryKeyManagementClient { } } +/// Factory for creating [`MemoryKeyManagementClient`] instances. +/// +/// The factory owns the master-key table and its key size; seed it with +/// [`add_master_key`](Self::add_master_key) / +/// [`add_master_key_bytes`](Self::add_master_key_bytes), then every client it +/// produces via [`create_kms_client`](KmsClientFactory::create_kms_client) +/// shares that same table. Useful for testing encryption flows without a real +/// KMS backend. +#[derive(Debug, Clone, Default)] +pub struct MemoryKmsClientFactory { + master_keys: Arc>>, + master_key_size: AesKeySize, +} + +impl MemoryKmsClientFactory { + /// Creates a new factory with 128-bit AES master keys. + pub fn new() -> Self { + Self::default() + } + + /// Creates a new factory whose clients use the given master key size. + pub fn with_master_key_size(master_key_size: AesKeySize) -> Self { + Self { + master_keys: Arc::new(RwLock::new(HashMap::new())), + master_key_size, + } + } + + /// A client view over this factory's shared master-key table. + fn client(&self) -> MemoryKeyManagementClient { + MemoryKeyManagementClient { + master_keys: Arc::clone(&self.master_keys), + master_key_size: self.master_key_size, + } + } + + /// Adds a randomly generated master key with the given ID. + pub fn add_master_key(&self, key_id: impl Into) -> Result<()> { + self.client().add_master_key(key_id) + } + + /// Adds a master key with explicit key bytes. + /// + /// Use this to seed the factory with known key material, e.g. for + /// cross-language integration tests where both Java and Rust must + /// share the same master key bytes. + pub fn add_master_key_bytes( + &self, + key_id: impl Into, + key_bytes: SensitiveBytes, + ) -> Result<()> { + self.client().add_master_key_bytes(key_id, key_bytes) + } +} + +#[async_trait] +impl KmsClientFactory for MemoryKmsClientFactory { + async fn create_kms_client( + &self, + _properties: &HashMap, + ) -> Result> { + Ok(Arc::new(self.client())) + } +} + #[async_trait] impl KeyManagementClient for MemoryKeyManagementClient { async fn wrap_key(&self, key: &[u8], wrapping_key_id: &str) -> Result> { @@ -293,4 +359,46 @@ mod tests { kms1.add_master_key("shared-key").unwrap(); assert!(kms2.has_key("shared-key")); } + + #[tokio::test] + async fn test_factory_seeds_produced_clients() { + // Seed the factory, then every client it produces sees those keys and + // can wrap/unwrap with them. + let factory = MemoryKmsClientFactory::new(); + factory + .add_master_key_bytes("master-1", SensitiveBytes::new([9u8; 16])) + .unwrap(); + + let client = factory.create_kms_client(&HashMap::new()).await.unwrap(); + + let dek = vec![3u8; 16]; + let wrapped = client.wrap_key(&dek, "master-1").await.unwrap(); + let unwrapped = client.unwrap_key(&wrapped, "master-1").await.unwrap(); + assert_eq!(unwrapped.as_bytes(), dek.as_slice()); + } + + #[tokio::test] + async fn test_factory_seeding_after_client_creation_is_visible() { + // The factory owns the shared table, so keys added after a client is + // produced are still visible to that client. + let factory = MemoryKmsClientFactory::new(); + let client = factory.create_kms_client(&HashMap::new()).await.unwrap(); + + factory.add_master_key("late-key").unwrap(); + + let dek = vec![1u8; 16]; + assert!(client.wrap_key(&dek, "late-key").await.is_ok()); + } + + #[tokio::test] + async fn test_factory_with_master_key_size() { + let factory = MemoryKmsClientFactory::with_master_key_size(AesKeySize::Bits256); + factory.add_master_key("master-256").unwrap(); + + let client = factory.create_kms_client(&HashMap::new()).await.unwrap(); + let dek = vec![0u8; 16]; + let wrapped = client.wrap_key(&dek, "master-256").await.unwrap(); + let unwrapped = client.unwrap_key(&wrapped, "master-256").await.unwrap(); + assert_eq!(unwrapped.as_bytes(), dek.as_slice()); + } } diff --git a/crates/iceberg/src/encryption/kms/mod.rs b/crates/iceberg/src/encryption/kms/mod.rs index 160e692550..0b4f331d13 100644 --- a/crates/iceberg/src/encryption/kms/mod.rs +++ b/crates/iceberg/src/encryption/kms/mod.rs @@ -21,7 +21,9 @@ //! integration and implementations for different key management systems. mod client; +mod factory; mod memory; pub use client::{GeneratedKey, KeyManagementClient}; -pub use memory::MemoryKeyManagementClient; +pub use factory::KmsClientFactory; +pub use memory::{MemoryKeyManagementClient, MemoryKmsClientFactory};