File size: 3,393 Bytes
1380717 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# 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 shutil
import os
from datetime import timedelta
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet.encryption as pe
from pyarrow.tests.parquet.encryption import InMemoryKmsClient
""" A sample to demonstrate parquet dataset encryption and decryption"""
# create a list of dictionaries that will represent our dataset
table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021],
'n_legs': [2, 2, 4, 4, 5, 100],
'animal': ["Flamingo", "Parrot", "Dog", "Horse",
"Brittle stars", "Centipede"]})
# create a PyArrow dataset from the table
dataset = ds.dataset(table)
FOOTER_KEY = b"0123456789112345"
FOOTER_KEY_NAME = "footer_key"
COL_KEY = b"1234567890123450"
COL_KEY_NAME = "col_key"
encryption_config = pe.EncryptionConfiguration(
footer_key=FOOTER_KEY_NAME,
plaintext_footer=False,
# Use COL_KEY_NAME to encrypt `n_legs` and `animal` columns.
column_keys={
COL_KEY_NAME: ["n_legs", "animal"],
},
encryption_algorithm="AES_GCM_V1",
# requires timedelta or an assertion is raised
cache_lifetime=timedelta(minutes=5.0),
data_key_length_bits=256)
kms_connection_config = pe.KmsConnectionConfig(
custom_kms_conf={
FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"),
COL_KEY_NAME: COL_KEY.decode("UTF-8"),
}
)
decryption_config = pe.DecryptionConfiguration(cache_lifetime=300)
def kms_factory(kms_connection_configuration):
return InMemoryKmsClient(kms_connection_configuration)
crypto_factory = pe.CryptoFactory(kms_factory)
parquet_encryption_cfg = ds.ParquetEncryptionConfig(
crypto_factory, kms_connection_config, encryption_config)
parquet_decryption_cfg = ds.ParquetDecryptionConfig(crypto_factory,
kms_connection_config,
decryption_config)
# set encryption config for parquet fragment scan options
pq_scan_opts = ds.ParquetFragmentScanOptions()
pq_scan_opts.parquet_decryption_config = parquet_decryption_cfg
pformat = pa.dataset.ParquetFileFormat(default_fragment_scan_options=pq_scan_opts)
if os.path.exists('sample_dataset'):
shutil.rmtree('sample_dataset')
write_options = pformat.make_write_options(
encryption_config=parquet_encryption_cfg)
ds.write_dataset(data=dataset, base_dir="sample_dataset",
partitioning=['year'], format=pformat, file_options=write_options)
# read the dataset back
dataset = ds.dataset('sample_dataset', format=pformat)
# print the dataset
print(dataset.to_table())
|