encryption key

Author: rambleraptor

pyiceberg/table/encryption.py

@@ -0,0 +1,27 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional
+from pydantic import Field
+from pyiceberg.typedef import IcebergBaseModel
+
+
+class EncryptedKey(IcebergBaseModel):
+    key_id: str = Field(alias="key-id", description="ID of the encryption key")
+    encrypted_key_metadata: bytes = Field(alias="encrypted-key-metadata", description="Encrypted key and metadata, base64 encoded")
+    encrypted_by_id: Optional[str] = Field(alias="encrypted-by-id", description="Optional ID of the key used to encrypt or wrap `key-metadata`", default=None)
+    properties: Optional[dict[str, str]] = Field(alias="properties", description="A string to string map of additional metadata used by the table's encryption scheme", default=None)

pyiceberg/table/metadata.py

@@ -27,6 +27,7 @@
 from pyiceberg.exceptions import ValidationError
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec, assign_fresh_partition_spec_ids
 from pyiceberg.schema import Schema, assign_fresh_schema_ids
+from pyiceberg.table.encryption import EncryptedKey
 from pyiceberg.table.name_mapping import NameMapping, parse_mapping_from_json
 from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import MetadataLogEntry, Snapshot, SnapshotLogEntry
@@ -526,6 +527,7 @@ class TableMetadataV3(TableMetadataCommonFields, IcebergBaseModel):
         - Multi-argument transforms for partitioning and sorting
         - Row Lineage tracking
         - Binary deletion vectors
+        - Encryption Keys
 
     For more information:
     https://iceberg.apache.org/spec/?column-projection#version-3-extended-types-and-capabilities
@@ -562,6 +564,9 @@ def construct_refs(cls, table_metadata: TableMetadata) -> TableMetadata:
     next_row_id: Optional[int] = 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=[])
+    """The list of encryption keys for this table."""
+
     def model_dump_json(
         self, exclude_none: bool = True, exclude: Optional[Any] = None, by_alias: bool = True, **kwargs: Any
     ) -> str:

pyiceberg/table/snapshots.py

@@ -244,6 +244,7 @@ class Snapshot(IcebergBaseModel):
     manifest_list: str = Field(alias="manifest-list", description="Location of the snapshot's manifest list file")
     summary: Optional[Summary] = Field(default=None)
     schema_id: Optional[int] = Field(alias="schema-id", default=None)
+    key_id: Optional[str] = Field(alias="key-id", default=None, description="The id of the encryption key")
 
     def __str__(self) -> str:
         """Return the string representation of the Snapshot class."""

pyiceberg/table/update/__init__.py

@@ -28,6 +28,7 @@
 from pyiceberg.exceptions import CommitFailedException
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec
 from pyiceberg.schema import Schema
+from pyiceberg.table.encryption import EncryptedKey
 from pyiceberg.table.metadata import SUPPORTED_TABLE_FORMAT_VERSION, TableMetadata, TableMetadataUtil
 from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
@@ -89,6 +90,13 @@ class UpgradeFormatVersionUpdate(IcebergBaseModel):
     action: Literal["upgrade-format-version"] = Field(default="upgrade-format-version")
     format_version: int = Field(alias="format-version")
 
+class AddEncryptedKeyUpdate(IcebergBaseModel):
+    action: Literal["add-encryption-key"] = Field(default="add-encryption-key")
+    key: EncryptedKey = Field(alias="key")
+
+class RemoveEncryptedKeyUpdate(IcebergBaseModel):
+    action: Literal["remove-encryption-key"] = Field(default="remove-encryption-key")
+    key_id: str = Field(alias="key-id")
 
 class AddSchemaUpdate(IcebergBaseModel):
     action: Literal["add-schema"] = Field(default="add-schema")
@@ -230,6 +238,8 @@ class RemovePartitionStatisticsUpdate(IcebergBaseModel):
         RemoveSchemasUpdate,
         SetPartitionStatisticsUpdate,
         RemovePartitionStatisticsUpdate,
+        AddEncryptedKeyUpdate,
+        RemoveEncryptedKeyUpdate,
     ],
     Field(discriminator="action"),
 ]
