pfparsons commented on code in PR #48009:
URL: https://github.com/apache/arrow/pull/48009#discussion_r2505968147
##########
python/pyarrow/tests/parquet/test_encryption.py:
##########
@@ -514,31 +533,33 @@ def kms_factory(kms_connection_configuration):
# assert table.num_rows == result_table.num_rows
[email protected](reason="External key material not supported yet")
-def test_encrypted_parquet_write_external(tempdir, data_table):
- """Write an encrypted parquet, with external key
- material.
- Currently it's not implemented, so should throw
- an exception"""
+def test_encrypted_parquet_write_read_external(tempdir, data_table,
+ external_encryption_config):
+ """Write an encrypted parquet with external key material.
+ """
path = tempdir / PARQUET_NAME
- # Encrypt the file with the footer key
- encryption_config = pe.EncryptionConfiguration(
- footer_key=FOOTER_KEY_NAME,
- column_keys={},
- internal_key_material=False)
-
- kms_connection_config = pe.KmsConnectionConfig(
- custom_kms_conf={FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8")}
- )
-
- def kms_factory(kms_connection_configuration):
- return InMemoryKmsClient(kms_connection_configuration)
+ kms_connection_config, crypto_factory = write_encrypted_file(
+ path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
+ external_encryption_config)
- crypto_factory = pe.CryptoFactory(kms_factory)
- # Write with encryption properties
- write_encrypted_parquet(path, data_table, encryption_config,
- kms_connection_config, crypto_factory)
+ verify_file_encrypted(path)
+
+ decryption_config = pe.DecryptionConfiguration()
+ result_table = read_encrypted_parquet(
+ path, decryption_config, kms_connection_config, crypto_factory,
+ internal_key_material=False)
+
+ try:
+ store = pa._parquet_encryption.FileSystemKeyMaterialStore.for_file(
+ path)
+ assert len(key_ids := store.get_key_id_set()) == (
+ len(external_encryption_config.column_keys[COL_KEY_NAME]) + 1 )
+ assert all([store.get_key_material(k) is not None for k in key_ids ])
+ except Exception as e:
+ pytest.fail(f"Unable to read external key material store: {e}")
+
+ assert data_table.equals(result_table)
Review Comment:
@adamreeve Thank you very much for the review! I've incorporated all of
your changes and added a key rotation unit test. Sorry about missing the
pre-commit checks. I believe I've corrected all the formatter/linter errors
and have a better setup to avoid this in the future.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
