Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/catalog/hms/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ linkedbytes = { workspace = true }
metainfo = { workspace = true }
motore-macros = { workspace = true }
volo = { workspace = true }
whoami = "1.6.1"

[dev-dependencies]
ctor = { workspace = true }
Expand Down
198 changes: 192 additions & 6 deletions crates/catalog/hms/src/catalog.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,14 @@ pub const THRIFT_TRANSPORT_BUFFERED: &str = "buffered";
/// HMS Catalog warehouse location
pub const HMS_CATALOG_PROP_WAREHOUSE: &str = "warehouse";

/// Builder for [`RestCatalog`].
///HMS Hive Locks Disabled
pub const HMS_HIVE_LOCKS_DISABLED: &str = "hive_locks_disabled";

/// HMS Environment Context
const HMS_EXPECTED_PARAMETER_KEY: &str = "expected_parameter_key";
const HMS_EXPECTED_PARAMETER_VALUE: &str = "expected_parameter_value";

/// Builder for [`HmsCatalog`].
#[derive(Debug)]
pub struct HmsCatalogBuilder(HmsCatalogConfig);

Expand Down Expand Up @@ -167,6 +174,43 @@ impl Debug for HmsCatalog {
}
}

/// RAII guard for HMS table locks. Automatically releases the lock when dropped.
struct HmsLockGuard {
client: ThriftHiveMetastoreClient,
lockid: i64,
}

impl HmsLockGuard {
async fn acquire(
client: &ThriftHiveMetastoreClient,
db_name: &str,
tbl_name: &str,
) -> Result<Self> {
let lock = client
.lock(create_lock_request(db_name, tbl_name))
.await
.map(from_thrift_exception)
.map_err(from_thrift_error)??;

Ok(Self {
client: client.clone(),
lockid: lock.lockid,
})
}
}

impl Drop for HmsLockGuard {
fn drop(&mut self) {
let client = self.client.clone();
let lockid = self.lockid;
tokio::spawn(async move {
let _ = client
.unlock(hive_metastore::UnlockRequest { lockid })
.await;
});
}
}

impl HmsCatalog {
/// Create a new hms catalog.
fn new(config: HmsCatalogConfig) -> Result<Self> {
Expand Down Expand Up @@ -208,6 +252,64 @@ impl HmsCatalog {
pub fn file_io(&self) -> FileIO {
self.file_io.clone()
}

/// Applies a commit to a table and prepares the update for HMS.
/// # Returns
/// A tuple of (staged_table, new_hive_table) ready for HMS alter_table operation
async fn apply_and_prepare_update(
&self,
commit: TableCommit,
db_name: &str,
tbl_name: &str,
hive_table: &hive_metastore::Table,
) -> Result<(Table, hive_metastore::Table)> {
let metadata_location = get_metadata_location(&hive_table.parameters)?;

let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?;

let cur_table = Table::builder()
.file_io(self.file_io())
.metadata_location(metadata_location)
.metadata(metadata)
.identifier(TableIdent::new(
NamespaceIdent::new(db_name.to_string()),
tbl_name.to_string(),
))
.build()?;

let staged_table = commit.apply(cur_table)?;
staged_table
.metadata()
.write_to(
staged_table.file_io(),
staged_table.metadata_location_result()?,
)
.await?;

let new_hive_table = update_hive_table_from_table(hive_table, &staged_table)?;

Ok((staged_table, new_hive_table))
}

/// Builds an EnvironmentContext for optimistic locking with HMS.
///
/// The context includes the expected metadata_location, which HMS will use
/// to validate that the table hasn't been modified concurrently.
fn build_environment_context(metadata_location: &str) -> hive_metastore::EnvironmentContext {
let mut env_context_properties = pilota::AHashMap::new();
env_context_properties.insert(
HMS_EXPECTED_PARAMETER_KEY.into(),
"metadata_location".into(),
);
env_context_properties.insert(
HMS_EXPECTED_PARAMETER_VALUE.into(),
pilota::FastStr::from_string(metadata_location.to_string()),
);

hive_metastore::EnvironmentContext {
properties: Some(env_context_properties),
}
}
}

#[async_trait]
Expand Down Expand Up @@ -603,10 +705,94 @@ impl Catalog for HmsCatalog {
))
}