@@ -594,6 +604,14 @@ def _(update: RemoveStatisticsUpdate, base_metadata: TableMetadata, context: _Ta
 
     return base_metadata.model_copy(update={"statistics": statistics})
 
+@_apply_table_update.register(AddEncryptedKeyUpdate)
+def _(update: AddEncryptedKeyUpdate, 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")
+
+    return base_metadata.model_copy(update={"encryption_keys": base_metadata.encryption_keys + [update.key]})
 
 @_apply_table_update.register(RemoveSchemasUpdate)
 def _(update: RemoveSchemasUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:

tests/conftest.py

@@ -72,7 +72,7 @@
 from pyiceberg.schema import Accessor, Schema
 from pyiceberg.serializers import ToOutputFile
 from pyiceberg.table import FileScanTask, Table
-from pyiceberg.table.metadata import TableMetadataV1, TableMetadataV2
+from pyiceberg.table.metadata import TableMetadataV1, TableMetadataV2, TableMetadataV3
 from pyiceberg.transforms import DayTransform, IdentityTransform
 from pyiceberg.types import (
     BinaryType,
@@ -2467,6 +2467,17 @@ def table_v2(example_table_metadata_v2: Dict[str, Any]) -> Table:
         catalog=NoopCatalog("NoopCatalog"),
     )
 
+@pytest.fixture
+def table_v3(example_table_metadata_v3: Dict[str, Any]) -> Table:
+    table_metadata = TableMetadataV3(**example_table_metadata_v3)
+    return Table(
+        identifier=("database", "table"),
+        metadata=table_metadata,
+        metadata_location=f"{table_metadata.location}/uuid.metadata.json",
+        io=load_file_io(),
+        catalog=NoopCatalog("NoopCatalog"),
+    )
+
 
 @pytest.fixture
 def table_v2_orc(example_table_metadata_v2: Dict[str, Any]) -> Table:

tests/table/test_init.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint:disable=redefined-outer-name
+import base64
 import json
 import uuid
 from copy import copy
@@ -49,6 +50,7 @@
     TableIdentifier,
     _match_deletes_to_data_file,
 )
+from pyiceberg.table.encryption import EncryptedKey
 from pyiceberg.table.metadata import INITIAL_SEQUENCE_NUMBER, TableMetadataUtil, TableMetadataV2, _generate_snapshot_id
 from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
@@ -66,6 +68,7 @@
 )
 from pyiceberg.table.statistics import BlobMetadata, PartitionStatisticsFile, StatisticsFile
 from pyiceberg.table.update import (
+    AddEncryptedKeyUpdate,
     AddSnapshotUpdate,
     AddSortOrderUpdate,
     AssertCreate,
@@ -1419,6 +1422,13 @@ 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("hello".encode('utf-8'))
+        )
     )
 
     expected = """
@@ -1477,3 +1487,52 @@ def test_remove_partition_statistics_update_with_invalid_snapshot_id(table_v2_wi
             table_v2_with_statistics.metadata,
             (RemovePartitionStatisticsUpdate(snapshot_id=123456789),),
         )
+      "key-id": "test",
+      "encrypted-key-metadata": "aGVsbG8="
+    }"""
+
+    assert table_v3.metadata.encryption_keys == []
+    add_metadata = update_table_metadata(table_v3.metadata, (update,))
+    assert len(add_metadata.encryption_keys) == 1
+
+    assert json.loads(add_metadata.encryption_keys[0].model_dump_json()) == json.loads(expected)
+
+def test_remove_encryption_key(table_v3: Table) -> None:
+    update_add = AddEncryptedKeyUpdate(
+        key=EncryptedKey(
+            key_id="test",
+            encrypted_key_metadata=base64.b64encode("hello".encode('utf-8'))
+        )
+    )
+    add_metadata = update_table_metadata(table_v3.metadata, (update_add,))
+    assert len(add_metadata.encryption_keys) == 1
+
+    update_remove = RemoveEncryptedKeyUpdate(key_id="test")
+    remove_metadata = update_table_metadata(add_metadata, (update_remove,))
+    assert len(remove_metadata.encryption_keys) == 0
+
+
+def test_remove_non_existent_encryption_key(table_v3: Table) -> None:
+    update_add = AddEncryptedKeyUpdate(
+        key=EncryptedKey(
+            key_id="test",
+            encrypted_key_metadata=base64.b64encode("hello".encode('utf-8'))
+        )
+    )
+    add_metadata = update_table_metadata(table_v3.metadata, (update_add,))
+    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
+
+
+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("hello_v2".encode('utf-8'))
+        )
+    )
+    with pytest.raises(ValueError, match=r"Cannot add encryption keys from Iceberg v1 or v2 table"):
+        update_table_metadata(table_v2.metadata, (update_add,))