This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-4159-e3564a54c10960307147377c578680276865ed18 in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
commit 718c7908b36117974257e3d214959bffcada14a1 Author: Xander <[email protected]> AuthorDate: Fri Nov 6 23:40:36 2026 +0000 feat(encryption): [9/N] Read encrypted manifest lists (#4159) * Read encrypted manifest lists * Pass encryption manager to remaining manifest list reads inspect.partitions(), the CLI files output and purge_table read manifest lists without the encryption manager, so they failed on encrypted tables. purge_table now builds the manager before dropping the table, so a missing KMS client no longer leaves a dropped table with unpurged files. * Reuse the encryption manager until table metadata changes Building a manager on every access discarded its KEK cache, so each manifest list read called the KMS again. * Read encrypted and plain manifest lists through one path in Snapshot.manifests * Track the encryption manager's metadata in a NamedTuple Replace the bare (metadata, manager) tuple with a NamedTuple so the property reads as "rebuild when the metadata changed since it was built", while still swapping the pair in a single assignment. Co-authored-by: Copilot App <[email protected]> --------- Co-authored-by: Kevin Liu <[email protected]> Co-authored-by: Copilot App <[email protected]> --- pyiceberg/catalog/__init__.py | 3 +- pyiceberg/cli/output.py | 3 +- pyiceberg/encryption/io.py | 50 ++++++++++++++ pyiceberg/manifest.py | 7 +- pyiceberg/table/__init__.py | 54 +++++++++++++-- pyiceberg/table/inspect.py | 13 ++-- pyiceberg/table/snapshots.py | 19 +++++- tests/catalog/test_catalog_behaviors.py | 15 +++++ tests/cli/test_output.py | 9 ++- tests/conftest.py | 73 +++++++++++++++++++++ tests/encryption/test_stream.py | 31 +++++++++ tests/table/encryption/README.md | 25 +++++++ .../encryption/manifest-list-v3-encrypted.avro | Bin 0 -> 1568 bytes tests/table/test_init.py | 58 ++++++++++++++++ tests/utils/test_manifest.py | 42 ++++++++++++ 15 files changed, 384 insertions(+), 18 deletions(-) diff --git a/pyiceberg/catalog/__init__.py b/pyiceberg/catalog/__init__.py index cdb4d26b3..ce0902f33 100644 --- a/pyiceberg/catalog/__init__.py +++ b/pyiceberg/catalog/__init__.py @@ -1026,13 +1026,14 @@ class MetastoreCatalog(Catalog, ABC): @override def purge_table(self, identifier: str | Identifier) -> None: table = self.load_table(identifier) + encryption_manager = table.encryption_manager self.drop_table(identifier) io = load_file_io(self.properties, table.metadata_location) metadata = table.metadata manifest_lists_to_delete = set() manifests_to_delete: list[ManifestFile] = [] for snapshot in metadata.snapshots: - manifests_to_delete += snapshot.manifests(io) + manifests_to_delete += snapshot.manifests(io, encryption_manager) manifest_lists_to_delete.add(snapshot.manifest_list) manifest_paths_to_delete = {manifest.manifest_path for manifest in manifests_to_delete} diff --git a/pyiceberg/cli/output.py b/pyiceberg/cli/output.py index cfb00abb7..f2e6ddaed 100644 --- a/pyiceberg/cli/output.py +++ b/pyiceberg/cli/output.py @@ -165,13 +165,14 @@ class ConsoleOutput(Output): snapshot_tree = Tree(f"Snapshots: {'.'.join(table.name())}") io = table.io + encryption_manager = table.encryption_manager for snapshot in snapshots: list_tree = snapshot_tree.add( f"Snapshot {snapshot.snapshot_id}, schema {snapshot.schema_id}: {snapshot.manifest_list}" ) - manifest_list = snapshot.manifests(io) + manifest_list = snapshot.manifests(io, encryption_manager) for manifest in manifest_list: manifest_tree = list_tree.add(f"Manifest: {manifest.manifest_path}") for manifest_entry in manifest.fetch_manifest_entry(io, discard_deleted=False): diff --git a/pyiceberg/encryption/io.py b/pyiceberg/encryption/io.py new file mode 100644 index 000000000..1ba9dd5b7 --- /dev/null +++ b/pyiceberg/encryption/io.py @@ -0,0 +1,50 @@ +# 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. +"""Input files for metadata encrypted with AGS1 streams.""" + +from __future__ import annotations + +import io + +from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.stream import _decrypt_ags1 +from pyiceberg.io import InputFile, InputStream + + +class EncryptedInputFile(InputFile): + """Decrypts an AGS1 encrypted file, such as a manifest list, when it is opened.""" + + def __init__(self, input_file: InputFile, key_metadata: StandardKeyMetadata) -> None: + super().__init__(input_file.location) + self._input_file = input_file + self._key_metadata = key_metadata + + def __len__(self) -> int: + """Return the length of the decrypted file, in bytes.""" + return len(self._decrypt()) + + def exists(self) -> bool: + """Check whether the encrypted file exists.""" + return self._input_file.exists() + + def open(self, seekable: bool = True) -> InputStream: + """Read and decrypt the whole file, authenticating every block before returning any of it.""" + return io.BytesIO(self._decrypt()) + + def _decrypt(self) -> bytes: + with self._input_file.open(seekable=False) as stream: + return _decrypt_ags1(stream.read(), self._key_metadata) diff --git a/pyiceberg/manifest.py b/pyiceberg/manifest.py index 0881273c4..3975c9c2a 100644 --- a/pyiceberg/manifest.py +++ b/pyiceberg/manifest.py @@ -34,6 +34,8 @@ from pydantic_core import to_json from pyiceberg.avro.codecs import AVRO_CODEC_KEY, AvroCompressionCodec from pyiceberg.avro.file import AvroFile, AvroOutputFile from pyiceberg.conversions import to_bytes +from pyiceberg.encryption.io import EncryptedInputFile +from pyiceberg.encryption.key_metadata import StandardKeyMetadata from pyiceberg.exceptions import ValidationError from pyiceberg.io import FileIO, InputFile, OutputFile from pyiceberg.partitioning import PartitionSpec @@ -998,7 +1000,7 @@ def clear_manifest_cache() -> None: _manifest_cache.clear() -def _manifests(io: FileIO, manifest_list: str) -> tuple[ManifestFile, ...]: +def _manifests(io: FileIO, manifest_list: str, key_metadata: StandardKeyMetadata | None = None) -> tuple[ManifestFile, ...]: """Read manifests from a manifest list, reusing cached ManifestFile objects. Caches individual ManifestFile objects by manifest_path. This is memory-efficient @@ -1018,11 +1020,14 @@ def _manifests(io: FileIO, manifest_list: str) -> tuple[ManifestFile, ...]: Args: io: FileIO instance for reading the manifest list. manifest_list: Path to the manifest list file. + key_metadata: Decrypts the manifest list, if it is encrypted. Returns: A tuple of ManifestFile objects. """ file = io.new_input(manifest_list) + if key_metadata is not None: + file = EncryptedInputFile(file, key_metadata) manifest_files = list(read_manifest_list(file)) return tuple(_manifest_cache.get_or_cache(manifest_file) for manifest_file in manifest_files) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 73c435f07..df7aaf824 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -28,11 +28,12 @@ from dataclasses import dataclass from functools import cached_property from itertools import chain from types import TracebackType -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar from cachetools import LRUCache from pydantic import Field +from pyiceberg.encryption.kms import KeyManagementClient from pyiceberg.exceptions import CommitFailedException, ValidationException from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, IsNull, Or, Reference from pyiceberg.expressions.visitors import ( @@ -126,6 +127,7 @@ if TYPE_CHECKING: from pyiceberg.catalog import Catalog from pyiceberg.catalog.rest.scan_planning import RESTContentFile, RESTDeleteFile, RESTFileScanTask + from pyiceberg.encryption.manager import EncryptionManager from pyiceberg.table.update.snapshot import _SnapshotProducer logger = logging.getLogger(__name__) @@ -1327,6 +1329,13 @@ class CommitTableResponse(IcebergBaseModel): metadata_location: str = Field(alias="metadata-location") +class _CachedEncryptionManager(NamedTuple): + """An encryption manager, reused with its unwrapped KEKs until the table metadata it was built from changes.""" + + metadata: TableMetadata + manager: EncryptionManager | None + + class Table: """An Iceberg table.""" @@ -1345,6 +1354,7 @@ class Table: io: FileIO, catalog: Catalog, config: dict[str, str] = EMPTY_DICT, + kms_client: KeyManagementClient | None = None, ) -> None: self._identifier = identifier self.metadata = metadata @@ -1352,6 +1362,26 @@ class Table: self.io = io self.catalog = catalog self.config = config + self._kms_client = kms_client + self._encryption_manager: _CachedEncryptionManager | None = None + + @property + def encryption_manager(self) -> EncryptionManager | None: + """The table's encryption manager, or None if the table is not encrypted. + + The manager is built from the current metadata, and rebuilt when the metadata changes. + + Raises: + ValueError: If the table is encrypted but no key management client was provided. + """ + from pyiceberg.encryption.manager import EncryptionManager + + cached = self._encryption_manager + if cached is None or cached.metadata is not self.metadata: + cached = self._encryption_manager = _CachedEncryptionManager( + self.metadata, EncryptionManager.from_table_metadata(self._kms_client, self.metadata) + ) + return cached.manager def transaction(self) -> Transaction: """Create a new transaction object to first stage the changes, and then commit them to the catalog. @@ -1451,6 +1481,7 @@ class Table: catalog=self.catalog, table_identifier=self._identifier, table_config=self.config, + encryption_manager=self.encryption_manager, ) def incremental_append_scan( @@ -1506,6 +1537,7 @@ class Table: to_snapshot_id=to_snapshot_id_inclusive, options=options, limit=limit, + encryption_manager=self.encryption_manager, ) @property @@ -2073,6 +2105,7 @@ class BaseScan(ABC): case_sensitive: bool options: Properties limit: int | None + encryption_manager: EncryptionManager | None def __init__( self, @@ -2083,6 +2116,7 @@ class BaseScan(ABC): case_sensitive: bool = True, options: Properties = EMPTY_DICT, limit: int | None = None, + encryption_manager: EncryptionManager | None = None, ): self.table_metadata = table_metadata self.io = io @@ -2091,6 +2125,7 @@ class BaseScan(ABC): self.case_sensitive = case_sensitive self.options = options self.limit = limit + self.encryption_manager = encryption_manager @abstractmethod def plan_files(self) -> Iterable[ScanTask]: ... @@ -2190,6 +2225,7 @@ class TableScan(BaseScan): catalog: Catalog | None = None, table_identifier: Identifier | None = None, table_config: Properties = EMPTY_DICT, + encryption_manager: EncryptionManager | None = None, ): super().__init__( table_metadata=table_metadata, @@ -2199,6 +2235,7 @@ class TableScan(BaseScan): case_sensitive=case_sensitive, options=options, limit=limit, + encryption_manager=encryption_manager, ) self.snapshot_id = snapshot_id self.catalog = catalog @@ -2454,7 +2491,7 @@ class DataScan(TableScan): if not snapshot: return iter([]) - return self._manifest_planner.plan_manifest_entries(snapshot.manifests(self.io)) + return self._manifest_planner.plan_manifest_entries(snapshot.manifests(self.io, self.encryption_manager)) def _should_use_server_side_planning(self) -> bool: """Check if server-side scan planning should be used for this scan.""" @@ -2493,7 +2530,7 @@ class DataScan(TableScan): if self.snapshot_id is not None: raise ValueError(f"Snapshot not found: {self.snapshot_id}") return [] - return self._manifest_planner.plan_files(snapshot.manifests(self.io)) + return self._manifest_planner.plan_files(snapshot.manifests(self.io, self.encryption_manager)) def plan_files(self) -> Iterable[FileScanTask]: """Plans the relevant files by filtering on the PartitionSpecs. @@ -2609,6 +2646,8 @@ class IncrementalAppendScan(BaseScan): An integer representing the number of rows to return in the scan result. If None, fetches all matching rows. + encryption_manager: + Decrypts the manifest lists of encrypted snapshots. """ from_snapshot_id: int | None @@ -2627,6 +2666,7 @@ class IncrementalAppendScan(BaseScan): to_snapshot_id: int | None = None, options: Properties = EMPTY_DICT, limit: int | None = None, + encryption_manager: EncryptionManager | None = None, ): super().__init__( table_metadata=table_metadata, @@ -2636,6 +2676,7 @@ class IncrementalAppendScan(BaseScan): case_sensitive=case_sensitive, options=options, limit=limit, + encryption_manager=encryption_manager, ) self.from_snapshot_id = from_snapshot_id self.from_snapshot_inclusive = from_snapshot_inclusive @@ -2685,7 +2726,7 @@ class IncrementalAppendScan(BaseScan): { manifest_file for snapshot in append_snapshots - for manifest_file in snapshot.manifests(self.io) + for manifest_file in snapshot.manifests(self.io, self.encryption_manager) if manifest_file.content == ManifestContent.DATA and manifest_file.added_snapshot_id in append_snapshot_ids } ) @@ -2698,8 +2739,9 @@ class IncrementalAppendScan(BaseScan): options=self.options, ).plan_files( manifests=manifests, - manifest_entry_filter=lambda manifest_entry: manifest_entry.snapshot_id in append_snapshot_ids - and manifest_entry.status == ManifestEntryStatus.ADDED, + manifest_entry_filter=lambda manifest_entry: ( + manifest_entry.snapshot_id in append_snapshot_ids and manifest_entry.status == ManifestEntryStatus.ADDED + ), ) def to_arrow(self) -> pa.Table: diff --git a/pyiceberg/table/inspect.py b/pyiceberg/table/inspect.py index e580b66dc..d51079fc9 100644 --- a/pyiceberg/table/inspect.py +++ b/pyiceberg/table/inspect.py @@ -232,7 +232,7 @@ class InspectTable: entries = [] snapshot = self._get_snapshot(snapshot_id) - for manifest in snapshot.manifests(self.tbl.io): + for manifest in snapshot.manifests(self.tbl.io, self.tbl.encryption_manager): for entry in manifest.fetch_manifest_entry(io=self.tbl.io, discard_deleted=False): column_sizes = entry.data_file.column_sizes or {} value_counts = entry.data_file.value_counts or {} @@ -383,7 +383,7 @@ class InspectTable: ) snapshot = self._get_snapshot(snapshot_id) - spec_ids = {manifest.partition_spec_id for manifest in snapshot.manifests(self.tbl.io)} + spec_ids = {manifest.partition_spec_id for manifest in snapshot.manifests(self.tbl.io, self.tbl.encryption_manager)} partition_record = self.tbl.metadata.specs_struct(spec_ids=spec_ids) has_partitions = len(partition_record.fields) > 0 @@ -404,6 +404,7 @@ class InspectTable: row_filter=row_filter, case_sensitive=case_sensitive, snapshot_id=snapshot.snapshot_id, + encryption_manager=self.tbl.encryption_manager, ) partitions_map: dict[tuple[str, Any], Any] = {} @@ -590,7 +591,7 @@ class InspectTable: specs = self.tbl.metadata.specs() manifests = [] if snapshot: - for manifest in snapshot.manifests(self.tbl.io): + for manifest in snapshot.manifests(self.tbl.io, self.tbl.encryption_manager): is_data_file = manifest.content == ManifestContent.DATA is_delete_file = manifest.content == ManifestContent.DELETES manifest_row = { @@ -882,7 +883,8 @@ class InspectTable: executor = ExecutorFactory.get_or_create() results = list( executor.map( - lambda manifest_list: self._get_files_from_manifest(manifest_list, data_file_filter), snapshot.manifests(io) + lambda manifest_list: self._get_files_from_manifest(manifest_list, data_file_filter), + snapshot.manifests(io, self.tbl.encryption_manager), ) ) return pa.concat_tables(results) @@ -970,7 +972,8 @@ class InspectTable: return pa.Table.from_pylist([], schema=self._get_files_schema()) executor = ExecutorFactory.get_or_create() - manifest_lists = executor.map(lambda snapshot: snapshot.manifests(self.tbl.io), snapshots) + encryption_manager = self.tbl.encryption_manager + manifest_lists = executor.map(lambda snapshot: snapshot.manifests(self.tbl.io, encryption_manager), snapshots) unique_manifests = {(manifest.manifest_path, manifest) for manifest_list in manifest_lists for manifest in manifest_list} diff --git a/pyiceberg/table/snapshots.py b/pyiceberg/table/snapshots.py index 66d3d7f23..624ce486b 100644 --- a/pyiceberg/table/snapshots.py +++ b/pyiceberg/table/snapshots.py @@ -32,6 +32,8 @@ from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec from pyiceberg.schema import Schema if TYPE_CHECKING: + from pyiceberg.encryption.key_metadata import StandardKeyMetadata + from pyiceberg.encryption.manager import EncryptionManager from pyiceberg.table.metadata import TableMetadata from pyiceberg.typedef import IcebergBaseModel @@ -289,9 +291,20 @@ class Snapshot(IcebergBaseModel): filtered_fields = [field for field in fields if field is not None] return f"Snapshot({', '.join(filtered_fields)})" - def manifests(self, io: FileIO) -> list[ManifestFile]: - """Return the manifests for the given snapshot.""" - return list(_manifests(io, self.manifest_list)) + def manifests(self, io: FileIO, encryption_manager: EncryptionManager | None = None) -> list[ManifestFile]: + """Return the manifests for the given snapshot. + + Args: + io: Reads the manifest list. + encryption_manager: Decrypts the manifest list, required when the snapshot has a `key-id`. + """ + key_metadata: StandardKeyMetadata | None = None + if self.key_id is not None: + if encryption_manager is None: + raise ValueError(f"Snapshot {self.snapshot_id} is encrypted, but no encryption manager was provided") + key_metadata = encryption_manager.decrypt_manifest_list_key_metadata(self.key_id) + + return list(_manifests(io, self.manifest_list, key_metadata)) class MetadataLogEntry(IcebergBaseModel): diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 1148ca852..5bd756d74 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -19,6 +19,7 @@ Consolidated behavior tests for InMemoryCatalog and SqlCatalog. """ +import json import os from collections.abc import Generator from pathlib import Path @@ -505,6 +506,20 @@ def test_purge_table(catalog: Catalog, table_schema_simple: Schema, test_table_i catalog.load_table(test_table_identifier) +def test_purge_encrypted_table_without_kms_client( + catalog: Catalog, test_table_identifier: Identifier, example_table_metadata_v3_encrypted: dict[str, Any], tmp_path: Path +) -> None: + metadata_location = tmp_path / "00000-encrypted.metadata.json" + metadata_location.write_text(json.dumps(example_table_metadata_v3_encrypted)) + catalog.create_namespace(Catalog.namespace_from(test_table_identifier)) + catalog.register_table(test_table_identifier, str(metadata_location)) + + with pytest.raises(ValueError, match="no key management client was provided"): + catalog.purge_table(test_table_identifier) + + assert catalog.table_exists(test_table_identifier) + + # List tables tests diff --git a/tests/cli/test_output.py b/tests/cli/test_output.py index 7b27a02c5..8dd8e0dae 100644 --- a/tests/cli/test_output.py +++ b/tests/cli/test_output.py @@ -18,7 +18,8 @@ import json import pytest -from pyiceberg.cli.output import JsonOutput +from pyiceberg.cli.output import ConsoleOutput, JsonOutput +from pyiceberg.table import Table from pyiceberg.table.refs import SnapshotRefType @@ -46,3 +47,9 @@ def test_json_describe_refs_emits_one_object_per_ref(capsys: pytest.CaptureFixtu "max_ref_age_ms": "forever", } ] + + +def test_files_of_encrypted_table(table_v3_encrypted: Table, capsys: pytest.CaptureFixture[str]) -> None: + ConsoleOutput().files(table_v3_encrypted, history=True) + + assert "Snapshot 1" in capsys.readouterr().out diff --git a/tests/conftest.py b/tests/conftest.py index f243397a6..05d50925c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2652,6 +2652,79 @@ def table_v3(example_table_metadata_v3: dict[str, Any]) -> Table: ) [email protected] +def example_table_metadata_v3_encrypted() -> dict[str, Any]: + """Metadata of an encrypted table written by iceberg-rust, whose manifest list is in tests/table/encryption.""" + manifest_list = Path(__file__).parent / "table" / "encryption" / "manifest-list-v3-encrypted.avro" + return { + "format-version": 3, + "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1", + "location": "memory:///table", + "last-sequence-number": 0, + "last-updated-ms": 1602638573590, + "last-column-id": 1, + "current-schema-id": 0, + "schemas": [{"type": "struct", "schema-id": 0, "fields": [{"id": 1, "name": "x", "required": True, "type": "long"}]}], + "default-spec-id": 0, + "partition-specs": [{"spec-id": 0, "fields": []}], + "last-partition-id": 1000, + "default-sort-order-id": 0, + "sort-orders": [{"order-id": 0, "fields": []}], + "properties": {"encryption.key-id": "master-1"}, + "current-snapshot-id": 1, + "snapshots": [ + { + "snapshot-id": 1, + "sequence-number": 0, + "timestamp-ms": 0, + "manifest-list": str(manifest_list), + "summary": {"operation": "append"}, + "schema-id": 0, + "key-id": "1fb95f04-7f98-4b1f-a08f-e810cfb907ff", + } + ], + "encryption-keys": [ + { + "key-id": "5b29941b-b3b4-4d57-8758-a9d1bf340b39", + "encrypted-key-metadata": "WVMIPwafypGJRgXOICFKcL+m8K4B40gEIJFsgC7rlitqN9PUpDeT19OZPs8=", + "encrypted-by-id": "master-1", + "properties": {"KEY_TIMESTAMP": "1781172301703"}, + }, + { + "key-id": "1fb95f04-7f98-4b1f-a08f-e810cfb907ff", + "encrypted-key-metadata": ( + "XcrqLxxyf9a7FAFXM63u3MEH8CSWqh6g1H8fJD2/6zx0EjcPzPaP/5AtRpO2s6p2isV8s7a4m77zj+GIqNmIpbA2Xw==" + ), + "encrypted-by-id": "5b29941b-b3b4-4d57-8758-a9d1bf340b39", + }, + ], + "snapshot-log": [{"snapshot-id": 1, "timestamp-ms": 0}], + "metadata-log": [], + "refs": {"main": {"snapshot-id": 1, "type": "branch"}}, + "next-row-id": 0, + } + + [email protected] +def table_v3_encrypted(example_table_metadata_v3_encrypted: dict[str, Any]) -> Table: + from pyiceberg.catalog.noop import NoopCatalog + from pyiceberg.encryption.ciphers import SecureKey + from tests.encryption.memory_kms import MemoryKeyManagementClient + + kms = MemoryKeyManagementClient() + kms.add_master_key("master-1", SecureKey(bytes(range(16)))) + + table_metadata = TableMetadataV3(**example_table_metadata_v3_encrypted) + return Table( + identifier=("database", "table"), + metadata=table_metadata, + metadata_location=f"{table_metadata.location}/uuid.metadata.json", + io=load_file_io(), + catalog=NoopCatalog("NoopCatalog"), + kms_client=kms, + ) + + @pytest.fixture def table_v2_orc(example_table_metadata_v2: dict[str, Any]) -> Table: import copy diff --git a/tests/encryption/test_stream.py b/tests/encryption/test_stream.py index 2f96e7155..5211f39d1 100644 --- a/tests/encryption/test_stream.py +++ b/tests/encryption/test_stream.py @@ -15,8 +15,12 @@ # specific language governing permissions and limitations # under the License. +from dataclasses import replace +from pathlib import Path + import pytest +from pyiceberg.encryption.io import EncryptedInputFile from pyiceberg.encryption.key_metadata import StandardKeyMetadata from pyiceberg.encryption.stream import ( _BLOCK_OVERHEAD, @@ -29,6 +33,7 @@ from pyiceberg.encryption.stream import ( _encrypt_ags1, _stream_block_aad, ) +from pyiceberg.io.pyarrow import PyArrowFileIO KEY = b"0123456789012345" AAD_PREFIX = b"0123456789abcdef" @@ -203,3 +208,29 @@ def test_decrypt_rejects_the_wrong_key() -> None: with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): decrypt(stream, key=b"1123456789012345") + + +def test_encrypted_input_file(tmp_path: Path) -> None: + plaintext = b"manifest list contents" + encrypted = _encrypt_ags1(plaintext, KEY_METADATA) + path = tmp_path / "manifest-list.avro" + path.write_bytes(encrypted) + + input_file = EncryptedInputFile(PyArrowFileIO().new_input(str(path)), replace(KEY_METADATA, file_length=len(encrypted))) + + assert input_file.location == str(path) + assert input_file.exists() + assert len(input_file) == len(plaintext) + with input_file.open() as stream: + assert stream.read() == plaintext + + +def test_encrypted_input_file_rejects_a_truncated_file(tmp_path: Path) -> None: + encrypted = _encrypt_ags1(b"manifest list contents", KEY_METADATA) + path = tmp_path / "manifest-list.avro" + path.write_bytes(encrypted[:-1]) + + input_file = EncryptedInputFile(PyArrowFileIO().new_input(str(path)), replace(KEY_METADATA, file_length=len(encrypted))) + + with pytest.raises(ValueError, match="key metadata expects"): + input_file.open() diff --git a/tests/table/encryption/README.md b/tests/table/encryption/README.md new file mode 100644 index 000000000..ae20c650c --- /dev/null +++ b/tests/table/encryption/README.md @@ -0,0 +1,25 @@ +<!-- + ~ 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. +--> + +# Encrypted manifest list test data + +`manifest-list-v3-encrypted.avro` is an empty manifest list encrypted by iceberg-rust, copied from its +`crates/iceberg/testdata/manifests_lists/`. It was written by iceberg-rust's `EncryptionManager`, with a memory KMS +holding master key `master-1` = bytes `0x00` to `0x0f`. Its keys and snapshot are in the +`example_table_metadata_v3_encrypted` fixture in `tests/conftest.py`. diff --git a/tests/table/encryption/manifest-list-v3-encrypted.avro b/tests/table/encryption/manifest-list-v3-encrypted.avro new file mode 100644 index 000000000..4963f0eff Binary files /dev/null and b/tests/table/encryption/manifest-list-v3-encrypted.avro differ diff --git a/tests/table/test_init.py b/tests/table/test_init.py index ff7ecd7c3..481326b44 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -2058,6 +2058,64 @@ def test_serialize_encryption_key_updates() -> None: assert TypeAdapter(TableUpdate).validate_json(remove.model_dump_json()) == remove +def test_encryption_manager_of_unencrypted_table(table_v3: Table) -> None: + assert table_v3.encryption_manager is None + + +def test_encryption_manager_requires_kms_client(example_table_metadata_v3_encrypted: dict[str, Any]) -> None: + table_metadata = TableMetadataUtil.parse_obj(example_table_metadata_v3_encrypted) + table = Table(("database", "table"), table_metadata, "s3://bucket/metadata.json", load_file_io(), NoopCatalog("noop")) + + with pytest.raises(ValueError, match="no key management client was provided"): + table.scan() + + +def test_encryption_manager_follows_current_metadata(table_v3_encrypted: Table) -> None: + assert table_v3_encrypted.encryption_manager is not None + + # Catalogs must not let encryption.key-id be removed, but removing it shows the manager follows the current metadata + table_v3_encrypted.metadata = table_v3_encrypted.metadata.model_copy(update={"properties": {}}) + + assert table_v3_encrypted.encryption_manager is None + + +def test_encryption_manager_is_reused_until_metadata_changes(table_v3_encrypted: Table, monkeypatch: pytest.MonkeyPatch) -> None: + kms = table_v3_encrypted._kms_client + assert kms is not None + unwrap_key = kms.unwrap_key + calls = [] + + def _spy(wrapped_key: bytes, wrapping_key_id: str) -> bytes: + calls.append(wrapping_key_id) + return unwrap_key(wrapped_key, wrapping_key_id) + + monkeypatch.setattr(kms, "unwrap_key", _spy) + + for _ in range(3): + list(table_v3_encrypted.scan().plan_files()) + table_v3_encrypted.inspect.manifests() + table_v3_encrypted.inspect.all_manifests() + assert len(calls) == 1 + + table_v3_encrypted.metadata = table_v3_encrypted.metadata.model_copy() + list(table_v3_encrypted.scan().plan_files()) + assert len(calls) == 2 + + +def test_scan_encrypted_table(table_v3_encrypted: Table) -> None: + """The manifest list was encrypted by iceberg-rust, and lists no manifests.""" + assert list(table_v3_encrypted.scan().plan_files()) == [] + assert list(table_v3_encrypted.incremental_append_scan().plan_files()) == [] + assert table_v3_encrypted.scan().to_arrow().num_rows == 0 + + +def test_inspect_encrypted_table(table_v3_encrypted: Table) -> None: + assert table_v3_encrypted.inspect.manifests().num_rows == 0 + assert table_v3_encrypted.inspect.all_manifests().num_rows == 0 + assert table_v3_encrypted.inspect.entries().num_rows == 0 + assert table_v3_encrypted.inspect.partitions().num_rows == 0 + + def test_add_snapshot_update_fails_without_first_row_id(table_v3: Table) -> None: new_snapshot = Snapshot( snapshot_id=25, diff --git a/tests/utils/test_manifest.py b/tests/utils/test_manifest.py index 45e4e894e..accfb07f5 100644 --- a/tests/utils/test_manifest.py +++ b/tests/utils/test_manifest.py @@ -16,6 +16,7 @@ # under the License. # pylint: disable=redefined-outer-name,arguments-renamed,fixme import importlib +from dataclasses import replace from pathlib import Path from tempfile import TemporaryDirectory from typing import Any @@ -26,6 +27,8 @@ import pytest import pyiceberg.manifest as manifest_module from pyiceberg.avro.codecs import AvroCompressionCodec from pyiceberg.avro.file import DEFAULT_SYNC_INTERVAL, AvroFile, AvroOutputFile +from pyiceberg.encryption.manager import EncryptionManager +from pyiceberg.encryption.stream import _encrypt_ags1 from pyiceberg.io import load_file_io from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.manifest import ( @@ -52,6 +55,7 @@ from pyiceberg.schema import Schema from pyiceberg.table.snapshots import Operation, Snapshot, Summary from pyiceberg.typedef import Record, TableVersion from pyiceberg.types import IntegerType, ListType, LongType, NestedField, StructType +from tests.encryption.memory_kms import MemoryKeyManagementClient @pytest.fixture(autouse=True) @@ -506,6 +510,44 @@ def test_read_manifest_v2(generated_manifest_file_file_v2: str) -> None: assert entry.status == ManifestEntryStatus.ADDED +def test_read_encrypted_manifest_list(generated_manifest_file_file_v2: str, tmp_path: Path) -> None: + io = load_file_io() + kms = MemoryKeyManagementClient() + kms.add_master_key("master-1") + encryption_manager = EncryptionManager(kms, "master-1") + + key_metadata = encryption_manager.generate_key_metadata() + encrypted = _encrypt_ags1(Path(generated_manifest_file_file_v2).read_bytes(), key_metadata) + manifest_list = tmp_path / "manifest-list.avro" + manifest_list.write_bytes(encrypted) + key_id = encryption_manager.encrypt_manifest_list_key_metadata(replace(key_metadata, file_length=len(encrypted))) + + plaintext = Snapshot( + snapshot_id=25, + parent_snapshot_id=19, + timestamp_ms=1602638573590, + manifest_list=generated_manifest_file_file_v2, + summary=Summary(Operation.APPEND), + schema_id=3, + ) + snapshot = plaintext.model_copy(update={"manifest_list": str(manifest_list), "key_id": key_id}) + + assert snapshot.manifests(io, encryption_manager) == plaintext.manifests(io) + + +def test_read_encrypted_manifest_list_without_encryption_manager(generated_manifest_file_file_v2: str) -> None: + snapshot = Snapshot( + snapshot_id=25, + timestamp_ms=1602638573590, + manifest_list=generated_manifest_file_file_v2, + summary=Summary(Operation.APPEND), + key_id="key", + ) + + with pytest.raises(ValueError, match="Snapshot 25 is encrypted, but no encryption manager was provided"): + snapshot.manifests(load_file_io()) + + def test_read_manifest_cache(generated_manifest_file_file_v2: str) -> None: """Test that ManifestFile objects are cached and reused across multiple reads.
