diff --git a/python/pyarrow/_parquet_encryption.pxd b/python/pyarrow/_parquet_encryption.pxd index d52669501a40..ab14d842aeb0 100644 --- a/python/pyarrow/_parquet_encryption.pxd +++ b/python/pyarrow/_parquet_encryption.pxd @@ -38,6 +38,10 @@ cdef class EncryptionConfiguration(_Weakrefable): cdef shared_ptr[CEncryptionConfiguration] configuration cdef inline shared_ptr[CEncryptionConfiguration] unwrap(self) nogil +cdef class ExternalEncryptionConfiguration(EncryptionConfiguration): + cdef shared_ptr[CExternalEncryptionConfiguration] external_configuration + cdef inline shared_ptr[CExternalEncryptionConfiguration] unwrap_external(self) nogil + cdef class DecryptionConfiguration(_Weakrefable): cdef shared_ptr[CDecryptionConfiguration] configuration cdef inline shared_ptr[CDecryptionConfiguration] unwrap(self) nogil @@ -54,3 +58,4 @@ cdef shared_ptr[CCryptoFactory] pyarrow_unwrap_cryptofactory(object crypto_facto 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 * diff --git a/python/pyarrow/_parquet_encryption.pyx b/python/pyarrow/_parquet_encryption.pyx index 81bd421dcbce..00964e1d5a4f 100644 --- a/python/pyarrow/_parquet_encryption.pyx +++ b/python/pyarrow/_parquet_encryption.pyx @@ -27,6 +27,8 @@ from pyarrow.includes.libarrow cimport * from pyarrow.lib cimport _Weakrefable from pyarrow.lib import tobytes, frombytes +import json + cdef ParquetCipher cipher_from_name(name): name = name.upper() @@ -174,6 +176,130 @@ cdef class EncryptionConfiguration(_Weakrefable): cdef inline shared_ptr[CEncryptionConfiguration] unwrap(self) nogil: return self.configuration +cdef class ExternalEncryptionConfiguration(EncryptionConfiguration): + """ExternalEncryptionConfiguration is a Cython extension class that inherits from EncryptionConfiguration.""" + __slots__ = () + + def __init__(self, footer_key, *, column_keys=None, + encryption_algorithm=None, + plaintext_footer=None, double_wrapping=None, + cache_lifetime=None, internal_key_material=None, + data_key_length_bits=None, per_column_encryption=None, + app_context=None, connection_config=None): + + super().__init__(footer_key, + column_keys=column_keys, + encryption_algorithm=encryption_algorithm, + plaintext_footer=plaintext_footer, + double_wrapping=double_wrapping, + cache_lifetime=cache_lifetime, + internal_key_material=internal_key_material, + data_key_length_bits=data_key_length_bits) + + self.external_configuration.reset( + new CExternalEncryptionConfiguration(tobytes(footer_key))) + + if app_context is not None: + self.app_context = app_context + if connection_config is not None: + self.connection_config = connection_config + if per_column_encryption is not None: + self.per_column_encryption = per_column_encryption + + @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.""" + result = {} + + for pair in self.external_configuration.get().connection_config: + result[frombytes(pair.first)] = frombytes(pair.second) + + 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[c_string, c_string] cpp_map + for k, v in value.items(): + if not isinstance(k, str): + raise TypeError(f"Connection config key must be str, got {type(k).__name__}") + if not isinstance(v, str): + raise TypeError(f"Connection config value must be str, got {type(v).__name__}") + cpp_map[tobytes(k)] = tobytes(v) + + self.external_configuration.get().connection_config = cpp_map + + @property + def per_column_encryption(self): + """Get the per_column_encryption as a Python dictionary.""" + + py_dict = {} + + for pair in self.external_configuration.get().per_column_encryption: + py_dict[frombytes(pair.first)] = { + "encryption_algorithm": cipher_to_name(pair.second.parquet_cipher), + "encryption_key": frombytes(pair.second.key_id) + } + + return py_dict + + @per_column_encryption.setter + def per_column_encryption(self, dict py_column_encryption): + """Set the per_column_encryption from a Python dictionary.""" + if py_column_encryption is None: + raise TypeError("per_column_encryption cannot be None") + + # Clear the existing C++ map first + self.external_configuration.get().per_column_encryption.clear() + + cdef CColumnEncryptionAttributes cpp_attrs + # Iterate over the Python dictionary + for py_key, py_attrs in py_column_encryption.items(): + if not isinstance(py_key, str) or not isinstance(py_attrs, dict): + raise TypeError("column_encryption keys must be strings and values must be dictionaries.") + + # Convert encryption_algorithm string to C++ ParquetCipher enum + if "encryption_algorithm" not in py_attrs or not isinstance(py_attrs["encryption_algorithm"], str): + raise ValueError("Each column must have 'encryption_algorithm' (string).") + + # Convert encryption_key string to C++ c_string + if "encryption_key" not in py_attrs or not isinstance(py_attrs["encryption_key"], str): + raise ValueError("Each column must have 'encryption_key' (string).") + + cpp_attrs.parquet_cipher = cipher_from_name(py_attrs["encryption_algorithm"]) + cpp_attrs.key_id = tobytes(py_attrs["encryption_key"]) + + self.external_configuration.get().per_column_encryption[tobytes(py_key)] = cpp_attrs + + cdef inline shared_ptr[CExternalEncryptionConfiguration] unwrap_external(self) nogil: + return self.external_configuration + cdef class DecryptionConfiguration(_Weakrefable): """Configuration of the decryption, such as cache timeout.""" @@ -477,6 +603,10 @@ cdef shared_ptr[CEncryptionConfiguration] pyarrow_unwrap_encryptionconfig(object return ( encryptionconfig).unwrap() raise TypeError("Expected EncryptionConfiguration, got %s" % type(encryptionconfig)) +cdef shared_ptr[CExternalEncryptionConfiguration] pyarrow_unwrap_external_encryptionconfig(object encryptionconfig) except *: + if isinstance(encryptionconfig, ExternalEncryptionConfiguration): + return ( encryptionconfig).unwrap_external() + raise TypeError("Expected ExternalEncryptionConfiguration, got %s" % type(encryptionconfig)) cdef shared_ptr[CDecryptionConfiguration] pyarrow_unwrap_decryptionconfig(object decryptionconfig) except *: if isinstance(decryptionconfig, DecryptionConfiguration): diff --git a/python/pyarrow/includes/libparquet_encryption.pxd b/python/pyarrow/includes/libparquet_encryption.pxd index 2b40414ce538..a420707ba74f 100644 --- a/python/pyarrow/includes/libparquet_encryption.pxd +++ b/python/pyarrow/includes/libparquet_encryption.pxd @@ -79,6 +79,27 @@ cdef extern from "parquet/encryption/crypto_factory.h" \ c_bool internal_key_material int32_t data_key_length_bits + cdef cppclass CColumnEncryptionAttributes\ + " parquet::encryption::ColumnEncryptionAttributes": + CColumnEncryptionAttributes() except + + ParquetCipher parquet_cipher + c_string key_id + + cdef cppclass CExternalEncryptionConfiguration\ + " parquet::encryption::ExternalEncryptionConfiguration": + CExternalEncryptionConfiguration(const c_string& footer_key) except + + c_string footer_key + c_string column_keys + ParquetCipher encryption_algorithm + c_bool plaintext_footer + c_bool double_wrapping + double cache_lifetime_seconds + c_bool internal_key_material + int32_t data_key_length_bits + unordered_map[c_string, CColumnEncryptionAttributes] per_column_encryption + c_string app_context + unordered_map[c_string, c_string] connection_config + cdef cppclass CDecryptionConfiguration\ " parquet::encryption::DecryptionConfiguration": CDecryptionConfiguration() except + @@ -127,4 +148,4 @@ cdef extern from "arrow/python/parquet_encryption.h" \ CResult[shared_ptr[CFileDecryptionProperties]] \ SafeGetFileDecryptionProperties( const CKmsConnectionConfig& kms_connection_config, - const CDecryptionConfiguration& decryption_config) + const CDecryptionConfiguration& decryption_config) \ No newline at end of file diff --git a/python/pyarrow/parquet/encryption.py b/python/pyarrow/parquet/encryption.py index df6eed913fa5..af83738db7e3 100644 --- a/python/pyarrow/parquet/encryption.py +++ b/python/pyarrow/parquet/encryption.py @@ -17,6 +17,7 @@ # specific language governing permissions and limitations # under the License. from pyarrow._parquet_encryption import (CryptoFactory, # noqa + ExternalEncryptionConfiguration, EncryptionConfiguration, DecryptionConfiguration, KmsConnectionConfig, diff --git a/python/pyarrow/tests/parquet/test_external_encryption.py b/python/pyarrow/tests/parquet/test_external_encryption.py new file mode 100644 index 000000000000..5f9b6bf22aad --- /dev/null +++ b/python/pyarrow/tests/parquet/test_external_encryption.py @@ -0,0 +1,173 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import pytest +from datetime import timedelta +import pyarrow.parquet as pq +import pyarrow.parquet.encryption as pe + +def test_encryption_configuration_properties(tempdir): + """Test the standard EncryptionConfiguration properties to avoid regressions.""" + + config = pe.EncryptionConfiguration( + footer_key="footer-key-name", + column_keys={ + "col-key-id": ["a", "b"], + }, + 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 + ) + + assert isinstance(config, pe.EncryptionConfiguration) + + assert config.footer_key == "footer-key-name" + assert config.column_keys == { + "col-key-id": ["a", "b"] + } + assert config.encryption_algorithm == "AES_GCM_V1" + assert config.plaintext_footer is True + assert config.double_wrapping is True + assert config.cache_lifetime == timedelta(minutes=5.0) + assert config.internal_key_material is True + assert config.data_key_length_bits == 256 + +def test_external_encryption_configuration_properties(tempdir): + """Test the ExternalEncryptionConfiguration properties including external-specific fields.""" + + config_external = pe.ExternalEncryptionConfiguration( + footer_key="footer-key-name", + column_keys={ + "col-key-id": ["a", "b"], + }, + 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={ + "a": { + "encryption_algorithm": "AES_GCM_V1", + "encryption_key": "key_1" + }, + "b": { + "encryption_algorithm": "AES_GCM_CTR_V1", + "encryption_key": "key_n" + } + }, + app_context={ + "user_id": "Picard1701", + "location": "Presidio" + }, + connection_config={ + "config_file": "path/to/config/file", + "config_file_decryption_key": "some_key" + } + ) + + assert isinstance(config_external, pe.ExternalEncryptionConfiguration) + + assert config_external.footer_key == "footer-key-name" + assert config_external.column_keys == { + "col-key-id": ["a", "b"] + } + assert config_external.encryption_algorithm == "AES_GCM_V1" + assert config_external.plaintext_footer is True + assert config_external.double_wrapping is True + assert config_external.cache_lifetime == timedelta(minutes=5.0) + assert config_external.internal_key_material is True + assert config_external.data_key_length_bits == 256 + + assert config_external.app_context == { + "user_id": "Picard1701", + "location": "Presidio" + } + + assert config_external.connection_config == { + "config_file": "path/to/config/file", + "config_file_decryption_key": "some_key" + } + + assert config_external.per_column_encryption == { + "a": { + "encryption_algorithm": "AES_GCM_V1", + "encryption_key": "key_1" + }, + "b": { + "encryption_algorithm": "AES_GCM_CTR_V1", + "encryption_key": "key_n" + } + } + +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( + 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( + footer_key="key", + per_column_encryption={ + "a": { + "encryption_algorithm": "INVALID", + "encryption_key": "some_key" + } + } + ) + +def test_external_encryption_connection_config_invalid_types(): + """Ensure connection_config rejects non-string keys or values.""" + with pytest.raises(TypeError, match="Connection config key must be str, got int"): + config=pe.ExternalEncryptionConfiguration( + footer_key="key" + ) + config.connection_config={ + "config_file": "path/to/file", + 123: "should-fail" # Invalid: key is not a string + } + + with pytest.raises(TypeError, match="Connection config value must be str, got list"): + config = pe.ExternalEncryptionConfiguration( + footer_key="key" + ) + config.connection_config={ + "config_file": ["not", "a", "string"] # Invalid: value is not a string + } + +def test_external_encryption_rejects_none_values(): + config = pe.ExternalEncryptionConfiguration(footer_key="key") + + # per_column_encryption: expect ValueError + with pytest.raises(TypeError, match="per_column_encryption cannot be None"): + config.per_column_encryption = None + + # app_context: expect ValueError due to None not being JSON-serializable + with pytest.raises(ValueError, match="app_context must be JSON-serializable"): + config.app_context = None + + # connection_config: expect ValueError due to None not being iterable + with pytest.raises(ValueError, match="Connection config value cannot be None"): + config.connection_config = None \ No newline at end of file