it works

Author: rambleraptor

pyiceberg/table/encryption.py

@@ -16,12 +16,22 @@
 # 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)
+    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/update/__init__.py

@@ -85,14 +85,17 @@ 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")
     schema_: Schema = Field(alias="schema")
@@ -591,6 +594,7 @@ 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)
@@ -600,6 +604,7 @@ def _(update: AddEncryptedKeyUpdate, base_metadata: TableMetadata, context: _Tab
 
     return base_metadata.model_copy(update={"encryption_keys": base_metadata.encryption_keys + [update.key]})
 
+
 def update_table_metadata(
     base_metadata: TableMetadata,
     updates: Tuple[TableUpdate, ...],

tests/conftest.py

@@ -2341,6 +2341,7 @@ 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)

tests/table/test_init.py

@@ -79,6 +79,7 @@
     AssertLastAssignedPartitionId,
     AssertRefSnapshotId,
     AssertTableUUID,
+    RemoveEncryptedKeyUpdate,
     RemovePropertiesUpdate,
     RemoveSnapshotRefUpdate,
     RemoveSnapshotsUpdate,
@@ -1349,13 +1350,9 @@ def test_remove_statistics_update(table_v2_with_statistics: Table) -> None:
             (RemoveStatisticsUpdate(snapshot_id=123456789),),
         )
 
+
 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'))
-        )
-    )
+    update = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test", encrypted_key_metadata=base64.b64encode(b"hello")))
 
     expected = """
     {
@@ -1369,13 +1366,9 @@ def test_add_encryption_key(table_v3: Table) -> None:
 
     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'))
-        )
-    )
+    update_add = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test", encrypted_key_metadata=base64.b64encode(b"hello")))
     add_metadata = update_table_metadata(table_v3.metadata, (update_add,))
     assert len(add_metadata.encryption_keys) == 1
 
@@ -1385,12 +1378,7 @@ def test_remove_encryption_key(table_v3: Table) -> None:
 
 
 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'))
-        )
-    )
+    update_add = AddEncryptedKeyUpdate(key=EncryptedKey(key_id="test", encrypted_key_metadata=base64.b64encode(b"hello")))
     add_metadata = update_table_metadata(table_v3.metadata, (update_add,))
     assert len(add_metadata.encryption_keys) == 1
 
@@ -1400,11 +1388,6 @@ def test_remove_non_existent_encryption_key(table_v3: Table) -> None:
 
 
 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'))
-        )
-    )
+    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"):
         update_table_metadata(table_v2.metadata, (update_add,))