diff --git a/python/pyarrow/_parquet_encryption.pxd b/python/pyarrow/_parquet_encryption.pxd index 9f3bb78cf5c8..97897eb243fd 100644 --- a/python/pyarrow/_parquet_encryption.pxd +++ b/python/pyarrow/_parquet_encryption.pxd @@ -63,4 +63,9 @@ cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryp 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 * \ No newline at end of file +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[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 3ed854d69074..f5ad46146cab 100644 --- a/python/pyarrow/_parquet_encryption.pyx +++ b/python/pyarrow/_parquet_encryption.pyx @@ -39,7 +39,7 @@ cdef ParquetCipher cipher_from_name(name): elif name == 'EXTERNAL_DBPA_V1': return ParquetCipher_EXTERNAL_DBPA_V1 else: - raise ValueError(f'Invalid cipher name: {name!r}') + raise ValueError(f'Invalid cipher name: {name}') cdef cipher_to_name(ParquetCipher cipher): @@ -327,6 +327,85 @@ cdef class DecryptionConfiguration(_Weakrefable): cdef inline shared_ptr[CDecryptionConfiguration] unwrap(self) nogil: return self.configuration +cdef class ExternalDecryptionConfiguration(DecryptionConfiguration): + """Configuration of the external decryption""" + # Avoid mistakingly creating attributes + __slots__ = () + + def __init__(self, *, cache_lifetime=None, app_context=None, connection_config=None): + super().__init__(cache_lifetime=cache_lifetime) + self.external_configuration.reset(new CExternalDecryptionConfiguration()) + + if app_context is not None: + self.app_context = app_context + if connection_config is not None: + self.connection_config = connection_config + + @property + def app_context(self): + """Get the application context as a dictionary.""" + app_context_str = frombytes(self.external_configuration.get().app_context) + if not app_context_str: + return {} + try: + return json.loads(app_context_str) + except Exception: + raise ValueError(f"Invalid JSON stored in app_context: {app_context_str}") + + @app_context.setter + def app_context(self, dict value): + """Set the application context from a dictionary.""" + if value is None: + raise ValueError("app_context must be JSON-serializable") + + try: + serialized = json.dumps(value) + self.external_configuration.get().app_context = tobytes(serialized) + except Exception: + raise TypeError(f"Failed to serialize app_context: {repr(value)}") + + @property + def connection_config(self): + """Get the connection configuration as a Python dictionary.""" + + cdef pair[ParquetCipher, unordered_map[c_string, c_string]] outer_pair + cdef pair[c_string, c_string] inner_pair + result = {} + + for outer_pair in self.external_configuration.get().connection_config: + cipher_name = cipher_to_name(outer_pair.first) + inner_map = {} + for inner_pair in outer_pair.second: + inner_map[frombytes(inner_pair.first)] = frombytes(inner_pair.second) + result[cipher_name] = inner_map + + return result + + @connection_config.setter + def connection_config(self, dict value): + """Set the connection configuration from a Python dictionary.""" + if value is None: + raise ValueError("Connection config value cannot be None") + + cdef unordered_map[ParquetCipher, unordered_map[c_string, c_string]] cpp_map + cdef unordered_map[c_string, c_string] inner_cpp_map + cdef ParquetCipher cipher_enum + + for cipher_name, inner_dict in value.items(): + cipher_enum = cipher_from_name(cipher_name) + if not isinstance(inner_dict, dict): + raise TypeError(f"Inner value for cipher {cipher_name} must be a dict") + inner_cpp_map.clear() + for k, v in inner_dict.items(): + if not isinstance(k, str) or not isinstance(v, str): + raise TypeError("All inner config keys/values must be str") + inner_cpp_map[tobytes(k)] = tobytes(v) + cpp_map[cipher_enum] = inner_cpp_map + + self.external_configuration.get().connection_config = cpp_map + + cdef inline shared_ptr[CExternalDecryptionConfiguration] unwrap_external(self) nogil: + return self.external_configuration cdef class KmsConnectionConfig(_Weakrefable): """Configuration of the connection to the Key Management Service (KMS)""" @@ -630,4 +709,9 @@ cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryp cdef shared_ptr[CDecryptionConfiguration] pyarrow_unwrap_decryptionconfig(object decryptionconfig) except *: if isinstance(decryptionconfig, DecryptionConfiguration): return ( decryptionconfig).unwrap() + raise TypeError("Expected DecryptionConfiguration, got %s" % type(decryptionconfig)) + +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 diff --git a/python/pyarrow/includes/libparquet_encryption.pxd b/python/pyarrow/includes/libparquet_encryption.pxd index 06242f263117..263b09aac19f 100644 --- a/python/pyarrow/includes/libparquet_encryption.pxd +++ b/python/pyarrow/includes/libparquet_encryption.pxd @@ -107,6 +107,13 @@ cdef extern from "parquet/encryption/crypto_factory.h" \ CDecryptionConfiguration() except + double cache_lifetime_seconds + cdef cppclass CExternalDecryptionConfiguration\ + " parquet::encryption::ExternalDecryptionConfiguration": + CExternalDecryptionConfiguration() except + + double cache_lifetime_seconds + c_string app_context + unordered_map[ParquetCipher, unordered_map[c_string, c_string]] connection_config + cdef cppclass CCryptoFactory" parquet::encryption::CryptoFactory": void RegisterKmsClientFactory( shared_ptr[CKmsClientFactory] kms_client_factory) except + diff --git a/python/pyarrow/parquet/encryption.py b/python/pyarrow/parquet/encryption.py index af83738db7e3..a9633710c719 100644 --- a/python/pyarrow/parquet/encryption.py +++ b/python/pyarrow/parquet/encryption.py @@ -20,5 +20,6 @@ ExternalEncryptionConfiguration, EncryptionConfiguration, DecryptionConfiguration, + ExternalDecryptionConfiguration, KmsConnectionConfig, KmsClient) diff --git a/python/pyarrow/tests/parquet/test_external_encryption.py b/python/pyarrow/tests/parquet/test_external_encryption.py index d6af033e6527..5a01cf3742ab 100644 --- a/python/pyarrow/tests/parquet/test_external_encryption.py +++ b/python/pyarrow/tests/parquet/test_external_encryption.py @@ -54,6 +54,24 @@ def external_encryption_config(): } ) +@pytest.fixture +def external_decryption_config(): + config = pe.ExternalDecryptionConfiguration( + cache_lifetime=timedelta(minutes=5.0), + app_context={ + "user_id": "Picard1701", + "location": "Presidio" + }, + connection_config={ + "AES_GCM_V1": { + "config_file": "path/to/config/file", + "config_file_decryption_key": "some_key" + } + } + ) + config.cache_lifetime=timedelta(minutes=5.0) + return config + def test_encryption_configuration_properties(): """Test the standard EncryptionConfiguration properties to avoid regressions.""" @@ -131,7 +149,7 @@ def test_external_encryption_app_context_invalid_json(): 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'"): + with pytest.raises(ValueError, match="Invalid cipher name: INVALID"): pe.ExternalEncryptionConfiguration( footer_key="key", per_column_encryption={ @@ -217,3 +235,66 @@ def unwrap_key(self, wrapped_key, master_key_identifier): 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) + + assert isinstance(config, pe.DecryptionConfiguration) + assert config.cache_lifetime == 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 + + assert isinstance(config, pe.ExternalDecryptionConfiguration) + assert config.cache_lifetime == timedelta(minutes=5.0) + assert config.app_context == { + "user_id": "Picard1701", + "location": "Presidio" + } + assert config.connection_config == { + "AES_GCM_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.connection_config = { + 123: { # invalid outer key + "config_file": "should-fail" + } + } + + # 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.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.connection_config = { + "AES_GCM_V1": { + 123: "should-fail" # invalid inner key + } + } + + # Inner value is not a string + with pytest.raises(TypeError, match="All inner config keys/values must be str"): + config = pe.ExternalDecryptionConfiguration() + config.connection_config = { + "AES_GCM_V1": { + "config_file": ["not", "a", "string"] # invalid inner value + } + } \ No newline at end of file