使用PyArrow 23.0.1生成带AAD前缀的加密Parquet文件报错求助

问题描述 #

使用PyArrow 23.0.1生成带AAD前缀的加密Parquet文件时，运行脚本出现如下错误：

Traceback (most recent call last):
  File "..../encrypted-parquet.py", line 65, in <module>
    enc_config = EncryptionConfiguration(
                 ^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/_parquet_encryption.pyx", line 57, in pyarrow._parquet_encryption.EncryptionConfiguration.__init__
TypeError: __init__() got an unexpected keyword argument 'aad_prefix'

问题原因 #

PyArrow 23.0.1版本的EncryptionConfiguration类并未提供aad_prefix参数，该参数需要通过CryptoFactory.file_encryption_properties方法直接传入，而非放在加密配置类中。

修改后的完整代码 #

import pyarrow as pa
import pyarrow.parquet as pq
from pyarrow.parquet.encryption import (
    CryptoFactory,
    KmsConnectionConfig,
    EncryptionConfiguration,
    KmsClient
)
import base64

# Sample data
data = pa.table({
    "id": [1, 2, 3],
    "name": ["Mark", "Mary", "Mike"],
    "salary": [10001, 45345, 34345]
})

KEY = base64.b64encode(b"0123456789abcdef0123456789abcdef").decode("utf-8")

# Minimal KMS Client
class SimpleKmsClient(KmsClient):
    def __init__(self, conf):
        super().__init__()
        self.conf = conf

    def wrap_key(self, key_bytes, master_key_identifier):
        return key_bytes

    def unwrap_key(self, wrapped_key, master_key_identifier):
        return base64.b64decode(self.conf[master_key_identifier])

# KMS factory
def kms_client_factory(kms_connection_config):
    return SimpleKmsClient({
        "footer_key": KEY,
        "col_key": KEY
    })

crypto_factory = CryptoFactory(kms_client_factory)

# Encryption config - 移除aad_prefix参数
enc_config = EncryptionConfiguration(
    footer_key="footer_key",
    column_keys={"col_key": ["salary"]},
    plaintext_footer=True,
    double_wrapping=False,
    internal_key_material=True
)

kms_conf = KmsConnectionConfig(kms_instance_id="local")

# 在file_encryption_properties中添加aad_prefix参数
file_encryption_props = crypto_factory.file_encryption_properties(
    kms_conf,
    enc_config,
    aad_prefix="test_prefix"
)

# Write file
pq.write_table(
    data,
    "encrypt_columns_plaintext_footer.parquet",
    encryption_properties=file_encryption_props
)

print("File written")

关键修改点 #

• 从EncryptionConfiguration的初始化参数中移除了aad_prefix="test_prefix"
• 在调用crypto_factory.file_encryption_properties时，新增aad_prefix="test_prefix"参数，这是PyArrow 23.0.1版本中设置AAD前缀的正确方式

内容的提问来源于stack exchange，提问作者Mandroid