diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd index 69226bf0c344..d361c5dc6a26 100644 --- a/python/pyarrow/_parquet.pxd +++ b/python/pyarrow/_parquet.pxd @@ -681,6 +681,10 @@ cdef extern from "parquet/encryption/encryption.h" namespace "parquet" nogil: " parquet::FileDecryptionProperties": pass + cdef cppclass CExternalFileDecryptionProperties\ + " parquet::ExternalFileDecryptionProperties": + pass + cdef cppclass CFileEncryptionProperties\ " parquet::FileEncryptionProperties": pass @@ -704,3 +708,19 @@ cdef class FileDecryptionProperties: cdef inline shared_ptr[CFileDecryptionProperties] unwrap(self): return self.properties + +cdef class ExternalFileDecryptionProperties(FileDecryptionProperties): + """File-level decryption properties for the low-level API""" + cdef: + shared_ptr[CExternalFileDecryptionProperties] properties + + @staticmethod + cdef inline ExternalFileDecryptionProperties wrap_external( + shared_ptr[CExternalFileDecryptionProperties] properties): + + result = ExternalFileDecryptionProperties() + result.properties = properties + return result + + cdef inline shared_ptr[CExternalFileDecryptionProperties] unwrap_external(self): + return self.properties diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index f92d2f27fcbf..f89d2b07bd7a 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -1507,8 +1507,12 @@ cdef class ParquetReader(_Weakrefable): thrift_container_size_limit) if decryption_properties is not None: - properties.file_decryption_properties( - decryption_properties.unwrap()) + if isinstance(decryption_properties, ExternalFileDecryptionProperties): + properties.file_decryption_properties( + static_pointer_cast[CFileDecryptionProperties, CExternalFileDecryptionProperties] ( + (decryption_properties).unwrap_external())) + else: + properties.file_decryption_properties((decryption_properties).unwrap()) arrow_props.set_pre_buffer(pre_buffer) diff --git a/python/pyarrow/_parquet_encryption.pxd b/python/pyarrow/_parquet_encryption.pxd index 97897eb243fd..81c93a776a8e 100644 --- a/python/pyarrow/_parquet_encryption.pxd +++ b/python/pyarrow/_parquet_encryption.pxd @@ -24,9 +24,11 @@ from pyarrow._parquet cimport (ParquetCipher, CFileEncryptionProperties, CExternalFileEncryptionProperties, CFileDecryptionProperties, + CExternalFileDecryptionProperties, FileEncryptionProperties, ExternalFileEncryptionProperties, FileDecryptionProperties, + ExternalFileDecryptionProperties, ParquetCipher_AES_GCM_V1, ParquetCipher_AES_GCM_CTR_V1, ParquetCipher_EXTERNAL_DBPA_V1) @@ -52,20 +54,16 @@ cdef class KmsConnectionConfig(_Weakrefable): @staticmethod cdef wrap(const CKmsConnectionConfig& config) - -cdef shared_ptr[CCryptoFactory] pyarrow_unwrap_cryptofactory(object crypto_factory) except * -cdef shared_ptr[CKmsConnectionConfig] pyarrow_unwrap_kmsconnectionconfig(object kmsconnectionconfig) except * -cdef shared_ptr[CEncryptionConfiguration] pyarrow_unwrap_encryptionconfig(object encryptionconfig) except * -cdef shared_ptr[CDecryptionConfiguration] pyarrow_unwrap_decryptionconfig(object decryptionconfig) except * -cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryptionconfig(object encryptionconfig) except * - - cdef class ExternalEncryptionConfiguration(EncryptionConfiguration): cdef shared_ptr[CExternalEncryptionConfiguration] external_configuration cdef inline shared_ptr[CExternalEncryptionConfiguration] unwrap_external(self) nogil -cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryptionconfig(object externalencryptionconfig) except * - cdef class ExternalDecryptionConfiguration(DecryptionConfiguration): cdef shared_ptr[CExternalDecryptionConfiguration] external_configuration cdef inline shared_ptr[CExternalDecryptionConfiguration] unwrap_external(self) nogil + +cdef shared_ptr[CCryptoFactory] pyarrow_unwrap_cryptofactory(object crypto_factory) except * +cdef shared_ptr[CKmsConnectionConfig] pyarrow_unwrap_kmsconnectionconfig(object kmsconnectionconfig) except * +cdef shared_ptr[CEncryptionConfiguration] pyarrow_unwrap_encryptionconfig(object encryptionconfig) except * +cdef shared_ptr[CDecryptionConfiguration] pyarrow_unwrap_decryptionconfig(object decryptionconfig) except * +cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryptionconfig(object externalencryptionconfig) except * cdef shared_ptr[CExternalDecryptionConfiguration] pyarrow_unwrap_external_decryptionconfig(object externaldecryptionconfig) except * \ No newline at end of file diff --git a/python/pyarrow/_parquet_encryption.pyx b/python/pyarrow/_parquet_encryption.pyx index 966e9d187503..88fc908e9558 100644 --- a/python/pyarrow/_parquet_encryption.pyx +++ b/python/pyarrow/_parquet_encryption.pyx @@ -419,13 +419,29 @@ cdef class ExternalDecryptionConfiguration(DecryptionConfiguration): __slots__ = () def __init__(self, *, cache_lifetime=None, app_context=None, connection_config=None): - super().__init__(cache_lifetime=cache_lifetime) + # Initialize the pointer first so the get/set forwards work. + # Super init will run the setters/getters below so we need the pointer to exist. self.external_configuration.reset(new CExternalDecryptionConfiguration()) + super().__init__(cache_lifetime=cache_lifetime) + + self.external_configuration.get().cache_lifetime_seconds = \ + self.configuration.get().cache_lifetime_seconds if app_context is not None: self.app_context = app_context if connection_config is not None: self.connection_config = connection_config + + """ Forward all attributes get/set methods to the superclass """ + """ The superclass already converts to/from bytes and does additional processing needed """ + @property + def cache_lifetime(self): + return DecryptionConfiguration.cache_lifetime.__get__(self) + + @cache_lifetime.setter + def cache_lifetime(self, value): + DecryptionConfiguration.cache_lifetime.__set__(self, value) + self.external_configuration.get().cache_lifetime_seconds = value.total_seconds() @property def app_context(self): @@ -759,6 +775,39 @@ cdef class CryptoFactory(_Weakrefable): c_file_decryption_properties) return FileDecryptionProperties.wrap(file_decryption_properties) + def external_file_decryption_properties( + self, + KmsConnectionConfig kms_connection_config, + ExternalDecryptionConfiguration decryption_config): + """Create file decryption properties. + Parameters + ---------- + kms_connection_config : KmsConnectionConfig + Configuration of connection to KMS + decryption_config : ExternalDecryptionConfiguration + Configuration of the decryption, such as cache timeout and the information on how to connect the external decryption service. + Returns + ------- + file_decryption_properties : ExternalFileDecryptionProperties + File decryption properties. + """ + cdef: + CExternalDecryptionConfiguration c_decryption_config + CResult[shared_ptr[CExternalFileDecryptionProperties]] \ + c_file_decryption_properties + if decryption_config is None: + c_decryption_config = CExternalDecryptionConfiguration() + else: + c_decryption_config = deref(decryption_config.unwrap_external().get()) + with nogil: + c_file_decryption_properties = \ + self.factory.get().SafeGetExternalFileDecryptionProperties( + deref(kms_connection_config.unwrap().get()), + c_decryption_config) + file_decryption_properties = GetResultValue( + c_file_decryption_properties) + return ExternalFileDecryptionProperties.wrap_external(file_decryption_properties) + def remove_cache_entries_for_token(self, access_token): self.factory.get().RemoveCacheEntriesForToken(tobytes(access_token)) @@ -800,4 +849,4 @@ cdef shared_ptr[CDecryptionConfiguration] pyarrow_unwrap_decryptionconfig(object cdef shared_ptr[CExternalDecryptionConfiguration] pyarrow_unwrap_external_decryptionconfig(object decryptionconfig) except *: if isinstance(decryptionconfig, ExternalDecryptionConfiguration): return ( decryptionconfig).unwrap_external() - raise TypeError("Expected DecryptionConfiguration, got %s" % type(decryptionconfig)) \ No newline at end of file + raise TypeError("Expected ExternalDecryptionConfiguration, got %s" % type(decryptionconfig)) \ No newline at end of file diff --git a/python/pyarrow/includes/libparquet_encryption.pxd b/python/pyarrow/includes/libparquet_encryption.pxd index 690c79cece8b..8550a442717d 100644 --- a/python/pyarrow/includes/libparquet_encryption.pxd +++ b/python/pyarrow/includes/libparquet_encryption.pxd @@ -22,6 +22,7 @@ from pyarrow._parquet cimport (ParquetCipher, CFileEncryptionProperties, CExternalFileEncryptionProperties, CFileDecryptionProperties, + CExternalFileDecryptionProperties, ParquetCipher_AES_GCM_V1, ParquetCipher_AES_GCM_CTR_V1, ParquetCipher_EXTERNAL_DBPA_V1) @@ -126,6 +127,9 @@ cdef extern from "parquet/encryption/crypto_factory.h" \ shared_ptr[CFileDecryptionProperties] GetFileDecryptionProperties( const CKmsConnectionConfig& kms_connection_config, const CDecryptionConfiguration& decryption_config) except +* + shared_ptr[CExternalFileDecryptionProperties] GetExternalFileDecryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CExternalDecryptionConfiguration& decryption_config) except +* void RemoveCacheEntriesForToken(const c_string& access_token) except + void RemoveCacheEntriesForAllTokens() except + @@ -164,4 +168,8 @@ cdef extern from "arrow/python/parquet_encryption.h" \ CResult[shared_ptr[CFileDecryptionProperties]] \ SafeGetFileDecryptionProperties( const CKmsConnectionConfig& kms_connection_config, - const CDecryptionConfiguration& decryption_config) \ No newline at end of file + const CDecryptionConfiguration& decryption_config) + CResult[shared_ptr[CExternalFileDecryptionProperties]] \ + SafeGetExternalFileDecryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CExternalDecryptionConfiguration& decryption_config) \ No newline at end of file diff --git a/python/pyarrow/tests/parquet/test_external_encryption.py b/python/pyarrow/tests/parquet/test_external_encryption.py index 76f5c7052713..70607efe1275 100644 --- a/python/pyarrow/tests/parquet/test_external_encryption.py +++ b/python/pyarrow/tests/parquet/test_external_encryption.py @@ -15,60 +15,101 @@ # specific language governing permissions and limitations # under the License. +import base64 +import datetime +import pyarrow +import pyarrow.parquet as pp +import pyarrow.parquet.encryption as ppe import pytest -import pyarrow.parquet.encryption as pe import re -from datetime import timedelta - -@pytest.fixture -def dummy_kms_client(): - class DummyKmsClient(pe.KmsClient): - def __init__(self, kms_connection_config): - super().__init__() - - def wrap_key(self, key_bytes, master_key_identifier): - # dummy wrap just returns key_bytes - return key_bytes - - def unwrap_key(self, wrapped_key, master_key_identifier): - # dummy unwrap just returns wrapped_key - return wrapped_key - return DummyKmsClient - -@pytest.fixture -def kms_config(): - return pe.KmsConnectionConfig( + +class FooKmsClient(ppe.KmsClient): + + def __init__(self, kms_connection_config): + ppe.KmsClient.__init__(self) + self.master_keys_map = kms_connection_config.custom_kms_conf + + def wrap_key(self, key_bytes, master_key_identifier): + master_key_bytes = self.master_keys_map[master_key_identifier].encode('utf-8') + joint_key = b"".join([master_key_bytes, key_bytes]) + return base64.b64encode(joint_key) + + def unwrap_key(self, wrapped_key, master_key_identifier): + expected_master = self.master_keys_map[master_key_identifier] + decoded_key = base64.b64decode(wrapped_key) + master_key_bytes = decoded_key[:16] + decrypted_key = decoded_key[16:] + if (expected_master == master_key_bytes.decode('utf-8')): + return decrypted_key + raise ValueError(f"Bad master key used [{master_key_bytes}] - [{decrypted_key}]") + + +def get_data_table(): + sample_data = { + "orderId": [1001, 1002, 1003], + "productId": [152, 268, 6548], + "price": [3.25, 6.48, 2.12], + "vat": [0.0, 0.2, 0.05] + } + return pyarrow.Table.from_pydict(sample_data) + + +def kms_client_factory(kms_connection_config): + return FooKmsClient(kms_connection_config) + + +def get_kms_connection_config(): + return ppe.KmsConnectionConfig( custom_kms_conf={ "footer_key": "012footer_secret", - "key_1": "column_secret001", - "key_2": "column_secret002" + "orderid_key": "column_secret001", + "productid_key": "column_secret002" } ) -@pytest.fixture -def external_encryption_config(): - return pe.ExternalEncryptionConfiguration( - footer_key="footer_key", - column_keys={ - "key_1": ["a"], + +def get_encryption_config(): + return ppe.EncryptionConfiguration( + footer_key = "footer_key", + column_keys = { + "orderid_key": ["orderId"], + "productid_key": ["productId"] }, - encryption_algorithm="AES_GCM_V1", - plaintext_footer=True, - double_wrapping=True, - cache_lifetime=timedelta(minutes=5.0), - internal_key_material=True, - data_key_length_bits=256, - per_column_encryption={ - "b": { + encryption_algorithm = "AES_GCM_V1", + cache_lifetime=datetime.timedelta(minutes=2.0), + data_key_length_bits = 128, + plaintext_footer=True + ) + + +def get_encryption_properties(): + encryption_config = get_encryption_config() + crypto_factory = ppe.CryptoFactory(kms_client_factory) + return crypto_factory.file_encryption_properties( + get_kms_connection_config(), encryption_config) + + +def get_external_encryption_config(plaintext_footer=True): + return ppe.ExternalEncryptionConfiguration( + footer_key = "footer_key", + column_keys = { + "productid_key": ["productId"] + }, + encryption_algorithm = "AES_GCM_V1", + cache_lifetime=datetime.timedelta(minutes=2.0), + data_key_length_bits = 128, + plaintext_footer=plaintext_footer, + per_column_encryption = { + "orderId": { "encryption_algorithm": "AES_GCM_CTR_V1", - "encryption_key": "key_2" - } + "encryption_key": "orderid_key" + }, }, - app_context={ + app_context = { "user_id": "Picard1701", "location": "Presidio" }, - connection_config={ + connection_config = { "EXTERNAL_DBPA_V1": { "config_file": "path/to/config/file", "config_file_decryption_key": "some_key" @@ -76,68 +117,104 @@ def external_encryption_config(): } ) -@pytest.fixture -def external_decryption_config(): - config = pe.ExternalDecryptionConfiguration( - cache_lifetime=timedelta(minutes=5.0), - app_context={ + +def get_external_encryption_properties(): + encryption_config = get_external_encryption_config() + crypto_factory = ppe.CryptoFactory(kms_client_factory) + return crypto_factory.external_file_encryption_properties( + get_kms_connection_config(), encryption_config) + + +def get_decryption_config(): + return ppe.DecryptionConfiguration(cache_lifetime=datetime.timedelta(minutes=10.0)) + + +def get_decryption_properties(): + decryption_config = get_decryption_config() + crypto_factory = ppe.CryptoFactory(kms_client_factory) + return crypto_factory.file_decryption_properties( + get_kms_connection_config(), decryption_config) + + +def get_external_decryption_config(): + return ppe.ExternalDecryptionConfiguration( + cache_lifetime=datetime.timedelta(minutes=10.0), + app_context = { "user_id": "Picard1701", "location": "Presidio" }, - connection_config={ - "AES_GCM_V1": { + connection_config = { + "EXTERNAL_DBPA_V1": { "config_file": "path/to/config/file", "config_file_decryption_key": "some_key" } } ) - config.cache_lifetime=timedelta(minutes=5.0) - return config + + +def get_external_decryption_properties(): + decryption_config = get_external_decryption_config() + crypto_factory = ppe.CryptoFactory(kms_client_factory) + return crypto_factory.external_file_decryption_properties( + get_kms_connection_config(), decryption_config) + + +def write_parquet(table, location, encryption_properties): + writer = pp.ParquetWriter(location, table.schema, encryption_properties=encryption_properties) + writer.write_table(table) + + +def read_parquet(location, decryption_properties): + reader = pp.ParquetFile(location, decryption_properties=decryption_properties) + return reader.read() + def test_encryption_configuration_properties(): """Test the standard EncryptionConfiguration properties to avoid regressions.""" - config = pe.EncryptionConfiguration( + config = ppe.EncryptionConfiguration( footer_key="footer-key-name", column_keys={ "key_1": ["a"], }, - encryption_algorithm="AES_GCM_V1", + encryption_algorithm="EXTERNAL_DBPA_V1", plaintext_footer=True, double_wrapping=True, - cache_lifetime=timedelta(minutes=5.0), + cache_lifetime=datetime.timedelta(minutes=5.0), internal_key_material=True, data_key_length_bits=256 ) - assert isinstance(config, pe.EncryptionConfiguration) + assert isinstance(config, ppe.EncryptionConfiguration) assert config.footer_key == "footer-key-name" assert config.column_keys == { "key_1": ["a"] } - assert config.encryption_algorithm == "AES_GCM_V1" + assert config.encryption_algorithm == "EXTERNAL_DBPA_V1" assert config.plaintext_footer is True assert config.double_wrapping is True - assert config.cache_lifetime == timedelta(minutes=5.0) + assert config.cache_lifetime == datetime.timedelta(minutes=5.0) assert config.internal_key_material is True assert config.data_key_length_bits == 256 -def test_external_encryption_configuration_properties(external_encryption_config): + +def test_external_encryption_configuration_properties(): """Test the ExternalEncryptionConfiguration properties including external-specific fields.""" - assert isinstance(external_encryption_config, pe.ExternalEncryptionConfiguration) + external_encryption_config = get_external_encryption_config() + assert isinstance(external_encryption_config, ppe.ExternalEncryptionConfiguration) assert external_encryption_config.footer_key == "footer_key" assert external_encryption_config.column_keys == { - "key_1": ["a"] + "productid_key": ["productId"] } assert external_encryption_config.encryption_algorithm == "AES_GCM_V1" assert external_encryption_config.plaintext_footer is True assert external_encryption_config.double_wrapping is True - assert external_encryption_config.cache_lifetime == timedelta(minutes=5.0) + assert external_encryption_config.cache_lifetime == datetime.timedelta(minutes=2.0) assert external_encryption_config.internal_key_material is True - assert external_encryption_config.data_key_length_bits == 256 + assert external_encryption_config.data_key_length_bits == 128 assert external_encryption_config.app_context == { "user_id": "Picard1701", @@ -152,25 +229,27 @@ def test_external_encryption_configuration_properties(external_encryption_config } assert external_encryption_config.per_column_encryption == { - "b": { + "orderId": { "encryption_algorithm": "AES_GCM_CTR_V1", - "encryption_key": "key_2" - } + "encryption_key": "orderid_key" + }, } + def test_external_encryption_app_context_invalid_json(): """Ensure app_context raises TypeError for non-JSON-serializable input.""" with pytest.raises(TypeError, match="Failed to serialize app_context: {'invalid': {1, 2, 3}}"): - pe.ExternalEncryptionConfiguration( + ppe.ExternalEncryptionConfiguration( footer_key="key", app_context={"invalid": set([1, 2, 3])} # sets are not JSON-serializable ) + def test_external_encryption_per_column_encryption_invalid_algorithm(): """Ensure invalid encryption_algorithm raises a ValueError or is rejected.""" with pytest.raises(ValueError, match="Invalid cipher name: INVALID"): - pe.ExternalEncryptionConfiguration( + ppe.ExternalEncryptionConfiguration( footer_key="key", per_column_encryption={ "a": { @@ -180,10 +259,11 @@ def test_external_encryption_per_column_encryption_invalid_algorithm(): } ) + def test_external_encryption_per_column_encryption_new_algorithm(): """Ensure new encryption_algorithm is accepted.""" - pe.ExternalEncryptionConfiguration( + ppe.ExternalEncryptionConfiguration( footer_key="key", per_column_encryption={ "a": { @@ -193,10 +273,11 @@ def test_external_encryption_per_column_encryption_new_algorithm(): } ) + def test_external_encryption_connection_config_invalid_types(): """Ensure connection_config rejects non-string keys or values.""" with pytest.raises(TypeError, match="All inner config keys/values must be str"): - config=pe.ExternalEncryptionConfiguration( + config=ppe.ExternalEncryptionConfiguration( footer_key="key" ) config.connection_config={ @@ -207,7 +288,7 @@ def test_external_encryption_connection_config_invalid_types(): } with pytest.raises(TypeError, match="All inner config keys/values must be str"): - config = pe.ExternalEncryptionConfiguration( + config = ppe.ExternalEncryptionConfiguration( footer_key="key" ) config.connection_config={ @@ -217,7 +298,8 @@ def test_external_encryption_connection_config_invalid_types(): } def test_external_encryption_rejects_none_values(): - config = pe.ExternalEncryptionConfiguration(footer_key="key") + """Ensure None values are rejected.""" + config = ppe.ExternalEncryptionConfiguration(footer_key="key") # per_column_encryption: expect ValueError with pytest.raises(TypeError, match="per_column_encryption cannot be None"): @@ -231,60 +313,64 @@ def test_external_encryption_rejects_none_values(): with pytest.raises(ValueError, match="Connection config value cannot be None"): config.connection_config = None -def test_external_file_encryption_properties_rejects_column_in_two_places( - dummy_kms_client, kms_config): - config = pe.ExternalEncryptionConfiguration( - footer_key="key", - column_keys={"key_1": ["a"]}, - per_column_encryption={"a": {"encryption_algorithm": "AES_GCM_V1", "encryption_key": "key_2"}}, + +def test_external_file_encryption_properties_rejects_column_in_two_places(): + """Ensure a column cannot be defined in both column_keys and per_column_encryption.""" + config = ppe.ExternalEncryptionConfiguration( + footer_key="footer_key", + column_keys={"orderid_key": ["a"]}, + per_column_encryption={"a": { + "encryption_algorithm": "AES_GCM_V1", + "encryption_key": "key_2" + }}, ) - factory = pe.CryptoFactory(dummy_kms_client) - with pytest.raises(OSError, match=re.escape("Multiple keys defined for column [a]")): - factory.external_file_encryption_properties(kms_config, config) + factory = ppe.CryptoFactory(kms_client_factory) + with pytest.raises(OSError, match=re.escape("Multiple keys defined for column [a]")): + factory.external_file_encryption_properties(get_kms_connection_config(), config) + + +def test_external_file_encryption_properties_valid(): + """Check class name and module because ExternalFileEncryptionProperties is not visible.""" + external_encryption_properties = get_external_encryption_properties() -def test_external_file_encryption_properties_valid( - dummy_kms_client, kms_config, external_encryption_config): - factory = pe.CryptoFactory(dummy_kms_client) - result = factory.external_file_encryption_properties(kms_config, external_encryption_config) + assert external_encryption_properties.__class__.__name__ == "ExternalFileEncryptionProperties" + assert external_encryption_properties.__class__.__module__ == "pyarrow._parquet" - # Instead of isinstance, check class name and module dynamically because - # ExternalFileEncryptionProperties is not visible - assert result.__class__.__name__ == "ExternalFileEncryptionProperties" - assert result.__class__.__module__ == "pyarrow._parquet" def test_decryption_configuration_properties(): """Test the standard DecryptionConfiguration properties to avoid regressions.""" - config = pe.DecryptionConfiguration() - config.cache_lifetime=timedelta(minutes=5.0) + config = ppe.DecryptionConfiguration() + config.cache_lifetime=datetime.timedelta(minutes=5.0) - assert isinstance(config, pe.DecryptionConfiguration) - assert config.cache_lifetime == timedelta(minutes=5.0) + assert isinstance(config, ppe.DecryptionConfiguration) + assert config.cache_lifetime == datetime.timedelta(minutes=5.0) -def test_external_decryption_configuration_properties(external_decryption_config): - """Test the ExternalDecryptionConfiguration properties including external-specific fields.""" - config = external_decryption_config +def test_external_decryption_configuration_properties(): + """Test the ExternalDecryptionConfiguration properties including external-specific fields.""" - assert isinstance(config, pe.ExternalDecryptionConfiguration) - assert config.cache_lifetime == timedelta(minutes=5.0) - assert config.app_context == { + external_decryption_config = get_external_decryption_config() + assert isinstance(external_decryption_config, ppe.ExternalDecryptionConfiguration) + assert external_decryption_config.cache_lifetime == datetime.timedelta(minutes=10.0) + assert external_decryption_config.app_context == { "user_id": "Picard1701", "location": "Presidio" } - assert config.connection_config == { - "AES_GCM_V1": { + assert external_decryption_config.connection_config == { + "EXTERNAL_DBPA_V1": { "config_file": "path/to/config/file", "config_file_decryption_key": "some_key" } } - + + def test_external_decryption_connection_config_invalid_types(): """Ensure connection_config rejects non-string keys or values.""" # Outer key is not a string (int instead of cipher name string) with pytest.raises(AttributeError, match="'int' object has no attribute 'upper'"): - config = pe.ExternalDecryptionConfiguration() + config = ppe.ExternalDecryptionConfiguration() config.connection_config = { 123: { # invalid outer key "config_file": "should-fail" @@ -293,14 +379,14 @@ def test_external_decryption_connection_config_invalid_types(): # Outer value is not a dict with pytest.raises(TypeError, match="Inner value for cipher AES_GCM_V1 must be a dict"): - config = pe.ExternalDecryptionConfiguration() + config = ppe.ExternalDecryptionConfiguration() config.connection_config = { "AES_GCM_V1": ["not", "a", "dict"] # invalid outer value (should be dict) } # Inner key is not a string with pytest.raises(TypeError, match="All inner config keys/values must be str"): - config = pe.ExternalDecryptionConfiguration() + config = ppe.ExternalDecryptionConfiguration() config.connection_config = { "AES_GCM_V1": { 123: "should-fail" # invalid inner key @@ -309,9 +395,48 @@ def test_external_decryption_connection_config_invalid_types(): # Inner value is not a string with pytest.raises(TypeError, match="All inner config keys/values must be str"): - config = pe.ExternalDecryptionConfiguration() + config = ppe.ExternalDecryptionConfiguration() config.connection_config = { "AES_GCM_V1": { "config_file": ["not", "a", "string"] # invalid inner value } } + + +def test_external_file_decryption_properties_valid(): + """Check class name and module because ExternalFileDecryptionProperties is not visible.""" + + external_decryption_properties = get_external_decryption_properties() + + assert external_decryption_properties.__class__.__name__ == "ExternalFileDecryptionProperties" + assert external_decryption_properties.__class__.__module__ == "pyarrow._parquet" + + +def test_read_and_write_standard_encryption(tmp_path): + """ Test a roundtrip encryption and decryption using standard encryption. """ + + data_table = get_data_table() + parquet_path = tmp_path / "test.parquet" + write_parquet(data_table, parquet_path, get_encryption_properties()) + + read_data_table = read_parquet(parquet_path, get_decryption_properties()) + assert read_data_table.equals(data_table) + assert read_data_table.num_rows == data_table.num_rows + assert read_data_table.num_columns == data_table.num_columns + assert read_data_table.schema.equals(data_table.schema) + assert read_data_table.column_names == data_table.column_names + + +def test_read_and_write_external_encryption(tmp_path): + """ Test a roundtrip encryption and decryption using external encryption. """ + + data_table = get_data_table() + parquet_path = tmp_path / "test.parquet" + write_parquet(data_table, parquet_path, get_external_encryption_properties()) + + read_data_table = read_parquet(parquet_path, get_external_decryption_properties()) + assert read_data_table.equals(data_table) + assert read_data_table.num_rows == data_table.num_rows + assert read_data_table.num_columns == data_table.num_columns + assert read_data_table.schema.equals(data_table.schema) + assert read_data_table.column_names == data_table.column_names diff --git a/python/scripts/base_app.py b/python/scripts/base_app.py index 59858f25eb21..c95dd12d5e29 100644 --- a/python/scripts/base_app.py +++ b/python/scripts/base_app.py @@ -158,6 +158,20 @@ def get_encryption_config(plaintext_footer=True): def get_decryption_config(): return ppe.DecryptionConfiguration(cache_lifetime=datetime.timedelta(minutes=2.0)) +def get_external_decryption_config(): + return ppe.ExternalDecryptionConfiguration( + cache_lifetime=datetime.timedelta(minutes=2.0), + app_context = { + "user_id": "Picard1701", + "location": "Presidio" + }, + connection_config = { + "EXTERNAL_DBPA_V1": { + "config_file": "path/to/config/file", + "config_file_decryption_key": "some_key" + } + } +) if __name__ == "__main__": create_and_encrypt_parquet()