adamreeve commented on code in PR #51350:
URL: https://github.com/apache/arrow/pull/51350#discussion_r4022007129
##########
python/pyarrow/tests/parquet/test_encryption.py:
##########
@@ -650,6 +667,108 @@ def check_rotated_external_keys(master_key_id: str) ->
None:
assert data_table.equals(table_read_after_rotation)
+def recording_kms_factory(created_configs, client_class=InMemoryKmsClient):
+ """Create a KMS client factory that appends the KMS instance ID and URL of
+ each connection configuration it is given to created_configs"""
+ def kms_factory(kms_connection_configuration):
+ created_configs.append(
+ (kms_connection_configuration.kms_instance_id,
+ kms_connection_configuration.kms_instance_url))
+ return client_class(kms_connection_configuration)
+ return kms_factory
+
+
[email protected]("read_kms_url", [False, True])
+def test_read_kms_url_from_file(
+ tempdir, data_table, basic_encryption_config, read_kms_url):
+ """Read a file written with KMS connection properties configured, using a
+ KmsConnectionConfig that doesn't specify them"""
+ path = tempdir / PARQUET_NAME
+ custom_kms_conf = {
+ FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"),
+ COL_KEY_NAME: COL_KEY.decode("UTF-8"),
+ }
+
+ write_config = pe.KmsConnectionConfig(
+ kms_instance_id=KMS_INSTANCE_ID,
+ kms_instance_url=KMS_INSTANCE_URL,
+ custom_kms_conf=custom_kms_conf)
+ write_crypto_factory = pe.CryptoFactory(InMemoryKmsClient)
+ write_encrypted_parquet(path, data_table, basic_encryption_config,
+ write_config, write_crypto_factory)
+ verify_file_encrypted(path)
+
+ # Leave the KMS instance ID and URL unset when reading
+ read_config = pe.KmsConnectionConfig(custom_kms_conf=custom_kms_conf)
+ created_configs = []
+ read_crypto_factory = pe.CryptoFactory(
+ recording_kms_factory(created_configs))
+ decryption_config = pe.DecryptionConfiguration(
+ read_kms_url=read_kms_url)
+ result_table = read_encrypted_parquet(
+ path, decryption_config, read_config, read_crypto_factory)
+ assert data_table.equals(result_table)
+
+ if read_kms_url:
+ # The URL is read from the file key material
+ assert created_configs == [(KMS_INSTANCE_ID, KMS_INSTANCE_URL)]
+ else:
+ # The URL in the key material is ignored and the default provided
+ # instead.
+ assert created_configs == [(KMS_INSTANCE_ID, "DEFAULT")]
+
+
[email protected]("read_kms_url", [False, True])
+def test_key_rotation_reads_kms_url_from_file(reusable_tempdir, data_table,
+ read_kms_url):
+ """Rotate the keys of a file written with KMS connection properties
+ configured, using a KmsConnectionConfig that doesn't specify them"""
+ path = reusable_tempdir / PARQUET_NAME
+ encryption_config = pe.EncryptionConfiguration(
+ footer_key=FOOTER_KEY_NAME,
+ column_keys={COL_KEY_NAME: ["a", "b"]},
+ internal_key_material=False)
+
+ # Write initial encrypted file with external key material
+ write_config = pe.KmsConnectionConfig(
+ kms_instance_id=KMS_INSTANCE_ID,
+ kms_instance_url=KMS_INSTANCE_URL,
+ key_access_token="1")
+ write_crypto_factory = pe.CryptoFactory(MockVersioningKmsClient)
+ write_encrypted_parquet(path, data_table, encryption_config, write_config,
+ write_crypto_factory)
+
+ # Rotate keys without specifying the KMS instance ID and URL
+ rotation_config = pe.KmsConnectionConfig(key_access_token="2")
+ created_configs = []
+ rotation_crypto_factory = pe.CryptoFactory(
+ recording_kms_factory(created_configs, MockVersioningKmsClient))
+ rotation_crypto_factory.rotate_master_keys(
+ rotation_config, path, read_kms_url=read_kms_url)
+
+ if read_kms_url:
+ # The empty config provided is used to wrap new keys,
+ # and the config from the file was used to unwrap the original keys.
+ assert created_configs == [("", ""), (KMS_INSTANCE_ID,
KMS_INSTANCE_URL)]
Review Comment:
Same response as above, this is not true and the tests do pass.
--
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]
