rebase and test fixes

Author: rambleraptor

pyiceberg/table/update/__init__.py

@@ -617,6 +617,21 @@ def _(update: AddEncryptedKeyUpdate, base_metadata: TableMetadata, context: _Tab
 
     return base_metadata.model_copy(update={"encryption_keys": base_metadata.encryption_keys + [update.key]})
 
+
+@_apply_table_update.register(RemoveEncryptedKeyUpdate)
+def _(update: RemoveEncryptedKeyUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
+    context.add_update(update)
+
+    if base_metadata.format_version <= 2:
+        raise ValueError("Cannot add encryption keys to Iceberg v1 or v2 tables")
+
+    encryption_keys = [key for key in base_metadata.encryption_keys if key.key_id != update.key_id]
+    if len(encryption_keys) == len(base_metadata.encryption_keys):
+        raise ValueError(f"Encryption key {update.key_id} not found")
+
+    return base_metadata.model_copy(update={"encryption_keys": encryption_keys})
+
+
 @_apply_table_update.register(RemoveSchemasUpdate)
 def _(update: RemoveSchemasUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
     # This method should error if any schemas do not exist.

tests/table/test_init.py

@@ -79,8 +79,8 @@
     AssertLastAssignedPartitionId,
     AssertRefSnapshotId,
     AssertTableUUID,
-    RemovePartitionStatisticsUpdate,
     RemoveEncryptedKeyUpdate,
+    RemovePartitionStatisticsUpdate,
     RemovePropertiesUpdate,
     RemoveSchemasUpdate,
     RemoveSnapshotRefUpdate,
@@ -1423,9 +1423,7 @@ def test_set_partition_statistics_update(table_v2_with_statistics: Table) -> Non
     new_metadata = update_table_metadata(
         table_v2_with_statistics.metadata,
         (update,),
-
-def test_add_encryption_key(table_v3: Table) -> None:
-    update = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test", encrypted_key_metadata=base64.b64encode(b"hello")))
+    )
 
     expected = """
     {
@@ -1483,6 +1481,13 @@ def test_remove_partition_statistics_update_with_invalid_snapshot_id(table_v2_wi
             table_v2_with_statistics.metadata,
             (RemovePartitionStatisticsUpdate(snapshot_id=123456789),),
         )
+
+
+def test_add_encryption_key(table_v3: Table) -> None:
+    update = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test", encrypted_key_metadata=base64.b64encode(b"hello")))
+
+    expected = """
+    {
       "key-id": "test",
       "encrypted-key-metadata": "aGVsbG8="
     }"""
@@ -1510,11 +1515,11 @@ def test_remove_non_existent_encryption_key(table_v3: Table) -> None:
     assert len(add_metadata.encryption_keys) == 1
 
     update_remove = RemoveEncryptedKeyUpdate(key_id="non_existent_key")
-    remove_metadata = update_table_metadata(add_metadata, (update_remove,))
-    assert len(remove_metadata.encryption_keys) == 1  # Should be a no-op
+    with pytest.raises(ValueError, match=r"Encryption key non_existent_key not found"):
+        update_table_metadata(add_metadata, (update_remove,))
 
 
 def test_add_remove_encryption_key_v2_table(table_v2: Table) -> None:
     update_add = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test_v2", encrypted_key_metadata=base64.b64encode(b"hello_v2")))
-    with pytest.raises(ValueError, match=r"Cannot add encryption keys from Iceberg v1 or v2 table"):
+    with pytest.raises(ValueError, match=r"Cannot add encryption keys to Iceberg v1 or v2 table"):
         update_table_metadata(table_v2.metadata, (update_add,))

tests/table/test_snapshots.py

@@ -139,15 +139,15 @@ def test_deserialize_snapshot_with_properties(snapshot_with_properties: Snapshot
 def test_snapshot_repr(snapshot: Snapshot) -> None:
     assert (
         repr(snapshot)
-        == """Snapshot(snapshot_id=25, parent_snapshot_id=19, sequence_number=200, timestamp_ms=1602638573590, manifest_list='s3:/a/b/c.avro', summary=Summary(Operation.APPEND), schema_id=3)"""
+        == """Snapshot(snapshot_id=25, parent_snapshot_id=19, sequence_number=200, timestamp_ms=1602638573590, manifest_list='s3:/a/b/c.avro', summary=Summary(Operation.APPEND), schema_id=3, key_id=None)"""
     )
     assert snapshot == eval(repr(snapshot))
 
 
 def test_snapshot_with_properties_repr(snapshot_with_properties: Snapshot) -> None:
     assert (
         repr(snapshot_with_properties)
-        == """Snapshot(snapshot_id=25, parent_snapshot_id=19, sequence_number=200, timestamp_ms=1602638573590, manifest_list='s3:/a/b/c.avro', summary=Summary(Operation.APPEND, **{'foo': 'bar'}), schema_id=3)"""
+        == """Snapshot(snapshot_id=25, parent_snapshot_id=19, sequence_number=200, timestamp_ms=1602638573590, manifest_list='s3:/a/b/c.avro', summary=Summary(Operation.APPEND, **{'foo': 'bar'}), schema_id=3, key_id=None)"""
     )
     assert snapshot_with_properties == eval(repr(snapshot_with_properties))