wgtmac commented on code in PR #34616:
URL: https://github.com/apache/arrow/pull/34616#discussion_r1254530922
##########
cpp/src/arrow/dataset/file_parquet.cc:
##########
@@ -621,11 +639,38 @@ Result<std::shared_ptr<FileWriter>>
ParquetFileFormat::MakeWriter(
auto parquet_options =
checked_pointer_cast<ParquetFileWriteOptions>(options);
- std::unique_ptr<parquet::arrow::FileWriter> parquet_writer;
- ARROW_ASSIGN_OR_RAISE(parquet_writer, parquet::arrow::FileWriter::Open(
- *schema, default_memory_pool(),
destination,
- parquet_options->writer_properties,
-
parquet_options->arrow_writer_properties));
+ std::unique_ptr<parquet::arrow::FileWriter> parquet_writer = NULLPTR;
+
+#ifdef PARQUET_REQUIRE_ENCRYPTION
+ std::shared_ptr<DatasetEncryptionConfiguration> dataset_encrypt_config =
+ GetDatasetEncryptionConfig();
Review Comment:
```suggestion
auto dataset_encrypt_config = GetDatasetEncryptionConfig();
```
##########
cpp/src/arrow/dataset/file_parquet.cc:
##########
@@ -621,11 +639,38 @@ Result<std::shared_ptr<FileWriter>>
ParquetFileFormat::MakeWriter(
auto parquet_options =
checked_pointer_cast<ParquetFileWriteOptions>(options);
- std::unique_ptr<parquet::arrow::FileWriter> parquet_writer;
- ARROW_ASSIGN_OR_RAISE(parquet_writer, parquet::arrow::FileWriter::Open(
- *schema, default_memory_pool(),
destination,
- parquet_options->writer_properties,
-
parquet_options->arrow_writer_properties));
+ std::unique_ptr<parquet::arrow::FileWriter> parquet_writer = NULLPTR;
+
+#ifdef PARQUET_REQUIRE_ENCRYPTION
+ std::shared_ptr<DatasetEncryptionConfiguration> dataset_encrypt_config =
+ GetDatasetEncryptionConfig();
+
+ if (dataset_encrypt_config != nullptr) {
+ auto file_encryption_prop =
+ dataset_encrypt_config->crypto_factory->GetFileEncryptionProperties(
+ *dataset_encrypt_config->kms_connection_config,
+ *dataset_encrypt_config->encryption_config,
destination_locator.path,
+ destination_locator.filesystem);
+
+ auto writer_properties =
+ parquet::WriterProperties::Builder(*parquet_options->writer_properties)
+ .encryption(std::move(file_encryption_prop))
+ ->build();
+
+ ARROW_ASSIGN_OR_RAISE(
+ parquet_writer, parquet::arrow::FileWriter::Open(
+ *schema, writer_properties->memory_pool(),
destination,
+ writer_properties,
parquet_options->arrow_writer_properties));
+ }
+#endif
+
+ if (parquet_writer == NULLPTR) {
Review Comment:
```suggestion
if (parquet_writer == nullptr) {
```
nit: it would be good to be consistent to use `nullptr` in the source file.
##########
cpp/src/arrow/dataset/file_parquet.cc:
##########
@@ -621,11 +639,38 @@ Result<std::shared_ptr<FileWriter>>
ParquetFileFormat::MakeWriter(
auto parquet_options =
checked_pointer_cast<ParquetFileWriteOptions>(options);
- std::unique_ptr<parquet::arrow::FileWriter> parquet_writer;
- ARROW_ASSIGN_OR_RAISE(parquet_writer, parquet::arrow::FileWriter::Open(
- *schema, default_memory_pool(),
destination,
- parquet_options->writer_properties,
-
parquet_options->arrow_writer_properties));
+ std::unique_ptr<parquet::arrow::FileWriter> parquet_writer = NULLPTR;
+
+#ifdef PARQUET_REQUIRE_ENCRYPTION
+ std::shared_ptr<DatasetEncryptionConfiguration> dataset_encrypt_config =
+ GetDatasetEncryptionConfig();
+
+ if (dataset_encrypt_config != nullptr) {
Review Comment:
It will not be nullptr, right? We can simply remove this check.
##########
cpp/src/arrow/dataset/parquet_encryption_config.h:
##########
@@ -0,0 +1,70 @@
+// 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.
+
+#pragma once
+
+#include "arrow/dataset/type_fwd.h"
+#include "parquet/encryption/crypto_factory.h"
+#include "parquet/encryption/encryption.h"
+#include "parquet/encryption/kms_client.h"
+
+namespace arrow {
+namespace dataset {
+
+/// core class, that translates the parameters of high level encryption
+struct ARROW_DS_EXPORT DatasetEncryptionConfiguration {
+ DatasetEncryptionConfiguration()
+ : kms_connection_config(
+ std::make_shared<parquet::encryption::KmsConnectionConfig>()) {}
+
+ void Setup(
+ std::shared_ptr<parquet::encryption::CryptoFactory> crypto_factory,
+ std::shared_ptr<parquet::encryption::KmsConnectionConfig>
kms_connection_config,
+ std::shared_ptr<parquet::encryption::EncryptionConfiguration>
encryption_config) {
+ this->crypto_factory = std::move(crypto_factory);
+ this->kms_connection_config = std::move(kms_connection_config);
+ this->encryption_config = std::move(encryption_config);
+ }
+
+ std::shared_ptr<parquet::encryption::CryptoFactory> crypto_factory;
+ std::shared_ptr<parquet::encryption::KmsConnectionConfig>
kms_connection_config;
+ std::shared_ptr<parquet::encryption::EncryptionConfiguration>
encryption_config;
+};
Review Comment:
```suggestion
struct ARROW_DS_EXPORT DatasetEncryptionConfiguration {
std::shared_ptr<parquet::encryption::CryptoFactory> crypto_factory;
std::shared_ptr<parquet::encryption::KmsConnectionConfig>
kms_connection_config;
std::shared_ptr<parquet::encryption::EncryptionConfiguration>
encryption_config;
};
```
Have you considered to simply define it as the above? From you test cases,
it seems that not all parameters are required. So `Setup` function is not that
helpful if none of checks is done inside. By doing this, it enables aggregate
initialization in C++20 like this:
```
auto conf = {
.crypto_factory = xxx,
.kms_connection_config =
std::make_shared<parquet::encryption::KmsConnectionConfig>()
};
```
I am not saying we have to do this. Just a suggestion.
##########
cpp/src/arrow/dataset/dataset_encryption_test.cc:
##########
@@ -0,0 +1,373 @@
+// 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.
+
+#include <string_view>
+
+#include "gtest/gtest.h"
+
+#include "arrow/api.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/builder.h"
+#include "arrow/dataset/api.h"
+#include "arrow/dataset/parquet_encryption_config.h"
+#include "arrow/dataset/partition.h"
+#include "arrow/filesystem/mockfs.h"
+#include "arrow/io/api.h"
+#include "arrow/status.h"
+#include "arrow/table.h"
+#include "arrow/testing/gtest_util.h"
+#include "parquet/arrow/reader.h"
+#include "parquet/encryption/test_in_memory_kms.h"
+
+constexpr std::string_view kFooterKeyMasterKey = "0123456789012345";
+constexpr std::string_view kFooterKeyMasterKeyId = "footer_key";
+constexpr std::string_view kColumnMasterKeys[] = {"1234567890123450"};
+constexpr std::string_view kColumnMasterKeysIds[] = {"col_key"};
+constexpr std::string_view kBaseDir = "";
+const int kNumColumns = 1;
+namespace arrow {
Review Comment:
```suggestion
const int kNumColumns = 1;
namespace arrow {
```
##########
cpp/src/arrow/dataset/file_parquet.h:
##########
@@ -28,6 +28,7 @@
#include "arrow/dataset/discovery.h"
#include "arrow/dataset/file_base.h"
+#include "arrow/dataset/parquet_encryption_config.h"
Review Comment:
I think so, it is generally the convention to use forward declaration
whenever possible.
##########
cpp/src/arrow/dataset/CMakeLists.txt:
##########
@@ -175,6 +175,12 @@ endif()
if(ARROW_PARQUET)
add_arrow_dataset_test(file_parquet_test)
+ if(PARQUET_REQUIRE_ENCRYPTION AND ARROW_DATASET)
+ add_arrow_dataset_test(dataset_encryption_test
+ SOURCES
+
${PROJECT_SOURCE_DIR}/src/parquet/encryption/test_in_memory_kms.cc
Review Comment:
Thanks for confirming!
##########
cpp/src/arrow/dataset/file_parquet.cc:
##########
@@ -621,11 +639,38 @@ Result<std::shared_ptr<FileWriter>>
ParquetFileFormat::MakeWriter(
auto parquet_options =
checked_pointer_cast<ParquetFileWriteOptions>(options);
- std::unique_ptr<parquet::arrow::FileWriter> parquet_writer;
- ARROW_ASSIGN_OR_RAISE(parquet_writer, parquet::arrow::FileWriter::Open(
- *schema, default_memory_pool(),
destination,
- parquet_options->writer_properties,
-
parquet_options->arrow_writer_properties));
+ std::unique_ptr<parquet::arrow::FileWriter> parquet_writer = NULLPTR;
Review Comment:
```suggestion
std::unique_ptr<parquet::arrow::FileWriter> parquet_writer;
```
That's redundant.
##########
cpp/src/arrow/dataset/file_parquet.h:
##########
@@ -136,6 +137,40 @@ class ARROW_DS_EXPORT ParquetFileFormat : public
FileFormat {
fs::FileLocator destination_locator) const override;
std::shared_ptr<FileWriteOptions> DefaultWriteOptions() override;
+
+ /// \brief A getter function to retrieve the dataset encryption configuration
+ std::shared_ptr<DatasetEncryptionConfiguration> GetDatasetEncryptionConfig()
const {
+#ifdef PARQUET_REQUIRE_ENCRYPTION
+ return dataset_encryption_config_;
+#else
+ return NULLPTR;
+#endif
+ }
+ /// \brief A getter function to retrieve the dataset decryption configuration
+ std::shared_ptr<DatasetDecryptionConfiguration> GetDatasetDecryptionConfig()
const {
+#ifdef PARQUET_REQUIRE_ENCRYPTION
+ return dataset_decryption_config_;
+#else
+ return NULLPTR;
+#endif
Review Comment:
It seems that these two functions are always called in the scope of
`PARQUET_REQUIRE_ENCRYPTION`. So we can remove the branch inside, right?
##########
cpp/src/arrow/dataset/dataset_encryption_test.cc:
##########
@@ -0,0 +1,388 @@
+// 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.
+
+#include <string_view>
+
+#include "gtest/gtest.h"
+
+#include "arrow/api.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/builder.h"
+#include "arrow/dataset/api.h"
+#include "arrow/dataset/parquet_encryption_config.h"
+#include "arrow/dataset/partition.h"
+#include "arrow/filesystem/mockfs.h"
+#include "arrow/io/api.h"
+#include "arrow/status.h"
+#include "arrow/table.h"
+#include "arrow/testing/gtest_util.h"
+#include "parquet/arrow/reader.h"
+#include "parquet/encryption/test_in_memory_kms.h"
+
+constexpr std::string_view ds_footer_master_key = "0123456789012345";
+constexpr std::string_view ds_footer_master_key_id = "footer_key";
+constexpr std::string_view ds_column_master_keys[] = {"1234567890123450"};
+constexpr std::string_view ds_column_master_key_ids[] = {"col_key"};
+const int ds_num_columns = 1;
+namespace arrow {
+namespace dataset {
+
+class DatasetEncryptionTest : public ::testing::Test {
+ protected:
+ // setup the test
+ void SetUp() {}
+
+ // Create our parquetfileformat with encryption properties
+ std::shared_ptr<ParquetFileFormat> CreateFileFormat(
+ const std::string_view* column_ids, const std::string_view* column_keys,
+ int num_columns, std::string_view footer_id, std::string_view footer_key,
+ std::string_view footer_key_name = "footer_key",
+ std::string_view column_key_mapping = "col_key: a") {
+ auto key_list =
+ BuildKeyMap(column_ids, column_keys, num_columns, footer_id,
footer_key);
+
+ auto crypto_factory = SetupCryptoFactory(true, key_list);
+
+ auto encryption_config =
+ std::make_shared<::parquet::encryption::EncryptionConfiguration>(
+ std::string(footer_key_name));
+ encryption_config->column_keys = column_key_mapping;
+ if (footer_key_name.size() > 0) {
+ encryption_config->footer_key = footer_key_name;
+ }
+
+ // DatasetEncryptionConfiguration
+ DatasetEncryptionConfiguration dataset_encryption_config;
+ dataset_encryption_config.crypto_factory = crypto_factory;
+ dataset_encryption_config.encryption_config = encryption_config;
+
+ // DatasetDecryptionConfiguration
+ DatasetDecryptionConfiguration dataset_decryption_config;
+ dataset_decryption_config.crypto_factory = crypto_factory;
+
+ // create our Parquet file format object
+ auto file_format = std::make_shared<ParquetFileFormat>();
+
+ file_format->SetDatasetEncryptionConfig(
+
std::make_shared<DatasetEncryptionConfiguration>(dataset_encryption_config));
+ file_format->SetDatasetDecryptionConfig(
+
std::make_shared<DatasetDecryptionConfiguration>(dataset_decryption_config));
+
+ return file_format;
+ }
+
+ // utility to build the key map
+ std::unordered_map<std::string, std::string> BuildKeyMap(
+ const std::string_view* column_ids, const std::string_view* column_keys,
+ int num_columns, const std::string_view& footer_id,
+ const std::string_view& footer_key) {
+ std::unordered_map<std::string, std::string> key_map;
+ // add column keys
+ for (int i = 0; i < num_columns; i++) {
+ key_map.insert({std::string(column_ids[i]),
std::string(column_keys[i])});
+ }
+ // add footer key
+ key_map.insert({std::string(footer_id), std::string(footer_key)});
+
+ return key_map;
+ }
+
+ // A utility function to validate our files were written out */
+ void ValidateFilesExist(const
std::shared_ptr<arrow::fs::internal::MockFileSystem>& fs,
+ const std::vector<std::string>& files) {
+ for (const auto& file_path : files) {
+ ASSERT_OK_AND_ASSIGN(auto result, fs->GetFileInfo(file_path));
+
+ ASSERT_NE(result.type(), arrow::fs::FileType::NotFound);
+ }
+ }
+
+ // Build a dummy table
+ std::shared_ptr<::arrow::dataset::InMemoryDataset> BuildTable() {
+ // Create an Arrow Table
+ auto schema = arrow::schema(
+ {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64()),
+ arrow::field("c", arrow::int64()), arrow::field("part",
arrow::utf8())});
+
+ std::vector<std::shared_ptr<arrow::Array>> arrays(4);
+ arrow::NumericBuilder<arrow::Int64Type> builder;
+ ARROW_EXPECT_OK(builder.AppendValues({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}));
+ ARROW_EXPECT_OK(builder.Finish(&arrays[0]));
+ builder.Reset();
+ ARROW_EXPECT_OK(builder.AppendValues({9, 8, 7, 6, 5, 4, 3, 2, 1, 0}));
+ ARROW_EXPECT_OK(builder.Finish(&arrays[1]));
+ builder.Reset();
+ ARROW_EXPECT_OK(builder.AppendValues({1, 2, 1, 2, 1, 2, 1, 2, 1, 2}));
+ ARROW_EXPECT_OK(builder.Finish(&arrays[2]));
+ arrow::StringBuilder string_builder;
+ ARROW_EXPECT_OK(
+ string_builder.AppendValues({"a", "b", "c", "d", "e", "f", "g", "h",
"i", "j"}));
+ ARROW_EXPECT_OK(string_builder.Finish(&arrays[3]));
+
+ auto table = arrow::Table::Make(schema, arrays);
+
+ // Write it using Datasets
+ std::shared_ptr<::arrow::dataset::InMemoryDataset> dataset =
+ std::make_shared<::arrow::dataset::InMemoryDataset>(table);
+
+ return dataset;
+ }
+
+ // Helper function to create crypto factory and setup
+ std::shared_ptr<::parquet::encryption::CryptoFactory> SetupCryptoFactory(
+ bool wrap_locally, const std::unordered_map<std::string, std::string>&
key_list) {
+ auto crypto_factory =
std::make_shared<::parquet::encryption::CryptoFactory>();
+
+ std::shared_ptr<::parquet::encryption::KmsClientFactory>
kms_client_factory =
+
std::make_shared<::parquet::encryption::TestOnlyInMemoryKmsClientFactory>(
+ wrap_locally, key_list);
+
+ crypto_factory->RegisterKmsClientFactory(kms_client_factory);
+
+ return crypto_factory;
+ }
+
+ // Write dataset to disk with encryption
+ void WriteReadDatasetWithEncryption() {
+ auto file_format =
+ CreateFileFormat(ds_column_master_key_ids, ds_column_master_keys,
ds_num_columns,
+ ds_footer_master_key_id, ds_footer_master_key);
+
+ // create our mock file system
+ ::arrow::fs::TimePoint mock_now = std::chrono::system_clock::now();
+ ASSERT_OK_AND_ASSIGN(auto file_system,
+ ::arrow::fs::internal::MockFileSystem::Make(mock_now,
{}));
+ // create filesystem
+ ASSERT_OK(file_system->CreateDir(""));
+
+ auto mock_fs =
+
std::dynamic_pointer_cast<::arrow::fs::internal::MockFileSystem>(file_system);
+
+ auto partition_schema = ::arrow::schema({::arrow::field("part",
::arrow::utf8())});
+ auto partitioning =
+ std::make_shared<::arrow::dataset::HivePartitioning>(partition_schema);
+
+ // ----- Write the Dataset ----
+ auto dataset_out = BuildTable();
+ ASSERT_OK_AND_ASSIGN(auto scanner_builder_out, dataset_out->NewScan());
+ ASSERT_OK_AND_ASSIGN(auto scanner_out, scanner_builder_out->Finish());
+
+ ::arrow::dataset::FileSystemDatasetWriteOptions write_options;
+ write_options.file_write_options = file_format->DefaultWriteOptions();
+ write_options.filesystem = file_system;
+ write_options.base_dir = "";
+ write_options.partitioning = partitioning;
+ write_options.basename_template = "part{i}.parquet";
+ ASSERT_OK(::arrow::dataset::FileSystemDataset::Write(write_options,
scanner_out));
+
+ std::vector<std::string> files = {"part=a/part0.parquet",
"part=b/part0.parquet",
Review Comment:
Thanks for the explanation! Would you mind adding these as a comment?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
