1use std::collections::HashMap;
19use std::fmt::{Debug, Formatter};
20use std::net::ToSocketAddrs;
21use std::sync::Arc;
22
23use anyhow::anyhow;
24use async_trait::async_trait;
25use hive_metastore::{
26 ThriftHiveMetastoreClient, ThriftHiveMetastoreClientBuilder,
27 ThriftHiveMetastoreGetDatabaseException, ThriftHiveMetastoreGetTableException,
28};
29use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory};
30use iceberg::io::{FileIO, FileIOBuilder, StorageFactory};
31use iceberg::spec::{TableMetadata, TableMetadataBuilder};
32use iceberg::table::Table;
33use iceberg::{
34 Catalog, CatalogBuilder, Error, ErrorKind, MetadataLocation, Namespace, NamespaceIdent, Result,
35 Runtime, TableCommit, TableCreation, TableIdent,
36};
37use volo_thrift::MaybeException;
38
39use super::utils::*;
40use crate::error::{from_io_error, from_thrift_error, from_thrift_exception};
41
42pub const HMS_CATALOG_PROP_URI: &str = "uri";
44
45pub const HMS_CATALOG_PROP_THRIFT_TRANSPORT: &str = "thrift_transport";
47pub const THRIFT_TRANSPORT_FRAMED: &str = "framed";
49pub const THRIFT_TRANSPORT_BUFFERED: &str = "buffered";
51
52pub const HMS_CATALOG_PROP_WAREHOUSE: &str = "warehouse";
54
55#[derive(Debug)]
57pub struct HmsCatalogBuilder {
58 config: HmsCatalogConfig,
59 storage_factory: Option<Arc<dyn StorageFactory>>,
60 kms_client_factory: Option<Arc<dyn KmsClientFactory>>,
61 runtime: Option<Runtime>,
62}
63
64impl Default for HmsCatalogBuilder {
65 fn default() -> Self {
66 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 },
74 storage_factory: None,
75 kms_client_factory: None,
76 runtime: None,
77 }
78 }
79}
80
81impl CatalogBuilder for HmsCatalogBuilder {
82 type C = HmsCatalog;
83
84 fn with_storage_factory(mut self, storage_factory: Arc<dyn StorageFactory>) -> Self {
85 self.storage_factory = Some(storage_factory);
86 self
87 }
88
89 fn with_kms_client_factory(mut self, kms_client_factory: Arc<dyn KmsClientFactory>) -> Self {
90 self.kms_client_factory = Some(kms_client_factory);
91 self
92 }
93
94 fn with_runtime(mut self, runtime: Runtime) -> Self {
95 self.runtime = Some(runtime);
96 self
97 }
98
99 fn load(
100 mut self,
101 name: impl Into<String>,
102 props: HashMap<String, String>,
103 ) -> 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();
133
134 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() {
147 return Err(Error::new(
148 ErrorKind::DataInvalid,
149 "Catalog address is required",
150 ));
151 }
152 if self.config.warehouse.is_empty() {
153 return Err(Error::new(
154 ErrorKind::DataInvalid,
155 "Catalog warehouse is required",
156 ));
157 }
158 let runtime = match self.runtime {
159 Some(rt) => rt,
160 None => Runtime::try_current()?,
161 };
162 HmsCatalog::new(self.config, self.storage_factory, runtime, kms_client)
163 }
164 }
165}
166
167#[derive(Debug, Default)]
170pub enum HmsThriftTransport {
171 Framed,
173 #[default]
175 Buffered,
176}
177
178#[derive(Debug)]
180pub(crate) struct HmsCatalogConfig {
181 name: Option<String>,
182 address: String,
183 thrift_transport: HmsThriftTransport,
184 warehouse: String,
185 props: HashMap<String, String>,
186}
187
188struct HmsClient(ThriftHiveMetastoreClient);
189
190pub struct HmsCatalog {
192 config: HmsCatalogConfig,
193 client: HmsClient,
194 file_io: FileIO,
195 runtime: Runtime,
196 kms_client: Option<Arc<dyn KeyManagementClient>>,
197}
198
199impl Debug for HmsCatalog {
200 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
201 f.debug_struct("HmsCatalog")
202 .field("config", &self.config)
203 .finish_non_exhaustive()
204 }
205}
206
207impl HmsCatalog {
208 fn new(
210 config: HmsCatalogConfig,
211 storage_factory: Option<Arc<dyn StorageFactory>>,
212 runtime: Runtime,
213 kms_client: Option<Arc<dyn KeyManagementClient>>,
214 ) -> Result<Self> {
215 let address = config
216 .address
217 .as_str()
218 .to_socket_addrs()
219 .map_err(from_io_error)?
220 .next()
221 .ok_or_else(|| {
222 Error::new(
223 ErrorKind::Unexpected,
224 format!("invalid address: {}", config.address),
225 )
226 })?;
227
228 let builder = ThriftHiveMetastoreClientBuilder::new("hms").address(address);
229
230 let client = match &config.thrift_transport {
231 HmsThriftTransport::Framed => builder
232 .make_codec(volo_thrift::codec::default::DefaultMakeCodec::framed())
233 .build(),
234 HmsThriftTransport::Buffered => builder
235 .make_codec(volo_thrift::codec::default::DefaultMakeCodec::buffered())
236 .build(),
237 };
238
239 let factory = storage_factory.ok_or_else(|| {
240 Error::new(
241 ErrorKind::Unexpected,
242 "StorageFactory must be provided for HmsCatalog. Use `with_storage_factory` to configure it.",
243 )
244 })?;
245 let file_io = FileIOBuilder::new(factory)
246 .with_props(&config.props)
247 .build();
248
249 Ok(Self {
250 config,
251 client: HmsClient(client),
252 file_io,
253 runtime,
254 kms_client,
255 })
256 }
257 pub fn file_io(&self) -> FileIO {
259 self.file_io.clone()
260 }
261}
262
263#[async_trait]
264impl Catalog for HmsCatalog {
265 async fn list_namespaces(
271 &self,
272 parent: Option<&NamespaceIdent>,
273 ) -> Result<Vec<NamespaceIdent>> {
274 let dbs = if parent.is_some() {
275 return Ok(vec![]);
276 } else {
277 self.client
278 .0
279 .get_all_databases()
280 .await
281 .map(from_thrift_exception)
282 .map_err(from_thrift_error)??
283 };
284
285 Ok(dbs
286 .into_iter()
287 .map(|v| NamespaceIdent::new(v.into()))
288 .collect())
289 }
290
291 async fn create_namespace(
307 &self,
308 namespace: &NamespaceIdent,
309 properties: HashMap<String, String>,
310 ) -> Result<Namespace> {
311 if self.namespace_exists(namespace).await? {
312 return Err(Error::new(
313 ErrorKind::NamespaceAlreadyExists,
314 format!("Namespace {namespace:?} already exists"),
315 ));
316 }
317 let database = convert_to_database(namespace, &properties)?;
318
319 self.client
320 .0
321 .create_database(database)
322 .await
323 .map_err(from_thrift_error)?;
324
325 Ok(Namespace::with_properties(namespace.clone(), properties))
326 }
327
328 async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result<Namespace> {
339 let name = validate_namespace(namespace)?;
340
341 let resp = self
342 .client
343 .0
344 .get_database(name.into())
345 .await
346 .map_err(from_thrift_error)?;
347
348 let db = match resp {
349 MaybeException::Ok(db) => db,
350 MaybeException::Exception(ThriftHiveMetastoreGetDatabaseException::O1(_)) => {
351 return Err(Error::new(
352 ErrorKind::NamespaceNotFound,
353 format!("Namespace {namespace:?} not found"),
354 ));
355 }
356 MaybeException::Exception(exception) => {
357 return Err(Error::new(
358 ErrorKind::Unexpected,
359 "Operation failed for hitting thrift error".to_string(),
360 )
361 .with_source(anyhow!("thrift error: {exception:?}")));
362 }
363 };
364
365 let ns = convert_to_namespace(&db)?;
366
367 Ok(ns)
368 }
369
370 async fn namespace_exists(&self, namespace: &NamespaceIdent) -> Result<bool> {
383 let name = validate_namespace(namespace)?;
384
385 let resp = self.client.0.get_database(name.into()).await;
386
387 match resp {
388 Ok(MaybeException::Ok(_)) => Ok(true),
389 Ok(MaybeException::Exception(ThriftHiveMetastoreGetDatabaseException::O1(_))) => {
390 Ok(false)
391 }
392 Ok(MaybeException::Exception(exception)) => Err(Error::new(
393 ErrorKind::Unexpected,
394 "Operation failed for hitting thrift error".to_string(),
395 )
396 .with_source(anyhow!("thrift error: {exception:?}"))),
397 Err(err) => Err(from_thrift_error(err)),
398 }
399 }
400
401 async fn update_namespace(
412 &self,
413 namespace: &NamespaceIdent,
414 properties: HashMap<String, String>,
415 ) -> Result<()> {
416 if !self.namespace_exists(namespace).await? {
417 return Err(Error::new(
418 ErrorKind::NamespaceNotFound,
419 format!("Namespace {namespace:?} does not exist"),
420 ));
421 }
422 let db = convert_to_database(namespace, &properties)?;
423
424 let name = match &db.name {
425 Some(name) => name,
426 None => {
427 return Err(Error::new(
428 ErrorKind::DataInvalid,
429 "Database name must be specified",
430 ));
431 }
432 };
433
434 self.client
435 .0
436 .alter_database(name.clone(), db)
437 .await
438 .map_err(from_thrift_error)?;
439
440 Ok(())
441 }
442
443 async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> {
451 let name = validate_namespace(namespace)?;
452
453 if !self.namespace_exists(namespace).await? {
454 return Err(Error::new(
455 ErrorKind::NamespaceNotFound,
456 format!("Namespace {namespace:?} does not exist"),
457 ));
458 }
459
460 self.client
461 .0
462 .drop_database(name.into(), false, false)
463 .await
464 .map_err(from_thrift_error)?;
465
466 Ok(())
467 }
468
469 async fn list_tables(&self, namespace: &NamespaceIdent) -> Result<Vec<TableIdent>> {
479 let name = validate_namespace(namespace)?;
480 if !self.namespace_exists(namespace).await? {
481 return Err(Error::new(
482 ErrorKind::NamespaceNotFound,
483 format!("Namespace {namespace:?} does not exist"),
484 ));
485 }
486
487 let tables = self
488 .client
489 .0
490 .get_all_tables(name.into())
491 .await
492 .map(from_thrift_exception)
493 .map_err(from_thrift_error)??;
494
495 let tables = tables
496 .iter()
497 .map(|table| TableIdent::new(namespace.clone(), table.to_string()))
498 .collect();
499
500 Ok(tables)
501 }
502
503 async fn create_table(
516 &self,
517 namespace: &NamespaceIdent,
518 mut creation: TableCreation,
519 ) -> Result<Table> {
520 let db_name = validate_namespace(namespace)?;
521 let table_name = creation.name.clone();
522
523 let location = match &creation.location {
524 Some(location) => location.clone(),
525 None => {
526 let ns = self.get_namespace(namespace).await?;
527 let location = get_default_table_location(&ns, &table_name, &self.config.warehouse);
528 creation.location = Some(location.clone());
529 location
530 }
531 };
532 let metadata = TableMetadataBuilder::from_table_creation(creation)?
533 .build()?
534 .metadata;
535
536 let metadata_location = MetadataLocation::new_with_metadata(location.clone(), &metadata);
537
538 metadata.write_to(&self.file_io, &metadata_location).await?;
539
540 let metadata_location_str = metadata_location.to_string();
541 let hive_table = convert_to_hive_table(
542 db_name.clone(),
543 metadata.current_schema(),
544 table_name.clone(),
545 location,
546 metadata_location_str.clone(),
547 metadata.properties(),
548 )?;
549
550 self.client
551 .0
552 .create_table(hive_table)
553 .await
554 .map_err(from_thrift_error)?;
555
556 let mut builder = Table::builder()
557 .file_io(self.file_io())
558 .metadata_location(metadata_location_str)
559 .metadata(metadata)
560 .identifier(TableIdent::new(NamespaceIdent::new(db_name), table_name))
561 .runtime(self.runtime.clone());
562 if let Some(kms_client) = self.kms_client.clone() {
563 builder = builder.kms_client(kms_client);
564 }
565 builder.build()
566 }
567
568 async fn load_table(&self, table: &TableIdent) -> Result<Table> {
581 let db_name = validate_namespace(table.namespace())?;
582
583 let hive_table = self
584 .client
585 .0
586 .get_table(db_name.clone().into(), table.name.clone().into())
587 .await
588 .map(from_thrift_exception)
589 .map_err(from_thrift_error)??;
590
591 let metadata_location = get_metadata_location(&hive_table.parameters)?;
592
593 let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?;
594
595 let mut builder = Table::builder()
596 .file_io(self.file_io())
597 .metadata_location(metadata_location)
598 .metadata(metadata)
599 .identifier(TableIdent::new(
600 NamespaceIdent::new(db_name),
601 table.name.clone(),
602 ))
603 .runtime(self.runtime.clone());
604 if let Some(kms_client) = self.kms_client.clone() {
605 builder = builder.kms_client(kms_client);
606 }
607 builder.build()
608 }
609
610 async fn drop_table(&self, table: &TableIdent) -> Result<()> {
621 let db_name = validate_namespace(table.namespace())?;
622 if !self.namespace_exists(table.namespace()).await? {
623 return Err(Error::new(
624 ErrorKind::NamespaceNotFound,
625 format!("Namespace {:?} does not exist", table.namespace()),
626 ));
627 }
628 if !self.table_exists(table).await? {
629 return Err(Error::new(
630 ErrorKind::TableNotFound,
631 format!("Table {table:?} does not exist"),
632 ));
633 }
634
635 self.client
636 .0
637 .drop_table(db_name.into(), table.name.clone().into(), false)
638 .await
639 .map_err(from_thrift_error)?;
640
641 Ok(())
642 }
643
644 async fn purge_table(&self, table: &TableIdent) -> Result<()> {
645 let table_info = self.load_table(table).await?;
646 self.drop_table(table).await?;
647 iceberg::drop_table_data(&table_info).await
648 }
649
650 async fn table_exists(&self, table: &TableIdent) -> Result<bool> {
658 let db_name = validate_namespace(table.namespace())?;
659 let table_name = table.name.clone();
660
661 let resp = self
662 .client
663 .0
664 .get_table(db_name.into(), table_name.into())
665 .await;
666
667 match resp {
668 Ok(MaybeException::Ok(_)) => Ok(true),
669 Ok(MaybeException::Exception(ThriftHiveMetastoreGetTableException::O2(_))) => Ok(false),
670 Ok(MaybeException::Exception(exception)) => Err(Error::new(
671 ErrorKind::Unexpected,
672 "Operation failed for hitting thrift error".to_string(),
673 )
674 .with_source(anyhow!("thrift error: {exception:?}"))),
675 Err(err) => Err(from_thrift_error(err)),
676 }
677 }
678
679 async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> {
686 let src_dbname = validate_namespace(src.namespace())?;
687 let dest_dbname = validate_namespace(dest.namespace())?;
688 if self.table_exists(dest).await? {
689 return Err(Error::new(
690 ErrorKind::TableAlreadyExists,
691 format!("Destination table {dest:?} already exists"),
692 ));
693 }
694
695 let src_tbl_name = src.name.clone();
696 let dest_tbl_name = dest.name.clone();
697
698 let mut tbl = self
699 .client
700 .0
701 .get_table(src_dbname.clone().into(), src_tbl_name.clone().into())
702 .await
703 .map(from_thrift_exception)
704 .map_err(from_thrift_error)??;
705
706 tbl.db_name = Some(dest_dbname.into());
707 tbl.table_name = Some(dest_tbl_name.into());
708
709 self.client
710 .0
711 .alter_table(src_dbname.into(), src_tbl_name.into(), tbl)
712 .await
713 .map_err(from_thrift_error)?;
714
715 Ok(())
716 }
717
718 async fn register_table(
719 &self,
720 _table_ident: &TableIdent,
721 _metadata_location: String,
722 ) -> Result<Table> {
723 Err(Error::new(
724 ErrorKind::FeatureUnsupported,
725 "Registering a table is not supported yet",
726 ))
727 }
728
729 async fn update_table(&self, _commit: TableCommit) -> Result<Table> {
730 Err(Error::new(
731 ErrorKind::FeatureUnsupported,
732 "Updating a table is not supported yet",
733 ))
734 }
735}