From 0c0a7ddd241c8e62a995e2b11fa91f64a6ddb626 Mon Sep 17 00:00:00 2001 From: Xander Date: Tue, 10 Nov 2026 13:05:58 -0500 Subject: [PATCH 1/2] Write encrypted Parquet data files Data files of a table with encryption.key-id are written with Parquet modular encryption, using a fresh key and AAD prefix per file, and the key metadata is recorded on the DataFile. Like Java, the AAD prefix is not stored in the file. --- pyiceberg/encryption/manager.py | 24 +++++++++++++----- pyiceberg/io/fileformat.py | 6 ++++- pyiceberg/io/pyarrow.py | 29 +++++++++++++++++++--- tests/io/test_pyarrow.py | 44 +++++++++++++++++++++++++++++++++ 4 files changed, 92 insertions(+), 11 deletions(-) diff --git a/pyiceberg/encryption/manager.py b/pyiceberg/encryption/manager.py index 73c57dfd5b..4c8c51691a 100644 --- a/pyiceberg/encryption/manager.py +++ b/pyiceberg/encryption/manager.py @@ -74,12 +74,7 @@ def from_table_metadata(cls, kms_client: KeyManagementClient | None, metadata: T if kms_client is None: raise ValueError(f"Table has {TableProperties.ENCRYPTION_KEY_ID} set but no key management client was provided") - key_length = property_as_int(metadata.properties, TableProperties.ENCRYPTION_DATA_KEY_LENGTH) - key_size = AesKeySize.from_key_length( - TableProperties.ENCRYPTION_DATA_KEY_LENGTH_DEFAULT if key_length is None else key_length - ) - - return cls(kms_client, table_key_id, metadata.encryption_keys, key_size) + return cls(kms_client, table_key_id, metadata.encryption_keys, _data_key_size(metadata)) @property def encryption_keys(self) -> list[EncryptedKey]: @@ -170,6 +165,23 @@ def _unwrap_kek(self, kek: EncryptedKey) -> SecureKey: return plaintext +def _new_data_file_key_metadata(metadata: TableMetadata) -> StandardKeyMetadata | None: + """Generate key metadata for a new data file, or return None if the table is not encrypted.""" + from pyiceberg.table import TableProperties + + if TableProperties.ENCRYPTION_KEY_ID not in metadata.properties: + return None + + return StandardKeyMetadata.generate(_data_key_size(metadata)) + + +def _data_key_size(metadata: TableMetadata) -> AesKeySize: + from pyiceberg.table import TableProperties + + key_length = property_as_int(metadata.properties, TableProperties.ENCRYPTION_DATA_KEY_LENGTH) + return AesKeySize.from_key_length(TableProperties.ENCRYPTION_DATA_KEY_LENGTH_DEFAULT if key_length is None else key_length) + + def _now_ms() -> int: return time.time_ns() // 1_000_000 diff --git a/pyiceberg/io/fileformat.py b/pyiceberg/io/fileformat.py index 65c0cbeff9..7d9345a6e9 100644 --- a/pyiceberg/io/fileformat.py +++ b/pyiceberg/io/fileformat.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: import pyarrow as pa + from pyiceberg.encryption.key_metadata import StandardKeyMetadata from pyiceberg.io.pyarrow import StatsAggregator @@ -159,7 +160,10 @@ def create_writer( output_file: OutputFile, file_schema: Schema, properties: Properties, - ) -> FileFormatWriter: ... + key_metadata: StandardKeyMetadata | None = None, + ) -> FileFormatWriter: + """Create a writer for one file, encrypting it with `key_metadata` when it is set.""" + ... def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None: """Add format-specific Arrow field metadata.""" diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 9f1e4c9b58..a107913f7c 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -69,6 +69,7 @@ from pyiceberg.conversions import to_bytes from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.manager import _new_data_file_key_metadata from pyiceberg.exceptions import ResolveError from pyiceberg.expressions import AlwaysTrue, BooleanExpression, BoundIsNaN, BoundIsNull, BoundTerm, Not, Or from pyiceberg.expressions.literals import Literal @@ -1175,6 +1176,16 @@ def _parquet_decryption_properties(key_metadata: bytes) -> pq.FileDecryptionProp return create_decryption_properties(metadata.encryption_key, aad_prefix=metadata.aad_prefix) +def _parquet_encryption_properties(key_metadata: StandardKeyMetadata) -> pq.FileEncryptionProperties: + try: + from pyarrow.parquet.encryption import create_encryption_properties + except ImportError as e: + raise ImportError(f"pyarrow >= 25.0.0 required to write encrypted Parquet files, found {pyarrow.__version__}") from e + + # Java does not store the AAD prefix in the file, readers take it from key metadata + return create_encryption_properties(key_metadata.encryption_key, aad_prefix=key_metadata.aad_prefix, store_aad_prefix=False) + + def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]: _check_encryption_supported(data_file) if data_file.file_format == FileFormat.PARQUET: @@ -2764,13 +2775,21 @@ def data_file_statistics_from_parquet_metadata( class ParquetFormatWriter(FileFormatWriter): """Writes Arrow tables to a Parquet file.""" - def __init__(self, output_file: OutputFile, file_schema: Schema, properties: Properties) -> None: + def __init__( + self, + output_file: OutputFile, + file_schema: Schema, + properties: Properties, + key_metadata: StandardKeyMetadata | None = None, + ) -> None: self._output_file = output_file self._file_schema = file_schema self._properties = properties self._writer: pq.ParquetWriter | None = None self._fos: OutputStream | None = None self._parquet_writer_kwargs = _get_parquet_writer_kwargs(properties) + if key_metadata is not None: + self._parquet_writer_kwargs["encryption_properties"] = _parquet_encryption_properties(key_metadata) self._row_group_size = property_as_int( properties=properties, property_name=TableProperties.PARQUET_ROW_GROUP_LIMIT, @@ -2823,8 +2842,9 @@ def create_writer( output_file: OutputFile, file_schema: Schema, properties: Properties, + key_metadata: StandardKeyMetadata | None = None, ) -> ParquetFormatWriter: - return ParquetFormatWriter(output_file, file_schema, properties) + return ParquetFormatWriter(output_file, file_schema, properties, key_metadata) def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None: if include_field_ids: @@ -2871,7 +2891,8 @@ def write_data_file(task: WriteTask) -> DataFile: partition_key=task.partition_key, ) fo = io.new_output(file_path) - writer = format_model.create_writer(fo, file_schema, table_metadata.properties) + key_metadata = _new_data_file_key_metadata(table_metadata) + writer = format_model.create_writer(fo, file_schema, table_metadata.properties, key_metadata) with writer: writer.write(arrow_table) statistics = writer.result() @@ -2889,7 +2910,7 @@ def write_data_file(task: WriteTask) -> DataFile: # Just copy these from the table for now spec_id=table_metadata.default_spec_id, equality_ids=None, - key_metadata=None, + key_metadata=key_metadata.to_bytes() if key_metadata is not None else None, **statistics.to_serialized_dict(), ) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 7fd809e1f1..041575d707 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -3678,6 +3678,50 @@ def test_write_file_parquet_round_trip(tmp_path: Path, table_schema_simple: Sche assert result.column_names == ["foo", "bar", "baz"] +def test_write_file_encrypts_parquet_for_encrypted_table(tmp_path: Path, table_schema_simple: Schema) -> None: + from pyarrow.parquet import encryption + + _, task = _simple_write_task_and_metadata(tmp_path, table_schema_simple, {}) + table_metadata = TableMetadataV3( + location=f"file://{tmp_path}", + last_column_id=3, + current_schema_id=1, + schemas=[table_schema_simple], + partition_specs=[PartitionSpec()], + properties={TableProperties.ENCRYPTION_KEY_ID: "master-1", TableProperties.ENCRYPTION_DATA_KEY_LENGTH: "32"}, + next_row_id=0, + ) + + data_file, other_data_file = write_file( + io=PyArrowFileIO(), table_metadata=table_metadata, tasks=iter([task, replace(task, task_id=1)]) + ) + + assert data_file.key_metadata is not None and other_data_file.key_metadata is not None + key_metadata = StandardKeyMetadata.from_bytes(data_file.key_metadata) + other_key_metadata = StandardKeyMetadata.from_bytes(other_data_file.key_metadata) + assert len(key_metadata.encryption_key) == 32 + assert key_metadata.encryption_key != other_key_metadata.encryption_key + assert key_metadata.aad_prefix != other_key_metadata.aad_prefix + file_path = data_file.file_path.removeprefix("file://") + with open(file_path, "rb") as f: + assert f.read(4) == b"PARE" + + with pytest.raises(OSError, match="AAD prefix used for file encryption, but not stored in file"): + pq.read_table(file_path, decryption_properties=encryption.create_decryption_properties(key_metadata.encryption_key)) + + batches = _task_to_record_batches( + PyArrowFileIO(), + FileScanTask(data_file), + bound_row_filter=AlwaysTrue(), + projected_schema=table_schema_simple, + table_schema=table_schema_simple, + projected_field_ids={1, 2, 3}, + positional_deletes=None, + case_sensitive=True, + ) + assert pa.Table.from_batches(batches).to_pydict() == {"foo": ["a", "b", "c"], "bar": [1, 2, 3], "baz": [True, False, True]} + + def test__to_requested_schema_timestamps( arrow_table_schema_with_all_timestamp_precisions: pa.Schema, arrow_table_with_all_timestamp_precisions: pa.Table, From 53e2f9436c73b909e69a9b0fd2558d4b5f793ff9 Mon Sep 17 00:00:00 2001 From: Xander Date: Tue, 10 Nov 2026 13:26:01 -0500 Subject: [PATCH 2/2] Only pass key_metadata to create_writer for encrypted tables Keeps plaintext writes working with FileFormatModel implementations that predate the key_metadata argument. --- pyiceberg/io/pyarrow.py | 6 ++++-- tests/io/test_fileformat.py | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index a107913f7c..0f30f29266 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2891,8 +2891,10 @@ def write_data_file(task: WriteTask) -> DataFile: partition_key=task.partition_key, ) fo = io.new_output(file_path) - key_metadata = _new_data_file_key_metadata(table_metadata) - writer = format_model.create_writer(fo, file_schema, table_metadata.properties, key_metadata) + if (key_metadata := _new_data_file_key_metadata(table_metadata)) is None: + writer = format_model.create_writer(fo, file_schema, table_metadata.properties) + else: + writer = format_model.create_writer(fo, file_schema, table_metadata.properties, key_metadata=key_metadata) with writer: writer.write(arrow_table) statistics = writer.result() diff --git a/tests/io/test_fileformat.py b/tests/io/test_fileformat.py index d5d487fa6d..aee71e3d67 100644 --- a/tests/io/test_fileformat.py +++ b/tests/io/test_fileformat.py @@ -48,7 +48,7 @@ def format(self) -> FileFormat: def file_extension(self) -> str: return "orc" - def create_writer(self, output_file: Any, file_schema: Any, properties: Any) -> Any: + def create_writer(self, output_file: Any, file_schema: Any, properties: Any, key_metadata: Any = None) -> Any: raise NotImplementedError def add_field_metadata(self, field: Any, metadata: Any, include_field_ids: bool) -> None: @@ -89,7 +89,7 @@ def format(self) -> FileFormat: def file_extension(self) -> str: return "orc" - def create_writer(self, output_file: Any, file_schema: Any, properties: Any) -> Any: + def create_writer(self, output_file: Any, file_schema: Any, properties: Any, key_metadata: Any = None) -> Any: raise NotImplementedError def add_field_metadata(self, field: Any, metadata: Any, include_field_ids: bool) -> None: