Merge pull request #1 from vladson/hms_update_table_refactoring

Hms update table refactoring

Author: vladson

crates/catalog/hms/src/catalog.rs

@@ -57,7 +57,7 @@ pub const HMS_HIVE_LOCKS_DISABLED: &str = "hive_locks_disabled";
 const HMS_EXPECTED_PARAMETER_KEY: &str = "expected_parameter_key";
 const HMS_EXPECTED_PARAMETER_VALUE: &str = "expected_parameter_value";
 
-/// Builder for [`RestCatalog`].
+/// Builder for [`HmsCatalog`].
 #[derive(Debug)]
 pub struct HmsCatalogBuilder(HmsCatalogConfig);
 
@@ -174,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> {
@@ -215,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]
@@ -610,25 +705,74 @@ impl Catalog for HmsCatalog {
         ))
     }
 
+    /// 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 HMS_HIVE_LOCKS_DISABLED is set
-        if let Some(tt) = &self.config.props.get(HMS_HIVE_LOCKS_DISABLED) {
-            // Do alter table with env context
-            Err(Error::new(ErrorKind::Unexpected, "Optimistic locks are not supported yet"))
-        } else {
-            // start with trying to acquire a lock
-            let lock = &self
+        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
-                .lock(create_lock_request(&db_name, &tbl_name))
+                .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
@@ -637,41 +781,18 @@ impl Catalog for HmsCatalog {
                 .map(from_thrift_exception)
                 .map_err(from_thrift_error)??;
 
-            let metadata_location = get_metadata_location(&hive_table.parameters)?;
-
-            let metadata = TableMetadata::read_from(&self.file_io, &metadata_location).await?;
+            let (staged_table, new_hive_table) =
+                self.apply_and_prepare_update(commit, &db_name, &tbl_name, &hive_table)
+                    .await?;
 
-            let cur_table = Table::builder()
-                .file_io(self.file_io())
-                .metadata_location(metadata_location)
-                .metadata(metadata)
-                .identifier(TableIdent::new(
-                    NamespaceIdent::new(db_name.clone()),
-                    tbl_name.clone(),
-                ))
-                .build()?;
-
-            let staged_table = commit.apply(cur_table)?;
-            staged_table
-                .metadata()
-                .write_to(
-                    staged_table.file_io(),
-                    staged_table.metadata().metadata_location(),
-                )
-                .await?;
-            let new_hive_table = update_hive_table_from_table(hive_table, &staged_table)?;
-
-            let updated = self.client.0.alter_table(
-                db_name.clone().into(),
-                tbl_name.clone().into(),
-                new_hive_table,
-            ).await
-            .map(from_thrift_exception)
-            .map_err(from_thrift_error)??;
+            self.client
+                .0
+                .alter_table(db_name.into(), tbl_name.into(), new_hive_table)
+                .await
+                .map_err(from_thrift_error)?;
 
-            // unlock the table after alter table
-            &self.client.0.unlock(lock.lockid).await.map(from_thrift_error)?;
             Ok(staged_table)
+            // Lock automatically released here via Drop
         }
     }
 }

crates/catalog/hms/src/utils.rs

@@ -187,7 +187,7 @@ pub(crate) fn update_hive_table_from_table(
         if k == METADATA_LOCATION || k == TABLE_TYPE || k == EXTERNAL {
             continue;
         }
-        params.insert(k.as_str().clone().into(), v.as_str().clone().into());
+        params.insert(FastStr::from_string(k.to_string()), FastStr::from_string(v.to_string()));
     }
 
     params.insert(FastStr::from(EXTERNAL), FastStr::from("TRUE"));
@@ -357,14 +357,14 @@ fn get_current_time() -> Result<i32> {
 }
 
 pub(crate) fn create_lock_request(
-    tbl_name: &String,
-    db_name: &String,
+    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: db_name.clone().into(),
-        tablename: Some(tbl_name.clone().into()),
+        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),
@@ -374,7 +374,7 @@ pub(crate) fn create_lock_request(
         component: vec![component],
         txnid: None,
         user: FastStr::from(whoami::username()),
-        hostname: FastStr::from(whoami::fallible::hostname()),
+        hostname: FastStr::from(whoami::fallible::hostname().unwrap()),
         agent_info: None,
     }
 }