async fn update_table(&self, _commit: TableCommit) -> Result<Table> {
Err(Error::new(
ErrorKind::FeatureUnsupported,
"Updating a table is not supported yet",
))
/// Updates an existing table by applying a commit operation.
///
/// This method supports two update strategies depending on the catalog configuration:
///
/// **Optimistic Locking** (when `hive_locks_disabled` is set):
/// - Retrieves the current table state from HMS without acquiring locks
/// - Constructs an `EnvironmentContext` with the expected metadata location
/// - Uses `alter_table_with_environment_context` to perform an atomic
/// compare-and-swap operation.
/// - HMS will reject the update if the metadata location has changed,
/// indicating a concurrent modification
///
/// **Traditional Locking** (default):
/// - Acquires an exclusive HMS lock on the table before making changes
/// - Retrieves the current table state
/// - Applies the commit and writes new metadata
/// - Updates the table in HMS using `alter_table`
/// - Releases the lock after the operation completes
///
/// # Returns
/// A `Result` wrapping the updated `Table` object with new metadata.
///
/// # Errors
/// This function may return an error in several scenarios:
/// - Failure to validate the namespace or table identifier
/// - Inability to acquire a lock (traditional locking mode)
/// - Failure to retrieve the table from HMS
/// - Errors reading or writing table metadata
/// - HMS rejects the update due to concurrent modification (optimistic locking)
/// - Errors from the underlying Thrift communication with HMS
async fn update_table(&self, commit: TableCommit) -> Result<Table> {
let ident = commit.identifier().clone();
let db_name = validate_namespace(ident.namespace())?;
let tbl_name = ident.name.clone();

if self.config.props.contains_key(HMS_HIVE_LOCKS_DISABLED) {
// Optimistic locking path: read first, then validate with EnvironmentContext
let hive_table = self
.client
.0
.get_table(db_name.clone().into(), tbl_name.clone().into())
.await
.map(from_thrift_exception)
.map_err(from_thrift_error)??;

let metadata_location = get_metadata_location(&hive_table.parameters)?;
let env_context = Self::build_environment_context(&metadata_location);

let (staged_table, new_hive_table) = self
.apply_and_prepare_update(commit, &db_name, &tbl_name, &hive_table)
.await?;

self.client
.0
.alter_table_with_environment_context(
db_name.into(),
tbl_name.into(),
new_hive_table,
env_context,
)
.await
.map_err(from_thrift_error)?;

Ok(staged_table)
} else {
// Traditional locking path: acquire lock first, then read
let _guard = HmsLockGuard::acquire(&self.client.0, &db_name, &tbl_name).await?;

let hive_table = self
.client
.0
.get_table(db_name.clone().into(), tbl_name.clone().into())
.await
.map(from_thrift_exception)
.map_err(from_thrift_error)??;

let (staged_table, new_hive_table) = self
.apply_and_prepare_update(commit, &db_name, &tbl_name, &hive_table)
.await?;

self.client
.0
.alter_table(db_name.into(), tbl_name.into(), new_hive_table)
.await
.map_err(from_thrift_error)?;

Ok(staged_table)
// Lock automatically released here via Drop
}
}
}
69 changes: 69 additions & 0 deletions crates/catalog/hms/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use std::collections::HashMap;
use chrono::Utc;
use hive_metastore::{Database, PrincipalType, SerDeInfo, StorageDescriptor};
use iceberg::spec::Schema;
use iceberg::table::Table;
use iceberg::{Error, ErrorKind, Namespace, NamespaceIdent, Result};
use pilota::{AHashMap, FastStr};

Expand Down Expand Up @@ -155,6 +156,54 @@ pub(crate) fn convert_to_database(
Ok(db)
}

pub(crate) fn update_hive_table_from_table(
hive_tbl: &hive_metastore::Table,
tbl: &Table,
) -> Result<hive_metastore::Table> {
let mut new_tbl = hive_tbl.clone();
let metadata = tbl.metadata();
let schema = metadata.current_schema();

let hive_schema = HiveSchemaBuilder::from_iceberg(schema)?.build();

match new_tbl.sd.as_mut() {
Some(sd) => {
sd.cols = Some(hive_schema);
}
None => {
// Highly unlikely for a real HMS table, but be defensive
new_tbl.sd = Some(StorageDescriptor {
cols: Some(hive_schema),
..Default::default()
});
}
}

let metadata_location = tbl.metadata_location_result()?.to_string();

let mut params: AHashMap<FastStr, FastStr> = new_tbl.parameters.take().unwrap_or_default();
for (k, v) in metadata.properties().iter() {
if k == METADATA_LOCATION || k == TABLE_TYPE || k == EXTERNAL {
continue;
}
params.insert(
FastStr::from_string(k.to_string()),
FastStr::from_string(v.to_string()),
);
}

params.insert(FastStr::from(EXTERNAL), FastStr::from("TRUE"));
params.insert(FastStr::from(TABLE_TYPE), FastStr::from("ICEBERG"));
params.insert(
FastStr::from(METADATA_LOCATION),
FastStr::from(metadata_location),
);

new_tbl.parameters = Some(params);

Ok(new_tbl)
}

pub(crate) fn convert_to_hive_table(
db_name: String,
schema: &Schema,
Expand Down Expand Up @@ -309,6 +358,26 @@ fn get_current_time() -> Result<i32> {
})
}

pub(crate) fn create_lock_request(db_name: &str, tbl_name: &str) -> hive_metastore::LockRequest {
let component = hive_metastore::LockComponent {
r#type: hive_metastore::LockType::EXCLUSIVE,
level: hive_metastore::LockLevel::TABLE,
dbname: FastStr::from_string(db_name.to_string()),
tablename: Some(FastStr::from_string(tbl_name.to_string())),
partitionname: None,
operation_type: None,
is_acid: Some(true),
is_dynamic_partition_write: None,
};
hive_metastore::LockRequest {
component: vec![component],
txnid: None,
user: FastStr::from(whoami::username()),
hostname: FastStr::from(whoami::fallible::hostname().unwrap()),
agent_info: None,
}
}

#[cfg(test)]
mod tests {
use iceberg::spec::{NestedField, PrimitiveType, Type};
Expand Down
Loading
Loading