This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 35bb9bb62 feat(encryption): [7/N] Add encryption key updates (#4148)
35bb9bb62 is described below
commit 35bb9bb62a00dcf8560d970403fef1a8047809e8
Author: Xander <[email protected]>
AuthorDate: Sat Oct 24 17:31:06 2026 +0000
feat(encryption): [7/N] Add encryption key updates (#4148)
feat(encryption): Add encryption key updates
---
pyiceberg/table/update/__init__.py | 47 ++++++++++++++++++++++-
tests/table/test_init.py | 76 +++++++++++++++++++++++++++++++++++++-
2 files changed, 119 insertions(+), 4 deletions(-)
diff --git a/pyiceberg/table/update/__init__.py
b/pyiceberg/table/update/__init__.py
index 64838b0bd..aaa4a51b7 100644
--- a/pyiceberg/table/update/__init__.py
+++ b/pyiceberg/table/update/__init__.py
@@ -29,7 +29,13 @@ from pydantic import Field, field_validator,
model_serializer, model_validator
from pyiceberg.exceptions import CommitFailedException
from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec
from pyiceberg.schema import Schema
-from pyiceberg.table.metadata import SUPPORTED_TABLE_FORMAT_VERSION,
TableMetadata, TableMetadataUtil
+from pyiceberg.table.metadata import (
+ SUPPORTED_TABLE_FORMAT_VERSION,
+ EncryptedKey,
+ TableMetadata,
+ TableMetadataUtil,
+ TableMetadataV3,
+)
from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
from pyiceberg.table.snapshots import (
MetadataLogEntry,
@@ -222,6 +228,16 @@ class RemovePartitionStatisticsUpdate(IcebergBaseModel):
snapshot_id: int = Field(alias="snapshot-id")
+class AddEncryptionKeyUpdate(IcebergBaseModel):
+ action: Literal["add-encryption-key"] = Field(default="add-encryption-key")
+ encryption_key: EncryptedKey = Field(alias="encryption-key")
+
+
+class RemoveEncryptionKeyUpdate(IcebergBaseModel):
+ action: Literal["remove-encryption-key"] =
Field(default="remove-encryption-key")
+ key_id: str = Field(alias="key-id")
+
+
TableUpdate = Annotated[
AssignUUIDUpdate
| UpgradeFormatVersionUpdate
@@ -243,7 +259,9 @@ TableUpdate = Annotated[
| RemovePartitionSpecsUpdate
| RemoveSchemasUpdate
| SetPartitionStatisticsUpdate
- | RemovePartitionStatisticsUpdate,
+ | RemovePartitionStatisticsUpdate
+ | AddEncryptionKeyUpdate
+ | RemoveEncryptionKeyUpdate,
Field(discriminator="action"),
]
@@ -680,6 +698,31 @@ def _(
return base_metadata.model_copy(update={"partition_statistics":
statistics})
+@_apply_table_update.register(AddEncryptionKeyUpdate)
+def _(update: AddEncryptionKeyUpdate, base_metadata: TableMetadata, context:
_TableMetadataUpdateContext) -> TableMetadata:
+ if not isinstance(base_metadata, TableMetadataV3):
+ raise ValueError(f"Encryption keys are not supported in format version
{base_metadata.format_version}")
+
+ if any(key.key_id == update.encryption_key.key_id for key in
base_metadata.encryption_keys):
+ return base_metadata
+
+ context.add_update(update)
+ return base_metadata.model_copy(update={"encryption_keys":
base_metadata.encryption_keys + [update.encryption_key]})
+
+
+@_apply_table_update.register(RemoveEncryptionKeyUpdate)
+def _(update: RemoveEncryptionKeyUpdate, base_metadata: TableMetadata,
context: _TableMetadataUpdateContext) -> TableMetadata:
+ if not isinstance(base_metadata, TableMetadataV3):
+ raise ValueError(f"Encryption keys are not supported in format version
{base_metadata.format_version}")
+
+ 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):
+ return base_metadata
+
+ context.add_update(update)
+ return base_metadata.model_copy(update={"encryption_keys":
encryption_keys})
+
+
def update_table_metadata(
base_metadata: TableMetadata,
updates: tuple[TableUpdate, ...],
diff --git a/tests/table/test_init.py b/tests/table/test_init.py
index 94614c9d6..5ee5e4566 100644
--- a/tests/table/test_init.py
+++ b/tests/table/test_init.py
@@ -22,7 +22,7 @@ from pathlib import Path
from typing import Any
import pytest
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
from pytest_lazy_fixtures import lf
from pyiceberg.catalog import Catalog
@@ -46,7 +46,7 @@ from pyiceberg.table import (
Table,
TableIdentifier,
)
-from pyiceberg.table.metadata import TableMetadataUtil, TableMetadataV1,
TableMetadataV2, _generate_snapshot_id
+from pyiceberg.table.metadata import EncryptedKey, TableMetadataUtil,
TableMetadataV1, TableMetadataV2, _generate_snapshot_id
from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
from pyiceberg.table.snapshots import (
MetadataLogEntry,
@@ -63,6 +63,7 @@ from pyiceberg.table.sorting import (
)
from pyiceberg.table.statistics import BlobMetadata, PartitionStatisticsFile,
StatisticsFile
from pyiceberg.table.update import (
+ AddEncryptionKeyUpdate,
AddPartitionSpecUpdate,
AddSnapshotUpdate,
AddSortOrderUpdate,
@@ -74,6 +75,7 @@ from pyiceberg.table.update import (
AssertLastAssignedPartitionId,
AssertRefSnapshotId,
AssertTableUUID,
+ RemoveEncryptionKeyUpdate,
RemovePartitionSpecsUpdate,
RemovePartitionStatisticsUpdate,
RemovePropertiesUpdate,
@@ -86,6 +88,7 @@ from pyiceberg.table.update import (
SetPropertiesUpdate,
SetSnapshotRefUpdate,
SetStatisticsUpdate,
+ TableUpdate,
_apply_table_update,
_TableMetadataUpdateContext,
update_table_metadata,
@@ -1951,6 +1954,75 @@ def
test_remove_partition_statistics_update_with_invalid_snapshot_id(table_v2_wi
)
+def test_add_encryption_key_update(table_v3: Table) -> None:
+ kek = EncryptedKey(key_id="kek", encrypted_key_metadata=b"wrapped-kek",
encrypted_by_id="master", properties={"a": "b"})
+ manifest_list_key = EncryptedKey(key_id="manifest-list-key",
encrypted_key_metadata=b"wrapped-dek", encrypted_by_id="kek")
+
+ new_metadata = update_table_metadata(
+ table_v3.metadata,
+ (AddEncryptionKeyUpdate(encryption_key=kek),
AddEncryptionKeyUpdate(encryption_key=manifest_list_key)),
+ )
+
+ assert new_metadata.encryption_keys == [kek, manifest_list_key]
+
+
+def test_add_encryption_key_update_with_existing_key_id(table_v3: Table) ->
None:
+ key = EncryptedKey(key_id="kek", encrypted_key_metadata=b"wrapped-kek")
+ base_metadata = update_table_metadata(table_v3.metadata,
(AddEncryptionKeyUpdate(encryption_key=key),))
+
+ context = _TableMetadataUpdateContext()
+ duplicate = EncryptedKey(key_id="kek", encrypted_key_metadata=b"other")
+ new_metadata =
_apply_table_update(AddEncryptionKeyUpdate(encryption_key=duplicate),
base_metadata, context)
+
+ assert new_metadata.encryption_keys == [key]
+ assert not context.has_changes()
+
+
+def test_remove_encryption_key_update(table_v3: Table) -> None:
+ first = EncryptedKey(key_id="first", encrypted_key_metadata=b"first")
+ second = EncryptedKey(key_id="second", encrypted_key_metadata=b"second")
+ base_metadata = update_table_metadata(
+ table_v3.metadata, (AddEncryptionKeyUpdate(encryption_key=first),
AddEncryptionKeyUpdate(encryption_key=second))
+ )
+
+ new_metadata = update_table_metadata(base_metadata,
(RemoveEncryptionKeyUpdate(key_id="first"),))
+
+ assert new_metadata.encryption_keys == [second]
+
+
+def test_remove_encryption_key_update_with_missing_key_id(table_v3: Table) ->
None:
+ context = _TableMetadataUpdateContext()
+ new_metadata =
_apply_table_update(RemoveEncryptionKeyUpdate(key_id="missing"),
table_v3.metadata, context)
+
+ assert new_metadata == table_v3.metadata
+ assert not context.has_changes()
+
+
[email protected](
+ "update",
+ [
+ AddEncryptionKeyUpdate(encryption_key=EncryptedKey(key_id="a",
encrypted_key_metadata=b"key")),
+ RemoveEncryptionKeyUpdate(key_id="a"),
+ ],
+)
+def test_encryption_key_updates_require_v3(table_v2: Table, update:
TableUpdate) -> None:
+ with pytest.raises(ValueError, match="Encryption keys are not supported in
format version 2"):
+ update_table_metadata(table_v2.metadata, (update,))
+
+
+def test_serialize_encryption_key_updates() -> None:
+ add = AddEncryptionKeyUpdate(encryption_key=EncryptedKey(key_id="a",
encrypted_key_metadata=b"key", encrypted_by_id="b"))
+ remove = RemoveEncryptionKeyUpdate(key_id="a")
+
+ assert json.loads(add.model_dump_json()) == {
+ "action": "add-encryption-key",
+ "encryption-key": {"key-id": "a", "encrypted-key-metadata": "a2V5",
"encrypted-by-id": "b", "properties": {}},
+ }
+ assert json.loads(remove.model_dump_json()) == {"action":
"remove-encryption-key", "key-id": "a"}
+ assert TypeAdapter(TableUpdate).validate_json(add.model_dump_json()) == add
+ assert TypeAdapter(TableUpdate).validate_json(remove.model_dump_json()) ==
remove
+
+
def test_add_snapshot_update_fails_without_first_row_id(table_v3: Table) ->
None:
new_snapshot = Snapshot(
snapshot_id=25,
