diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py index 64838b0bd6..aaa4a51b7d 100644 --- a/pyiceberg/table/update/__init__.py +++ b/pyiceberg/table/update/__init__.py @@ -29,7 +29,13 @@ 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 @@ class RemovePartitionStatisticsUpdate(IcebergBaseModel): | 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 94614c9d60..5ee5e4566a 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -22,7 +22,7 @@ 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 @@ 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.statistics import BlobMetadata, PartitionStatisticsFile, StatisticsFile from pyiceberg.table.update import ( + AddEncryptionKeyUpdate, AddPartitionSpecUpdate, AddSnapshotUpdate, AddSortOrderUpdate, @@ -74,6 +75,7 @@ AssertLastAssignedPartitionId, AssertRefSnapshotId, AssertTableUUID, + RemoveEncryptionKeyUpdate, RemovePartitionSpecsUpdate, RemovePartitionStatisticsUpdate, RemovePropertiesUpdate, @@ -86,6 +88,7 @@ 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() + + +@pytest.mark.parametrize( + "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,