feat(encryption): [9/N] Read encrypted manifest lists

pyiceberg/catalog/__init__.py

@@ -1026,13 +1026,14 @@ def namespace_exists(self, namespace: str | Identifier) -> bool:
     @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}

pyiceberg/cli/output.py

@@ -165,13 +165,14 @@ def files(self, table: Table, history: bool) -> None:
 
         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):

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)

pyiceberg/manifest.py

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

pyiceberg/table/__init__.py

@@ -28,11 +28,12 @@
 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 @@
 
     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,13 +1354,34 @@ def __init__(
         io: FileIO,
         catalog: Catalog,
         config: dict[str, str] = EMPTY_DICT,
+        kms_client: KeyManagementClient | None = None,
     ) -> None:
         self._identifier = identifier
         self.metadata = metadata
         self.metadata_location = metadata_location
         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 @@ def scan(
             catalog=self.catalog,
             table_identifier=self._identifier,
             table_config=self.config,
+            encryption_manager=self.encryption_manager,
         )
 
     def incremental_append_scan(
@@ -1506,6 +1537,7 @@ def incremental_append_scan(
             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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def __init__(
             case_sensitive=case_sensitive,
             options=options,
             limit=limit,
+            encryption_manager=encryption_manager,
         )
         self.snapshot_id = snapshot_id
         self.catalog = catalog
@@ -2454,7 +2491,7 @@ def _plan_manifest_entries(self) -> Iterator[list[ManifestEntry]]:
         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 @@ def _plan_files_local(self) -> Iterable[FileScanTask]:
             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 @@ def __init__(
         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 @@ def __init__(
             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 @@ def plan_files(self) -> Iterable[FileScanTask]:
             {
                 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 @@ def plan_files(self) -> Iterable[FileScanTask]:
             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:

pyiceberg/table/inspect.py

@@ -232,7 +232,7 @@ def _readable_metrics_struct(bound_type: PrimitiveType) -> pa.StructType:
 
         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 @@ def partitions(
         )
 
         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 @@ def partitions(
             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 @@ def _partition_summaries_to_rows(
         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 @@ def _files(self, snapshot_id: int | None = None, data_file_filter: set[DataFileC
         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 @@ def _all_files(self, data_file_filter: set[DataFileContent] | None = None) -> pa
             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}
 

pyiceberg/table/snapshots.py

@@ -32,6 +32,8 @@
 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 @@ def __repr__(self) -> str:
         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):