diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 9f1e4c9b58..b4a0bd894d 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1161,8 +1161,10 @@ def _create_arrow_read_format(file_format: FileFormat, key_metadata: bytes | Non def _check_encryption_supported(data_file: DataFile) -> None: - if data_file.key_metadata is not None and data_file.file_format != FileFormat.PARQUET: - raise ValueError(f"Cannot read encrypted {data_file.file_format} file {data_file.file_path}: only Parquet is supported") + if data_file.key_metadata is not None and data_file.file_format not in (FileFormat.PARQUET, FileFormat.PUFFIN): + raise ValueError( + f"Cannot read encrypted {data_file.file_format} file {data_file.file_path}: only Parquet and Puffin are supported" + ) def _parquet_decryption_properties(key_metadata: bytes) -> pq.FileDecryptionProperties: diff --git a/pyiceberg/table/deletion_vector.py b/pyiceberg/table/deletion_vector.py index 9c56af24a2..16c991c04c 100644 --- a/pyiceberg/table/deletion_vector.py +++ b/pyiceberg/table/deletion_vector.py @@ -21,12 +21,14 @@ from pyroaring import BitMap, FrozenBitMap +from pyiceberg.encryption.io import EncryptedInputFile +from pyiceberg.encryption.key_metadata import StandardKeyMetadata from pyiceberg.table.puffin import PuffinFile if TYPE_CHECKING: import pyarrow as pa - from pyiceberg.io import FileIO + from pyiceberg.io import FileIO, InputFile from pyiceberg.manifest import DataFile EMPTY_BITMAP = FrozenBitMap() @@ -152,10 +154,18 @@ def _has_deletion_vector_content_reference(dv: "DataFile") -> bool: return dv.content_offset is not None or dv.content_size_in_bytes is not None or dv.referenced_data_file is not None +def _new_input(io: "FileIO", dv: "DataFile") -> "InputFile": + """Return the Puffin file to read, decrypting it when it has key metadata so content offsets refer to the plaintext.""" + input_file = io.new_input(dv.file_path) + if dv.key_metadata is not None: + return EncryptedInputFile(input_file, StandardKeyMetadata.from_bytes(dv.key_metadata)) + return input_file + + def _read_deletion_vector(io: "FileIO", dv: "DataFile") -> DeletionVector: content_offset, content_size_in_bytes, referenced_data_file = _validate_deletion_vector_content(dv) - with io.new_input(dv.file_path).open() as fi: + with _new_input(io, dv).open() as fi: fi.seek(content_offset) payload = fi.read(content_size_in_bytes) @@ -173,7 +183,7 @@ def read_deletion_vectors(io: "FileIO", dv: "DataFile") -> list[DeletionVector]: if _has_deletion_vector_content_reference(dv): return [_read_deletion_vector(io, dv)] - with io.new_input(dv.file_path).open() as fi: + with _new_input(io, dv).open() as fi: return deletion_vectors_from_puffin_file(PuffinFile(fi.read())) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 7fd809e1f1..419d130965 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -44,6 +44,7 @@ from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTDeleteFile, RESTFileScanTask, RESTPositionDeleteFile from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.stream import _encrypt_ags1 from pyiceberg.exceptions import ResolveError from pyiceberg.expressions import ( AlwaysFalse, @@ -2083,6 +2084,68 @@ def test_read_deletion_vector_blob_from_content_range(tmp_path: Path) -> None: assert deletes == {referenced_data_file: pa.chunked_array([[1, 3, 5]])} +def _puffin_delete_file(file_path: str, **kwargs: Any) -> DataFile: + return DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=file_path, + file_format=FileFormat.PUFFIN, + record_count=3, + **kwargs, + ) + + +def test_read_deletes_from_puffin_file_encrypted_by_iceberg_rust() -> None: + """The deletion vector was encrypted by iceberg-rust, see tests/table/encryption/README.md.""" + file_path = str(Path(__file__).parent.parent / "table" / "encryption" / "dv-v3-encrypted.puffin") + referenced_data_file = "s3://bucket/table/data/00000.parquet" + key_metadata = bytes.fromhex("01203031323334353637383961626364656602206665646362613938373635343332313002ea04") + + for dv in [ + _puffin_delete_file(file_path, key_metadata=key_metadata), + _puffin_delete_file( + file_path, + key_metadata=key_metadata, + referenced_data_file=referenced_data_file, + content_offset=4, + content_size_in_bytes=46, + ), + ]: + assert _read_deletes(PyArrowFileIO(), dv) == {referenced_data_file: pa.chunked_array([[1, 3, 5]])} + + +def test_read_deletes_from_encrypted_puffin_file(tmp_path: Path) -> None: + referenced_data_file = f"{tmp_path}/data.parquet" + delete_file_path = f"{tmp_path}/deletes.puffin" + puffin_payload = _deletion_vector_puffin_payload(referenced_data_file) + blob = PuffinFile(puffin_payload).footer.blobs[0] + key_metadata = StandardKeyMetadata.generate() + encrypted = _encrypt_ags1(puffin_payload, key_metadata) + key_metadata = replace(key_metadata, file_length=len(encrypted)) + wrong_key = replace(key_metadata, encryption_key=bytes(16)) + + with open(delete_file_path, "wb") as f: + f.write(encrypted) + + for content_reference, plaintext_error in [ + ({}, "Incorrect magic bytes"), + ( + {"referenced_data_file": referenced_data_file, "content_offset": blob.offset, "content_size_in_bytes": blob.length}, + "Invalid bitmap data length", + ), + ]: + dv = _puffin_delete_file(delete_file_path, key_metadata=key_metadata.to_bytes(), **content_reference) + assert _read_deletes(PyArrowFileIO(), dv) == {referenced_data_file: pa.chunked_array([[1, 3, 5]])} + + with pytest.raises(ValueError, match="GCM tag check failed"): + _read_deletes( + PyArrowFileIO(), _puffin_delete_file(delete_file_path, key_metadata=wrong_key.to_bytes(), **content_reference) + ) + + with pytest.raises(ValueError, match=plaintext_error): + _read_deletes(PyArrowFileIO(), _puffin_delete_file(delete_file_path, **content_reference)) + + def test_read_all_delete_files_keeps_multiple_dv_content_ranges_for_same_path(tmp_path: Path) -> None: referenced_data_files = [f"{tmp_path}/data-{index}.parquet" for index in range(2)] first_dv_blob = _deletion_vector_blob(_deletion_vector_bitmap_payload([1, 3])) @@ -4234,10 +4297,9 @@ def test_read_encrypted_parquet_fails_closed(tmp_path: Path) -> None: with open(path, "rb") as f: _create_arrow_read_format(FileFormat.PARQUET, key).make_fragment(f).to_table() - for file_format in [FileFormat.ORC, FileFormat.PUFFIN]: - delete_file = DataFile.from_args(file_path="deletes", file_format=file_format, key_metadata=key_metadata.to_bytes()) - with pytest.raises(ValueError, match="file deletes: only Parquet is supported"): - _read_deletes(PyArrowFileIO(), delete_file) + delete_file = DataFile.from_args(file_path="deletes", file_format=FileFormat.ORC, key_metadata=key_metadata.to_bytes()) + with pytest.raises(ValueError, match="file deletes: only Parquet and Puffin are supported"): + _read_deletes(PyArrowFileIO(), delete_file) def test_parse_location_defaults() -> None: diff --git a/tests/table/encryption/README.md b/tests/table/encryption/README.md index d88061c682..b1c4c0e124 100644 --- a/tests/table/encryption/README.md +++ b/tests/table/encryption/README.md @@ -31,3 +31,8 @@ entry would record is in `test_read_manifest_encrypted_by_iceberg_rust` in `test `data-v3-encrypted.parquet` is a Parquet data file with three rows, written and encrypted by iceberg-rust's `ParquetWriter` with a key from its `EncryptionManager`. The key metadata that its manifest entry would record is in `test_task_to_record_batches_encrypted_by_iceberg_rust` in `tests/io/test_pyarrow.py`. + +`dv-v3-encrypted.puffin` is a Puffin file with one deletion vector, deleting positions 1, 3 and 5 of +`s3://bucket/table/data/00000.parquet`. It was written and encrypted by iceberg-rust's `PuffinWriter` with key +`0123456789abcdef` and AAD prefix `fedcba9876543210`. The key metadata that its manifest entry would record is in +`test_read_deletes_from_puffin_file_encrypted_by_iceberg_rust` in `tests/io/test_pyarrow.py`. diff --git a/tests/table/encryption/dv-v3-encrypted.puffin b/tests/table/encryption/dv-v3-encrypted.puffin new file mode 100644 index 0000000000..31713f7380 Binary files /dev/null and b/tests/table/encryption/dv-v3-encrypted.puffin differ