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

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

@@ -32,6 +32,7 @@
 
 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 (
@@ -124,6 +125,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__)
@@ -1340,13 +1342,26 @@ 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
+
+    @property
+    def encryption_manager(self) -> EncryptionManager | None:
+        """The table's encryption manager, or None if the table is not encrypted.
+
+        Raises:
+            ValueError: If the table is encrypted but no key management client was provided.
+        """
+        from pyiceberg.encryption.manager import EncryptionManager
+
+        return EncryptionManager.from_table_metadata(self._kms_client, self.metadata)
 
     def transaction(self) -> Transaction:
         """Create a new transaction object to first stage the changes, and then commit them to the catalog.
@@ -1446,6 +1461,7 @@ def scan(
             catalog=self.catalog,
             table_identifier=self._identifier,
             table_config=self.config,
+            encryption_manager=self.encryption_manager,
         )
 
     def incremental_append_scan(
@@ -1501,6 +1517,7 @@ def incremental_append_scan(
             to_snapshot_id=to_snapshot_id_inclusive,
             options=options,
             limit=limit,
+            encryption_manager=self.encryption_manager,
         )
 
     @property
@@ -2068,6 +2085,7 @@ class BaseScan(ABC):
     case_sensitive: bool
     options: Properties
     limit: int | None
+    encryption_manager: EncryptionManager | None
 
     def __init__(
         self,
@@ -2078,6 +2096,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
@@ -2086,6 +2105,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]: ...
@@ -2185,6 +2205,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,
@@ -2194,6 +2215,7 @@ def __init__(
             case_sensitive=case_sensitive,
             options=options,
             limit=limit,
+            encryption_manager=encryption_manager,
         )
         self.snapshot_id = snapshot_id
         self.catalog = catalog
@@ -2449,7 +2471,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."""
@@ -2486,7 +2508,7 @@ def _plan_files_local(self) -> Iterable[FileScanTask]:
         snapshot = self.snapshot()
         if not snapshot:
             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.
@@ -2602,6 +2624,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
@@ -2620,6 +2644,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,
@@ -2629,6 +2654,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
@@ -2678,7 +2704,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
             }
         )
@@ -2691,8 +2717,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
 
@@ -590,7 +590,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 +882,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 +971,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,7 @@
 from pyiceberg.schema import Schema
 
 if TYPE_CHECKING:
+    from pyiceberg.encryption.manager import EncryptionManager
     from pyiceberg.table.metadata import TableMetadata
 from pyiceberg.typedef import IcebergBaseModel
 
@@ -289,9 +290,21 @@ 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`.
+        """
+        if self.key_id is None:
+            return list(_manifests(io, self.manifest_list))
+
+        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):