Merge branch 'hive-lock-refactor'

Author: jcellary

pyiceberg/catalog/hive.py

@@ -136,6 +136,9 @@
 DO_NOT_UPDATE_STATS = "DO_NOT_UPDATE_STATS"
 DO_NOT_UPDATE_STATS_DEFAULT = "true"
 
+NO_LOCK_EXPECTED_KEY = "expected_parameter_key"
+NO_LOCK_EXPECTED_VALUE = "expected_parameter_value"
+
 logger = logging.getLogger(__name__)
 
 
@@ -499,6 +502,66 @@ def _do_wait_for_lock() -> LockResponse:
 
         return _do_wait_for_lock()
 
+    @staticmethod
+    def _hive_lock_enabled(table_properties: Properties, catalog_properties: Properties) -> bool:
+        """Determine whether HMS locking is enabled for a commit.
+
+        Matches the Java implementation in HiveTableOperations: checks the table property first,
+        then falls back to catalog properties, then defaults to True.
+        """
+        if TableProperties.HIVE_LOCK_ENABLED in table_properties:
+            return property_as_bool(
+                table_properties, TableProperties.HIVE_LOCK_ENABLED, TableProperties.HIVE_LOCK_ENABLED_DEFAULT
+            )
+        return property_as_bool(catalog_properties, TableProperties.HIVE_LOCK_ENABLED, TableProperties.HIVE_LOCK_ENABLED_DEFAULT)
+
+    def commit_table(
+        self, table: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
+    ) -> CommitTableResponse:
+        """Commit updates to a table.
+
+        Args:
+            table (Table): The table to be updated.
+            requirements: (Tuple[TableRequirement, ...]): Table requirements.
+            updates: (Tuple[TableUpdate, ...]): Table updates.
+
+        Returns:
+            CommitTableResponse: The updated metadata.
+
+        Raises:
+            NoSuchTableError: If a table with the given identifier does not exist.
+            CommitFailedException: Requirement not met, or a conflict with a concurrent commit.
+        """
+        table_identifier = table.name()
+        database_name, table_name = self.identifier_to_database_and_table(table_identifier, NoSuchTableError)
+        lock_enabled = self._hive_lock_enabled(table.properties, self.properties)
+        # commit to hive
+        # https://github.com/apache/hive/blob/master/standalone-metastore/metastore-common/src/main/thrift/hive_metastore.thrift#L1232
+        with self._client as open_client:
+            if lock_enabled:
+                lock: LockResponse = open_client.lock(self._create_lock_request(database_name, table_name))
+
+                try:
+                    if lock.state != LockState.ACQUIRED:
+                        if lock.state == LockState.WAITING:
+                            self._wait_for_lock(database_name, table_name, lock.lockid, open_client)
+                        else:
+                            raise CommitFailedException(f"Failed to acquire lock for {table_identifier}, state: {lock.state}")
+
+                    return self._do_commit(
+                        open_client, table_identifier, database_name, table_name, requirements, updates,
+                        lock_enabled=True,
+                    )
+                except WaitingForLockException as e:
+                    raise CommitFailedException(f"Failed to acquire lock for {table_identifier}, state: {lock.state}") from e
+                finally:
+                    open_client.unlock(UnlockRequest(lockid=lock.lockid))
+            else:
+                return self._do_commit(
+                    open_client, table_identifier, database_name, table_name, requirements, updates,
+                    lock_enabled=False,
+                )
+
     def _do_commit(
         self,
         open_client: Client,
@@ -507,10 +570,13 @@ def _do_commit(
         table_name: str,
         requirements: tuple[TableRequirement, ...],
         updates: tuple[TableUpdate, ...],
+        lock_enabled: bool = True,
     ) -> CommitTableResponse:
         """Perform the actual commit logic (get table, update, write metadata, alter/create in HMS).
 
         This method contains the core commit logic, separated from locking concerns.
+        When lock_enabled is False, an optimistic concurrency check via the HMS EnvironmentContext
+        is used instead (requires HIVE-26882 on the server).
         """
         hive_table: HiveTable | None
         current_table: Table | None
@@ -566,11 +632,16 @@ def _do_commit(
                 updated_staged_table.location(),
                 property_as_bool(self.properties, HIVE2_COMPATIBLE, HIVE2_COMPATIBLE_DEFAULT),
             )
+            env_context_properties: dict[str, str] = {DO_NOT_UPDATE_STATS: DO_NOT_UPDATE_STATS_DEFAULT}
+            if not lock_enabled:
+                env_context_properties[NO_LOCK_EXPECTED_KEY] = PROP_METADATA_LOCATION
+                env_context_properties[NO_LOCK_EXPECTED_VALUE] = current_table.metadata_location
+
             open_client.alter_table_with_environment_context(
                 dbname=database_name,
                 tbl_name=table_name,
                 new_tbl=hive_table,
-                environment_context=EnvironmentContext(properties={DO_NOT_UPDATE_STATS: DO_NOT_UPDATE_STATS_DEFAULT}),
+                environment_context=EnvironmentContext(properties=env_context_properties),
             )
         else:
             # Table does not exist, create it.
@@ -589,60 +660,6 @@ def _do_commit(
             metadata=updated_staged_table.metadata, metadata_location=updated_staged_table.metadata_location
         )
 
-    @staticmethod
-    def _hive_lock_enabled(table_properties: Properties, catalog_properties: Properties) -> bool:
-        """Determine whether HMS locking is enabled for a commit.
-
-        Matches the Java implementation in HiveTableOperations: checks the table property first,
-        then falls back to catalog properties, then defaults to True.
-        """
-        if TableProperties.HIVE_LOCK_ENABLED in table_properties:
-            return property_as_bool(
-                table_properties, TableProperties.HIVE_LOCK_ENABLED, TableProperties.HIVE_LOCK_ENABLED_DEFAULT
-            )
-        return property_as_bool(catalog_properties, TableProperties.HIVE_LOCK_ENABLED, TableProperties.HIVE_LOCK_ENABLED_DEFAULT)
-
-    def commit_table(
-        self, table: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
-    ) -> CommitTableResponse:
-        """Commit updates to a table.
-
-        Args:
-            table (Table): The table to be updated.
-            requirements: (Tuple[TableRequirement, ...]): Table requirements.
-            updates: (Tuple[TableUpdate, ...]): Table updates.
-
-        Returns:
-            CommitTableResponse: The updated metadata.
-
-        Raises:
-            NoSuchTableError: If a table with the given identifier does not exist.
-            CommitFailedException: Requirement not met, or a conflict with a concurrent commit.
-        """
-        table_identifier = table.name()
-        database_name, table_name = self.identifier_to_database_and_table(table_identifier, NoSuchTableError)
-        lock_enabled = self._hive_lock_enabled(table.properties, self.properties)
-        # commit to hive
-        # https://github.com/apache/hive/blob/master/standalone-metastore/metastore-common/src/main/thrift/hive_metastore.thrift#L1232
-        with self._client as open_client:
-            if lock_enabled:
-                lock: LockResponse = open_client.lock(self._create_lock_request(database_name, table_name))
-
-                try:
-                    if lock.state != LockState.ACQUIRED:
-                        if lock.state == LockState.WAITING:
-                            self._wait_for_lock(database_name, table_name, lock.lockid, open_client)
-                        else:
-                            raise CommitFailedException(f"Failed to acquire lock for {table_identifier}, state: {lock.state}")
-
-                    return self._do_commit(open_client, table_identifier, database_name, table_name, requirements, updates)
-                except WaitingForLockException as e:
-                    raise CommitFailedException(f"Failed to acquire lock for {table_identifier}, state: {lock.state}") from e
-                finally:
-                    open_client.unlock(UnlockRequest(lockid=lock.lockid))
-            else:
-                return self._do_commit(open_client, table_identifier, database_name, table_name, requirements, updates)
-
     def load_table(self, identifier: str | Identifier) -> Table:
         """Load the table's metadata and return the table instance.
 

tests/catalog/test_hive.py

@@ -51,6 +51,9 @@
     LOCK_CHECK_MAX_WAIT_TIME,
     LOCK_CHECK_MIN_WAIT_TIME,
     LOCK_CHECK_RETRIES,
+    NO_LOCK_EXPECTED_KEY,
+    NO_LOCK_EXPECTED_VALUE,
+    PROP_METADATA_LOCATION,
     HiveCatalog,
     _construct_hive_storage_descriptor,
     _HiveClient,
@@ -1455,6 +1458,8 @@ def test_commit_table_skips_locking_when_table_property_disables_it() -> None:
         catalog.commit_table(mock_table, requirements=(), updates=())
 
     mock_do_commit.assert_called_once()
+    _, kwargs = mock_do_commit.call_args
+    assert kwargs["lock_enabled"] is False
     catalog._client.__enter__().lock.assert_not_called()
     catalog._client.__enter__().check_lock.assert_not_called()
     catalog._client.__enter__().unlock.assert_not_called()
@@ -1507,3 +1512,94 @@ def test_commit_table_uses_locking_by_default() -> None:
     mock_client.lock.assert_called_once()
     mock_client.unlock.assert_called_once()
     mock_do_commit.assert_called_once()
+    _, kwargs = mock_do_commit.call_args
+    assert kwargs["lock_enabled"] is True
+
+
+def test_do_commit_env_context_includes_expected_params_when_lock_disabled() -> None:
+    """When lock_enabled=False, alter_table_with_environment_context must include
+    expected_parameter_key and expected_parameter_value for optimistic concurrency."""
+    prop = {"uri": HIVE_METASTORE_FAKE_URL}
+    catalog = HiveCatalog(HIVE_CATALOG_NAME, **prop)
+
+    current_metadata_location = "s3://bucket/db/table/metadata/v1.metadata.json"
+
+    mock_client = MagicMock()
+    mock_hive_table = MagicMock()
+    mock_hive_table.parameters = {
+        "table_type": "ICEBERG",
+        "metadata_location": current_metadata_location,
+    }
+    mock_client.get_table.return_value = mock_hive_table
+
+    with (
+        patch.object(catalog, "_convert_hive_into_iceberg") as mock_convert,
+        patch.object(catalog, "_update_and_stage_table") as mock_stage,
+        patch.object(catalog, "_write_metadata"),
+        patch.object(catalog, "_convert_iceberg_into_hive"),
+        patch("pyiceberg.catalog.hive.CommitTableResponse"),
+    ):
+        mock_current_table = MagicMock()
+        mock_current_table.metadata_location = current_metadata_location
+        mock_current_table.metadata = MagicMock()
+        mock_current_table.properties = {}
+        mock_convert.return_value = mock_current_table
+
+        mock_staged = MagicMock()
+        mock_staged.metadata = MagicMock()
+        mock_staged.properties = {}
+        mock_stage.return_value = mock_staged
+
+        catalog._do_commit(
+            mock_client, ("default", "my_table"), "default", "my_table",
+            requirements=(), updates=(), lock_enabled=False,
+        )
+
+    mock_client.alter_table_with_environment_context.assert_called_once()
+    env_ctx = mock_client.alter_table_with_environment_context.call_args[1]["environment_context"]
+    assert env_ctx.properties[NO_LOCK_EXPECTED_KEY] == PROP_METADATA_LOCATION
+    assert env_ctx.properties[NO_LOCK_EXPECTED_VALUE] == current_metadata_location
+    assert env_ctx.properties[DO_NOT_UPDATE_STATS] == DO_NOT_UPDATE_STATS_DEFAULT
+
+
+def test_do_commit_env_context_excludes_expected_params_when_lock_enabled() -> None:
+    """When lock_enabled=True (default), alter_table_with_environment_context must NOT include
+    expected_parameter_key or expected_parameter_value."""
+    prop = {"uri": HIVE_METASTORE_FAKE_URL}
+    catalog = HiveCatalog(HIVE_CATALOG_NAME, **prop)
+
+    mock_client = MagicMock()
+    mock_hive_table = MagicMock()
+    mock_hive_table.parameters = {
+        "table_type": "ICEBERG",
+        "metadata_location": "s3://bucket/db/table/metadata/v1.metadata.json",
+    }
+    mock_client.get_table.return_value = mock_hive_table
+
+    with (
+        patch.object(catalog, "_convert_hive_into_iceberg") as mock_convert,
+        patch.object(catalog, "_update_and_stage_table") as mock_stage,
+        patch.object(catalog, "_write_metadata"),
+        patch.object(catalog, "_convert_iceberg_into_hive"),
+        patch("pyiceberg.catalog.hive.CommitTableResponse"),
+    ):
+        mock_current_table = MagicMock()
+        mock_current_table.metadata = MagicMock()
+        mock_current_table.properties = {}
+        mock_convert.return_value = mock_current_table
+
+        mock_staged = MagicMock()
+        mock_staged.metadata = MagicMock()
+        mock_staged.properties = {}
+        mock_stage.return_value = mock_staged
+
+        catalog._do_commit(
+            mock_client, ("default", "my_table"), "default", "my_table",
+            requirements=(), updates=(), lock_enabled=True,
+        )
+
+    mock_client.alter_table_with_environment_context.assert_called_once()
+    env_ctx = mock_client.alter_table_with_environment_context.call_args[1]["environment_context"]
+    assert NO_LOCK_EXPECTED_KEY not in env_ctx.properties
+    assert NO_LOCK_EXPECTED_VALUE not in env_ctx.properties
+    assert env_ctx.properties[DO_NOT_UPDATE_STATS] == DO_NOT_UPDATE_STATS_DEFAULT