From 22919d1514c5294b3a80a26d5d85c8a379740dbf Mon Sep 17 00:00:00 2001 From: Xander Date: Fri, 11 Sep 2026 12:44:34 +0100 Subject: [PATCH 1/4] feat(encryption): [1/N] spec model --- pyiceberg/table/__init__.py | 4 ++ pyiceberg/table/metadata.py | 31 +++++++++++++ pyiceberg/table/snapshots.py | 4 ++ tests/table/test_metadata.py | 85 +++++++++++++++++++++++++++++++++++ tests/table/test_snapshots.py | 42 +++++++++++++++++ 5 files changed, 166 insertions(+) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index fca718f5ec..303b3db135 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -198,6 +198,10 @@ class TableProperties: FORMAT_VERSION = "format-version" DEFAULT_FORMAT_VERSION: TableVersion = 2 + ENCRYPTION_KEY_ID = "encryption.key-id" + ENCRYPTION_DATA_KEY_LENGTH = "encryption.data-key-length" + ENCRYPTION_DATA_KEY_LENGTH_DEFAULT = 16 + MANIFEST_TARGET_SIZE_BYTES = "commit.manifest.target-size-bytes" MANIFEST_TARGET_SIZE_BYTES_DEFAULT = 8 * 1024 * 1024 # 8 MB diff --git a/pyiceberg/table/metadata.py b/pyiceberg/table/metadata.py index 84be7b07bf..8af73beeb9 100644 --- a/pyiceberg/table/metadata.py +++ b/pyiceberg/table/metadata.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import base64 import datetime import uuid from collections.abc import Iterable @@ -125,6 +126,33 @@ def construct_refs(table_metadata: TableMetadata) -> TableMetadata: return table_metadata +class EncryptedKey(IcebergBaseModel): + """A key used for table encryption, tracked in v3 metadata under `encryption-keys`. + + https://iceberg.apache.org/spec/#encryption-keys + """ + + key_id: str = Field(alias="key-id") + """ID of the encryption key.""" + + encrypted_key_metadata: bytes = Field(alias="encrypted-key-metadata") + """The encrypted key and metadata, base64 encoded in JSON.""" + + encrypted_by_id: str | None = Field(alias="encrypted-by-id", default=None) + """ID of the key used to encrypt or wrap `encrypted-key-metadata`.""" + + properties: dict[str, str] = Field(default_factory=dict) + """Additional metadata used by the table's encryption scheme.""" + + @field_validator("encrypted_key_metadata", mode="before") + def decode_encrypted_key_metadata(cls, encrypted_key_metadata: Any) -> Any: + return base64.b64decode(encrypted_key_metadata) if isinstance(encrypted_key_metadata, str) else encrypted_key_metadata + + @field_serializer("encrypted_key_metadata") + def serialize_encrypted_key_metadata(self, encrypted_key_metadata: bytes) -> str: + return base64.b64encode(encrypted_key_metadata).decode("utf-8") + + class TableMetadataCommonFields(IcebergBaseModel): """Metadata for an Iceberg table as specified in the Apache Iceberg spec. @@ -584,6 +612,9 @@ def construct_refs(self) -> TableMetadata: next_row_id: int | None = Field(alias="next-row-id", default=None) """A long higher than all assigned row IDs; the next snapshot's `first-row-id`.""" + encryption_keys: list[EncryptedKey] = Field(alias="encryption-keys", default_factory=list) + """An optional list of encryption keys used for table encryption.""" + def model_dump_json(self, exclude_none: bool = True, exclude: Any | None = None, by_alias: bool = True, **kwargs: Any) -> str: raise NotImplementedError("Writing V3 is not yet supported, see: https://github.com/apache/iceberg-python/issues/1551") diff --git a/pyiceberg/table/snapshots.py b/pyiceberg/table/snapshots.py index 5e9e519a01..f862529e83 100644 --- a/pyiceberg/table/snapshots.py +++ b/pyiceberg/table/snapshots.py @@ -259,6 +259,9 @@ class Snapshot(IcebergBaseModel): added_rows: int | None = Field( alias="added-rows", default=None, description="The upper bound of the number of rows with assigned row IDs" ) + key_id: str | None = Field( + alias="key-id", default=None, description="ID of the encryption key that encrypts the manifest list key metadata" + ) def __str__(self) -> str: """Return the string representation of the Snapshot class.""" @@ -280,6 +283,7 @@ def __repr__(self) -> str: f"schema_id={self.schema_id}" if self.schema_id is not None else None, f"first_row_id={self.first_row_id}" if self.first_row_id is not None else None, f"added_rows={self.added_rows}" if self.added_rows is not None else None, + f"key_id='{self.key_id}'" if self.key_id is not None else None, ] filtered_fields = [field for field in fields if field is not None] return f"Snapshot({', '.join(filtered_fields)})" diff --git a/tests/table/test_metadata.py b/tests/table/test_metadata.py index c163c90626..5063c97a4b 100644 --- a/tests/table/test_metadata.py +++ b/tests/table/test_metadata.py @@ -24,12 +24,14 @@ from uuid import UUID import pytest +from pydantic import ValidationError as PydanticValidationError from pyiceberg.exceptions import ValidationError from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema from pyiceberg.serializers import FromByteStream from pyiceberg.table.metadata import ( + EncryptedKey, TableMetadataUtil, TableMetadataV1, TableMetadataV2, @@ -876,3 +878,86 @@ def test_new_table_metadata_format_v2_with_v3_schema_fails(field_type: Primitive location="s3://some_v1_location/", properties={"format-version": "2"}, ) + + +def test_encrypted_key_minimal() -> None: + # Mirrors Java's TestEncryptedKeyParser, where the key metadata is base64 of b"key" + key = EncryptedKey.model_validate_json('{"key-id": "a", "encrypted-key-metadata": "a2V5"}') + + assert key.key_id == "a" + assert key.encrypted_key_metadata == b"key" + assert key.encrypted_by_id is None + assert key.properties == {} + + +def test_encrypted_key_full() -> None: + key = EncryptedKey.model_validate_json( + '{"key-id": "a", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "b", "properties": {"test": "value"}}' + ) + + assert key.key_id == "a" + assert key.encrypted_key_metadata == b"key" + assert key.encrypted_by_id == "b" + assert key.properties == {"test": "value"} + + +def test_encrypted_key_serialize() -> None: + key = EncryptedKey(key_id="a", encrypted_key_metadata=b"key", encrypted_by_id="b", properties={"test": "value"}) + + expected = '{"key-id":"a","encrypted-key-metadata":"a2V5","encrypted-by-id":"b","properties":{"test":"value"}}' + assert key.model_dump_json() == expected + + +def test_encrypted_key_serialize_minimal() -> None: + key = EncryptedKey(key_id="a", encrypted_key_metadata=b"key") + + assert key.model_dump_json() == '{"key-id":"a","encrypted-key-metadata":"a2V5","properties":{}}' + + +@pytest.mark.parametrize( + "payload, missing", + [ + ('{"encrypted-key-metadata": "a2V5"}', "key-id"), + ('{"key-id": "a"}', "encrypted-key-metadata"), + ], +) +def test_encrypted_key_missing_required_field(payload: str, missing: str) -> None: + with pytest.raises(PydanticValidationError) as exc_info: + EncryptedKey.model_validate_json(payload) + + assert missing in str(exc_info.value) + + +def test_v3_metadata_parsing_encryption_keys(example_table_metadata_v3: dict[str, Any]) -> None: + metadata = { + **example_table_metadata_v3, + "encryption-keys": [ + {"key-id": "kek-1", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "master-1", "properties": {"a": "b"}}, + {"key-id": "dek-1", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "kek-1"}, + ], + "snapshots": [ + {**snapshot, "key-id": "dek-1"} if snapshot["snapshot-id"] == 3055729675574597004 else snapshot + for snapshot in example_table_metadata_v3["snapshots"] + ], + } + + table_metadata = TableMetadataUtil.parse_obj(metadata) + + assert isinstance(table_metadata, TableMetadataV3) + assert [key.key_id for key in table_metadata.encryption_keys] == ["kek-1", "dek-1"] + assert table_metadata.encryption_keys[0].encrypted_key_metadata == b"key" + assert table_metadata.encryption_keys[0].properties == {"a": "b"} + assert table_metadata.encryption_keys[1].encrypted_by_id == "kek-1" + + current_snapshot = table_metadata.snapshot_by_id(3055729675574597004) + assert current_snapshot is not None + assert current_snapshot.key_id == "dek-1" + + +def test_v3_metadata_without_encryption_keys(example_table_metadata_v3: dict[str, Any]) -> None: + table_metadata = TableMetadataUtil.parse_obj(example_table_metadata_v3) + + assert isinstance(table_metadata, TableMetadataV3) + assert table_metadata.encryption_keys == [] + for snapshot in table_metadata.snapshots: + assert snapshot.key_id is None diff --git a/tests/table/test_snapshots.py b/tests/table/test_snapshots.py index 0f72b08087..4902f4d5ad 100644 --- a/tests/table/test_snapshots.py +++ b/tests/table/test_snapshots.py @@ -176,6 +176,48 @@ def test_snapshot_with_properties_repr(snapshot_with_properties: Snapshot) -> No assert snapshot_with_properties == eval(repr(snapshot_with_properties)) +@pytest.fixture +def snapshot_with_key_id() -> Snapshot: + return 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="dek-1", + ) + + +def test_serialize_snapshot_with_key_id(snapshot_with_key_id: Snapshot) -> None: + assert snapshot_with_key_id.model_dump_json() == ( + '{"snapshot-id":25,"parent-snapshot-id":19,"sequence-number":200,"timestamp-ms":1602638573590,' + '"manifest-list":"s3:/a/b/c.avro","summary":{"operation":"append"},"schema-id":3,"key-id":"dek-1"}' + ) + + +def test_deserialize_snapshot_with_key_id(snapshot_with_key_id: Snapshot) -> None: + payload = ( + '{"snapshot-id": 25, "parent-snapshot-id": 19, "sequence-number": 200, "timestamp-ms": 1602638573590, ' + '"manifest-list": "s3:/a/b/c.avro", "summary": {"operation": "append"}, "schema-id": 3, "key-id": "dek-1"}' + ) + assert Snapshot.model_validate_json(payload) == snapshot_with_key_id + + +def test_snapshot_without_key_id_omits_it(snapshot: Snapshot) -> None: + assert snapshot.key_id is None + assert "key-id" not in snapshot.model_dump_json() + + +def test_snapshot_with_key_id_repr(snapshot_with_key_id: Snapshot) -> None: + assert repr(snapshot_with_key_id) == ( + "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='dek-1')" + ) + assert snapshot_with_key_id == eval(repr(snapshot_with_key_id)) + + @pytest.fixture def manifest_file() -> ManifestFile: return ManifestFile.from_args( From 21d833fe327051b3a0a70614a9cc15892fec59a9 Mon Sep 17 00:00:00 2001 From: Xander Date: Fri, 11 Sep 2026 13:00:03 +0100 Subject: [PATCH 2/4] test --- tests/table/test_metadata.py | 17 +++++++++++------ tests/table/test_snapshots.py | 8 ++++---- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/tests/table/test_metadata.py b/tests/table/test_metadata.py index 5063c97a4b..c7d2825d5f 100644 --- a/tests/table/test_metadata.py +++ b/tests/table/test_metadata.py @@ -932,11 +932,16 @@ def test_v3_metadata_parsing_encryption_keys(example_table_metadata_v3: dict[str metadata = { **example_table_metadata_v3, "encryption-keys": [ - {"key-id": "kek-1", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "master-1", "properties": {"a": "b"}}, - {"key-id": "dek-1", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "kek-1"}, + { + "key-id": "table-key-1", + "encrypted-key-metadata": "a2V5", + "encrypted-by-id": "external-key-1", + "properties": {"a": "b"}, + }, + {"key-id": "table-key-2", "encrypted-key-metadata": "a2V5", "encrypted-by-id": "table-key-1"}, ], "snapshots": [ - {**snapshot, "key-id": "dek-1"} if snapshot["snapshot-id"] == 3055729675574597004 else snapshot + {**snapshot, "key-id": "table-key-2"} if snapshot["snapshot-id"] == 3055729675574597004 else snapshot for snapshot in example_table_metadata_v3["snapshots"] ], } @@ -944,14 +949,14 @@ def test_v3_metadata_parsing_encryption_keys(example_table_metadata_v3: dict[str table_metadata = TableMetadataUtil.parse_obj(metadata) assert isinstance(table_metadata, TableMetadataV3) - assert [key.key_id for key in table_metadata.encryption_keys] == ["kek-1", "dek-1"] + assert [key.key_id for key in table_metadata.encryption_keys] == ["table-key-1", "table-key-2"] assert table_metadata.encryption_keys[0].encrypted_key_metadata == b"key" assert table_metadata.encryption_keys[0].properties == {"a": "b"} - assert table_metadata.encryption_keys[1].encrypted_by_id == "kek-1" + assert table_metadata.encryption_keys[1].encrypted_by_id == "table-key-1" current_snapshot = table_metadata.snapshot_by_id(3055729675574597004) assert current_snapshot is not None - assert current_snapshot.key_id == "dek-1" + assert current_snapshot.key_id == "table-key-2" def test_v3_metadata_without_encryption_keys(example_table_metadata_v3: dict[str, Any]) -> None: diff --git a/tests/table/test_snapshots.py b/tests/table/test_snapshots.py index 4902f4d5ad..439375cc69 100644 --- a/tests/table/test_snapshots.py +++ b/tests/table/test_snapshots.py @@ -186,21 +186,21 @@ def snapshot_with_key_id() -> Snapshot: manifest_list="s3:/a/b/c.avro", summary=Summary(Operation.APPEND), schema_id=3, - key_id="dek-1", + key_id="table-key-1", ) def test_serialize_snapshot_with_key_id(snapshot_with_key_id: Snapshot) -> None: assert snapshot_with_key_id.model_dump_json() == ( '{"snapshot-id":25,"parent-snapshot-id":19,"sequence-number":200,"timestamp-ms":1602638573590,' - '"manifest-list":"s3:/a/b/c.avro","summary":{"operation":"append"},"schema-id":3,"key-id":"dek-1"}' + '"manifest-list":"s3:/a/b/c.avro","summary":{"operation":"append"},"schema-id":3,"key-id":"table-key-1"}' ) def test_deserialize_snapshot_with_key_id(snapshot_with_key_id: Snapshot) -> None: payload = ( '{"snapshot-id": 25, "parent-snapshot-id": 19, "sequence-number": 200, "timestamp-ms": 1602638573590, ' - '"manifest-list": "s3:/a/b/c.avro", "summary": {"operation": "append"}, "schema-id": 3, "key-id": "dek-1"}' + '"manifest-list": "s3:/a/b/c.avro", "summary": {"operation": "append"}, "schema-id": 3, "key-id": "table-key-1"}' ) assert Snapshot.model_validate_json(payload) == snapshot_with_key_id @@ -213,7 +213,7 @@ def test_snapshot_without_key_id_omits_it(snapshot: Snapshot) -> None: def test_snapshot_with_key_id_repr(snapshot_with_key_id: Snapshot) -> None: assert repr(snapshot_with_key_id) == ( "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='dek-1')" + "manifest_list='s3:/a/b/c.avro', summary=Summary(Operation.APPEND), schema_id=3, key_id='table-key-1')" ) assert snapshot_with_key_id == eval(repr(snapshot_with_key_id)) From cc9e32981f5bff366121a98aa4607149df57b0ca Mon Sep 17 00:00:00 2001 From: Xander Date: Fri, 11 Sep 2026 19:27:03 +0100 Subject: [PATCH 3/4] validate --- pyiceberg/table/metadata.py | 5 ++++- tests/table/test_metadata.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/pyiceberg/table/metadata.py b/pyiceberg/table/metadata.py index 8af73beeb9..7896d3e8eb 100644 --- a/pyiceberg/table/metadata.py +++ b/pyiceberg/table/metadata.py @@ -146,7 +146,10 @@ class EncryptedKey(IcebergBaseModel): @field_validator("encrypted_key_metadata", mode="before") def decode_encrypted_key_metadata(cls, encrypted_key_metadata: Any) -> Any: - return base64.b64decode(encrypted_key_metadata) if isinstance(encrypted_key_metadata, str) else encrypted_key_metadata + # validate=True so that malformed base64 raises instead of silently discarding characters + if isinstance(encrypted_key_metadata, str): + return base64.b64decode(encrypted_key_metadata, validate=True) + return encrypted_key_metadata @field_serializer("encrypted_key_metadata") def serialize_encrypted_key_metadata(self, encrypted_key_metadata: bytes) -> str: diff --git a/tests/table/test_metadata.py b/tests/table/test_metadata.py index c7d2825d5f..1b39ca1956 100644 --- a/tests/table/test_metadata.py +++ b/tests/table/test_metadata.py @@ -928,6 +928,20 @@ def test_encrypted_key_missing_required_field(payload: str, missing: str) -> Non assert missing in str(exc_info.value) +@pytest.mark.parametrize( + "encrypted_key_metadata", + [ + "a2V*5", # character outside the base64 alphabet + "a2V5=extra", + "a2V", # incorrect padding + "not base64", + ], +) +def test_encrypted_key_malformed_base64(encrypted_key_metadata: str) -> None: + with pytest.raises(PydanticValidationError, match="encrypted-key-metadata"): + EncryptedKey.model_validate({"key-id": "a", "encrypted-key-metadata": encrypted_key_metadata}) + + def test_v3_metadata_parsing_encryption_keys(example_table_metadata_v3: dict[str, Any]) -> None: metadata = { **example_table_metadata_v3, From a2fdf68c629ab34027632cf9bf1d0ef4820bebcc Mon Sep 17 00:00:00 2001 From: Xander Date: Fri, 11 Sep 2026 19:35:56 +0100 Subject: [PATCH 4/4] assert missing when empty --- pyiceberg/table/metadata.py | 15 +++++++++++++-- tests/table/test_metadata.py | 12 ++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/pyiceberg/table/metadata.py b/pyiceberg/table/metadata.py index 7896d3e8eb..8236f12229 100644 --- a/pyiceberg/table/metadata.py +++ b/pyiceberg/table/metadata.py @@ -21,9 +21,9 @@ import uuid from collections.abc import Iterable from copy import copy -from typing import Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, Literal -from pydantic import Field, field_serializer, field_validator, model_validator +from pydantic import Field, field_serializer, field_validator, model_serializer, model_validator from pydantic import ValidationError as PydanticValidationError from pyiceberg.exceptions import ValidationError @@ -49,6 +49,9 @@ from pyiceberg.utils.config import Config from pyiceberg.utils.datetime import datetime_to_millis +if TYPE_CHECKING: + from pydantic.functional_serializers import ModelWrapSerializerWithoutInfo + CURRENT_SNAPSHOT_ID = "current-snapshot-id" CURRENT_SCHEMA_ID = "current-schema-id" SCHEMAS = "schemas" @@ -618,6 +621,14 @@ def construct_refs(self) -> TableMetadata: encryption_keys: list[EncryptedKey] = Field(alias="encryption-keys", default_factory=list) """An optional list of encryption keys used for table encryption.""" + @model_serializer(mode="wrap") + def serialize_model(self, handler: ModelWrapSerializerWithoutInfo) -> dict[str, Any]: + """Set custom serializer to leave out `encryption-keys` when it is empty.""" + serialized: dict[str, Any] = handler(self) + if not self.encryption_keys: + serialized.pop("encryption-keys", None) + return serialized + def model_dump_json(self, exclude_none: bool = True, exclude: Any | None = None, by_alias: bool = True, **kwargs: Any) -> str: raise NotImplementedError("Writing V3 is not yet supported, see: https://github.com/apache/iceberg-python/issues/1551") diff --git a/tests/table/test_metadata.py b/tests/table/test_metadata.py index 1b39ca1956..d696a16205 100644 --- a/tests/table/test_metadata.py +++ b/tests/table/test_metadata.py @@ -980,3 +980,15 @@ def test_v3_metadata_without_encryption_keys(example_table_metadata_v3: dict[str assert table_metadata.encryption_keys == [] for snapshot in table_metadata.snapshots: assert snapshot.key_id is None + + assert "encryption-keys" not in table_metadata.model_dump(mode="json") + + +def test_v3_metadata_with_encryption_keys_serializes_them(example_table_metadata_v3: dict[str, Any]) -> None: + table_metadata = TableMetadataUtil.parse_obj( + {**example_table_metadata_v3, "encryption-keys": [{"key-id": "a", "encrypted-key-metadata": "a2V5"}]} + ) + + assert table_metadata.model_dump(mode="json")["encryption-keys"] == [ + {"key-id": "a", "encrypted-key-metadata": "a2V5", "properties": {}} + ]