Skip to content

Commit 7bd5511

Browse files
committed
refactor(hms): derive catalog configuration from properties
Use the shared property derive for HMS settings while retaining address and warehouse validation and transport fallback behavior. Tracks [#3098](#3098). Generated-by: Codex
1 parent 75c9201 commit 7bd5511

3 files changed

Lines changed: 183 additions & 73 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/catalog/hms/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ async-trait = { workspace = true }
3535
chrono = { workspace = true }
3636
hive_metastore = { workspace = true }
3737
iceberg = { workspace = true }
38+
iceberg-property-macro = { workspace = true }
3839
pilota = { workspace = true }
3940
serde_json = { workspace = true }
4041
tokio = { workspace = true }

‎crates/catalog/hms/src/catalog.rs‎

Lines changed: 181 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ use iceberg::{
3434
Catalog, CatalogBuilder, Error, ErrorKind, MetadataLocation, Namespace, NamespaceIdent, Result,
3535
Runtime, TableCommit, TableCreation, TableIdent,
3636
};
37+
use iceberg_property_macro::Properties;
3738
use volo_thrift::MaybeException;
3839

3940
use super::utils::*;
@@ -55,22 +56,15 @@ pub const HMS_CATALOG_PROP_WAREHOUSE: &str = "warehouse";
5556
/// Builder for [`HmsCatalog`].
5657
#[derive(Debug)]
5758
pub struct HmsCatalogBuilder {
58-
config: HmsCatalogConfig,
5959
storage_factory: Option<Arc<dyn StorageFactory>>,
6060
kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
6161
runtime: Option<Runtime>,
6262
}
6363

64+
#[allow(clippy::derivable_impls)]
6465
impl Default for HmsCatalogBuilder {
6566
fn default() -> Self {
6667
Self {
67-
config: HmsCatalogConfig {
68-
name: None,
69-
address: "".to_string(),
70-
thrift_transport: HmsThriftTransport::default(),
71-
warehouse: "".to_string(),
72-
props: HashMap::new(),
73-
},
7468
storage_factory: None,
7569
kms_client_factory: None,
7670
runtime: None,
@@ -97,59 +91,21 @@ impl CatalogBuilder for HmsCatalogBuilder {
9791
}
9892

9993
fn load(
100-
mut self,
94+
self,
10195
name: impl Into<String>,
10296
props: HashMap<String, String>,
10397
) -> impl Future<Output = Result<Self::C>> + Send {
104-
self.config.name = Some(name.into());
105-
106-
if props.contains_key(HMS_CATALOG_PROP_URI) {
107-
self.config.address = props.get(HMS_CATALOG_PROP_URI).cloned().unwrap_or_default();
108-
}
109-
110-
if let Some(tt) = props.get(HMS_CATALOG_PROP_THRIFT_TRANSPORT) {
111-
self.config.thrift_transport = match tt.to_lowercase().as_str() {
112-
THRIFT_TRANSPORT_FRAMED => HmsThriftTransport::Framed,
113-
THRIFT_TRANSPORT_BUFFERED => HmsThriftTransport::Buffered,
114-
_ => HmsThriftTransport::default(),
115-
};
116-
}
117-
118-
if props.contains_key(HMS_CATALOG_PROP_WAREHOUSE) {
119-
self.config.warehouse = props
120-
.get(HMS_CATALOG_PROP_WAREHOUSE)
121-
.cloned()
122-
.unwrap_or_default();
123-
}
124-
125-
self.config.props = props
126-
.into_iter()
127-
.filter(|(k, _)| {
128-
k != HMS_CATALOG_PROP_URI
129-
&& k != HMS_CATALOG_PROP_THRIFT_TRANSPORT
130-
&& k != HMS_CATALOG_PROP_WAREHOUSE
131-
})
132-
.collect();
98+
let name = name.into();
13399

134100
async move {
135-
let kms_client = match self.kms_client_factory {
136-
Some(factory) => Some(factory.create_kms_client(&self.config.props).await?),
137-
None => None,
138-
};
139-
140-
if self.config.name.is_none() {
141-
return Err(Error::new(
142-
ErrorKind::DataInvalid,
143-
"Catalog name is required",
144-
));
145-
}
146-
if self.config.address.is_empty() {
101+
let catalog_properties = HmsCatalogProperties::from_properties(&props)?;
102+
if catalog_properties.address.is_empty() {
147103
return Err(Error::new(
148104
ErrorKind::DataInvalid,
149105
"Catalog address is required",
150106
));
151107
}
152-
if self.config.warehouse.is_empty() {
108+
if catalog_properties.warehouse.is_empty() {
153109
return Err(Error::new(
154110
ErrorKind::DataInvalid,
155111
"Catalog warehouse is required",
@@ -159,11 +115,44 @@ impl CatalogBuilder for HmsCatalogBuilder {
159115
Some(rt) => rt,
160116
None => Runtime::try_current()?,
161117
};
162-
HmsCatalog::new(self.config, self.storage_factory, runtime, kms_client)
118+
let kms_client = match self.kms_client_factory {
119+
Some(factory) => Some(factory.create_kms_client(&props).await?),
120+
None => None,
121+
};
122+
HmsCatalog::new(
123+
name,
124+
catalog_properties,
125+
props,
126+
self.storage_factory,
127+
runtime,
128+
kms_client,
129+
)
163130
}
164131
}
165132
}
166133

134+
fn parse_thrift_transport(value: &str) -> Result<HmsThriftTransport> {
135+
Ok(match value.to_lowercase().as_str() {
136+
THRIFT_TRANSPORT_FRAMED => HmsThriftTransport::Framed,
137+
THRIFT_TRANSPORT_BUFFERED => HmsThriftTransport::Buffered,
138+
_ => HmsThriftTransport::default(),
139+
})
140+
}
141+
142+
#[derive(Debug, Properties)]
143+
pub(crate) struct HmsCatalogProperties {
144+
#[property(key = HMS_CATALOG_PROP_URI, default = "")]
145+
address: String,
146+
#[property(
147+
key = HMS_CATALOG_PROP_THRIFT_TRANSPORT,
148+
default = HmsThriftTransport::default(),
149+
parse_with = parse_thrift_transport
150+
)]
151+
thrift_transport: HmsThriftTransport,
152+
#[property(key = HMS_CATALOG_PROP_WAREHOUSE, default = "")]
153+
warehouse: String,
154+
}
155+
167156
/// Which variant of the thrift transport to communicate with HMS
168157
/// See: <https://github.com/apache/thrift/blob/master/doc/specs/thrift-rpc.md#framed-vs-unframed-transport>
169158
#[derive(Debug, Default)]
@@ -175,21 +164,12 @@ pub enum HmsThriftTransport {
175164
Buffered,
176165
}
177166

178-
/// Hive metastore Catalog configuration.
179-
#[derive(Debug)]
180-
pub(crate) struct HmsCatalogConfig {
181-
name: Option<String>,
182-
address: String,
183-
thrift_transport: HmsThriftTransport,
184-
warehouse: String,
185-
props: HashMap<String, String>,
186-
}
187-
188167
struct HmsClient(ThriftHiveMetastoreClient);
189168

190169
/// Hive metastore Catalog.
191170
pub struct HmsCatalog {
192-
config: HmsCatalogConfig,
171+
name: String,
172+
properties: HmsCatalogProperties,
193173
client: HmsClient,
194174
file_io: FileIO,
195175
runtime: Runtime,
@@ -199,20 +179,23 @@ pub struct HmsCatalog {
199179
impl Debug for HmsCatalog {
200180
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
201181
f.debug_struct("HmsCatalog")
202-
.field("config", &self.config)
182+
.field("name", &self.name)
183+
.field("properties", &self.properties)
203184
.finish_non_exhaustive()
204185
}
205186
}
206187

207188
impl HmsCatalog {
208189
/// Create a new hms catalog.
209190
fn new(
210-
config: HmsCatalogConfig,
191+
name: String,
192+
properties: HmsCatalogProperties,
193+
props: HashMap<String, String>,
211194
storage_factory: Option<Arc<dyn StorageFactory>>,
212195
runtime: Runtime,
213196
kms_client: Option<Arc<dyn KeyManagementClient>>,
214197
) -> Result<Self> {
215-
let address = config
198+
let address = properties
216199
.address
217200
.as_str()
218201
.to_socket_addrs()
@@ -221,13 +204,13 @@ impl HmsCatalog {
221204
.ok_or_else(|| {
222205
Error::new(
223206
ErrorKind::Unexpected,
224-
format!("invalid address: {}", config.address),
207+
format!("invalid address: {}", properties.address),
225208
)
226209
})?;
227210

228211
let builder = ThriftHiveMetastoreClientBuilder::new("hms").address(address);
229212

230-
let client = match &config.thrift_transport {
213+
let client = match &properties.thrift_transport {
231214
HmsThriftTransport::Framed => builder
232215
.make_codec(volo_thrift::codec::default::DefaultMakeCodec::framed())
233216
.build(),
@@ -242,12 +225,11 @@ impl HmsCatalog {
242225
"StorageFactory must be provided for HmsCatalog. Use `with_storage_factory` to configure it.",
243226
)
244227
})?;
245-
let file_io = FileIOBuilder::new(factory)
246-
.with_props(&config.props)
247-
.build();
228+
let file_io = FileIOBuilder::new(factory).with_props(props).build();
248229

249230
Ok(Self {
250-
config,
231+
name,
232+
properties,
251233
client: HmsClient(client),
252234
file_io,
253235
runtime,
@@ -524,7 +506,8 @@ impl Catalog for HmsCatalog {
524506
Some(location) => location.clone(),
525507
None => {
526508
let ns = self.get_namespace(namespace).await?;
527-
let location = get_default_table_location(&ns, &table_name, &self.config.warehouse);
509+
let location =
510+
get_default_table_location(&ns, &table_name, &self.properties.warehouse);
528511
creation.location = Some(location.clone());
529512
location
530513
}
@@ -733,3 +716,128 @@ impl Catalog for HmsCatalog {
733716
))
734717
}
735718
}
719+
720+
#[cfg(test)]
721+
mod catalog_properties_tests {
722+
use iceberg::io::MemoryStorageFactory;
723+
724+
use super::*;
725+
726+
#[test]
727+
fn test_catalog_properties() {
728+
let properties = HmsCatalogProperties::from_properties(&HashMap::from([
729+
(
730+
HMS_CATALOG_PROP_URI.to_string(),
731+
"localhost:9083".to_string(),
732+
),
733+
(
734+
HMS_CATALOG_PROP_THRIFT_TRANSPORT.to_string(),
735+
"FRAMED".to_string(),
736+
),
737+
(
738+
HMS_CATALOG_PROP_WAREHOUSE.to_string(),
739+
"s3://warehouse".to_string(),
740+
),
741+
]))
742+
.unwrap();
743+
744+
assert_eq!(properties.address, "localhost:9083");
745+
assert!(matches!(
746+
properties.thrift_transport,
747+
HmsThriftTransport::Framed
748+
));
749+
assert_eq!(properties.warehouse, "s3://warehouse");
750+
}
751+
752+
#[test]
753+
fn test_catalog_properties_defaults() {
754+
let properties = HmsCatalogProperties::from_properties(&HashMap::new()).unwrap();
755+
756+
assert_eq!(properties.address, "");
757+
assert!(matches!(
758+
properties.thrift_transport,
759+
HmsThriftTransport::Buffered
760+
));
761+
assert_eq!(properties.warehouse, "");
762+
}
763+
764+
#[test]
765+
fn test_invalid_thrift_transport_uses_default() {
766+
let properties = HmsCatalogProperties::from_properties(&HashMap::from([(
767+
HMS_CATALOG_PROP_THRIFT_TRANSPORT.to_string(),
768+
"unknown".to_string(),
769+
)]))
770+
.unwrap();
771+
772+
assert!(matches!(
773+
properties.thrift_transport,
774+
HmsThriftTransport::Buffered
775+
));
776+
}
777+
778+
#[tokio::test]
779+
async fn test_load_rejects_missing_address() {
780+
let error = HmsCatalogBuilder::default()
781+
.load(
782+
"hms",
783+
HashMap::from([(
784+
HMS_CATALOG_PROP_WAREHOUSE.to_string(),
785+
"memory:///warehouse".to_string(),
786+
)]),
787+
)
788+
.await
789+
.unwrap_err();
790+
791+
assert_eq!(error.kind(), ErrorKind::DataInvalid);
792+
assert_eq!(error.message(), "Catalog address is required");
793+
}
794+
795+
#[tokio::test]
796+
async fn test_load_rejects_missing_warehouse() {
797+
let error = HmsCatalogBuilder::default()
798+
.load(
799+
"hms",
800+
HashMap::from([(
801+
HMS_CATALOG_PROP_URI.to_string(),
802+
"localhost:9083".to_string(),
803+
)]),
804+
)
805+
.await
806+
.unwrap_err();
807+
808+
assert_eq!(error.kind(), ErrorKind::DataInvalid);
809+
assert_eq!(error.message(), "Catalog warehouse is required");
810+
}
811+
812+
#[tokio::test]
813+
async fn test_catalog_forwards_properties_to_file_io() {
814+
let catalog = HmsCatalogBuilder::default()
815+
.with_storage_factory(Arc::new(MemoryStorageFactory))
816+
.load(
817+
"hms",
818+
HashMap::from([
819+
(
820+
HMS_CATALOG_PROP_URI.to_string(),
821+
"localhost:9083".to_string(),
822+
),
823+
(
824+
HMS_CATALOG_PROP_WAREHOUSE.to_string(),
825+
"memory:///warehouse".to_string(),
826+
),
827+
("custom.property".to_string(), "value".to_string()),
828+
]),
829+
)
830+
.await
831+
.unwrap();
832+
let file_io_props = catalog.file_io.config().props();
833+
834+
assert_eq!(
835+
file_io_props.get("custom.property"),
836+
Some(&"value".to_string())
837+
);
838+
assert_eq!(
839+
file_io_props.get(HMS_CATALOG_PROP_WAREHOUSE),
840+
Some(&"memory:///warehouse".to_string())
841+
);
842+
}
843+
}

0 commit comments

Comments
 (0)