ARROW-9947: [Python] High-level Python API for Parquet encryption of files.

.github/workflows/python.yml

@@ -126,6 +126,7 @@ jobs:
       ARROW_JEMALLOC: ON
       ARROW_ORC: ON
       ARROW_PARQUET: ON
+      PARQUET_REQUIRE_ENCRYPTION: ON
       ARROW_PLASMA: ON
       ARROW_PYTHON: ON
       ARROW_S3: ON

ci/appveyor-cpp-build.bat

@@ -146,6 +146,7 @@ set PYARROW_WITH_DATASET=ON
 set PYARROW_WITH_FLIGHT=%ARROW_BUILD_FLIGHT%
 set PYARROW_WITH_GANDIVA=%ARROW_BUILD_GANDIVA%
 set PYARROW_WITH_PARQUET=ON
+set PYARROW_WITH_PARQUET_ENCRYPTION=ON
 set PYARROW_WITH_S3=%ARROW_S3%
 set PYARROW_WITH_STATIC_BOOST=ON
 

ci/docker/conda-python-hdfs.dockerfile

@@ -46,6 +46,7 @@ ENV CC=gcc \
     ARROW_GANDIVA=OFF \
     ARROW_PLASMA=OFF \
     ARROW_PARQUET=ON \
+    PARQUET_REQUIRE_ENCRYPTION=ON \
     ARROW_ORC=OFF \
     ARROW_HDFS=ON \
     ARROW_PYTHON=ON \

ci/scripts/python_build.sh

@@ -57,6 +57,7 @@ export PYARROW_WITH_FLIGHT=${ARROW_FLIGHT:-OFF}
 export PYARROW_WITH_PLASMA=${ARROW_PLASMA:-OFF}
 export PYARROW_WITH_GANDIVA=${ARROW_GANDIVA:-OFF}
 export PYARROW_WITH_PARQUET=${ARROW_PARQUET:-OFF}
+export PYARROW_WITH_PARQUET_ENCRYPTION=${ARROW_PARQUET:-OFF}
 export PYARROW_WITH_DATASET=${ARROW_DATASET:-OFF}
 
 export PYARROW_PARALLEL=${n_jobs}

ci/scripts/python_sdist_test.sh

@@ -35,6 +35,7 @@ export PYARROW_WITH_FLIGHT=${ARROW_FLIGHT:-OFF}
 export PYARROW_WITH_PLASMA=${ARROW_PLASMA:-OFF}
 export PYARROW_WITH_GANDIVA=${ARROW_GANDIVA:-OFF}
 export PYARROW_WITH_PARQUET=${ARROW_PARQUET:-OFF}
+export PYARROW_WITH_PARQUET_ENCRYPTION=${ARROW_PARQUET:-OFF}
 export PYARROW_WITH_DATASET=${ARROW_DATASET:-OFF}
 
 # TODO: Users should not require ARROW_HOME and pkg-config to find Arrow C++.

ci/scripts/python_wheel_macos_build.sh

@@ -70,6 +70,7 @@ echo "=== (${PYTHON_VERSION}) Building Arrow C++ libraries ==="
 : ${ARROW_MIMALLOC:=ON}
 : ${ARROW_ORC:=ON}
 : ${ARROW_PARQUET:=ON}
+: ${PARQUET_REQUIRE_ENCRYPTION:=ON}
 : ${ARROW_PLASMA:=ON}
 : ${ARROW_S3:=ON}
 : ${ARROW_SIMD_LEVEL:="SSE4_2"}
@@ -106,6 +107,7 @@ cmake \
     -DARROW_ORC=${ARROW_ORC} \
     -DARROW_PACKAGE_KIND="python-wheel-macos" \
     -DARROW_PARQUET=${ARROW_PARQUET} \
+    -DPARQUET_REQUIRE_ENCRYPTION=${PARQUET_REQUIRE_ENCRYPTION}
     -DARROW_PLASMA=${ARROW_PLASMA} \
     -DARROW_PYTHON=ON \
     -DARROW_RPATH_ORIGIN=ON \
@@ -147,6 +149,7 @@ export PYARROW_WITH_GANDIVA=${ARROW_GANDIVA}
 export PYARROW_WITH_HDFS=${ARROW_HDFS}
 export PYARROW_WITH_ORC=${ARROW_ORC}
 export PYARROW_WITH_PARQUET=${ARROW_PARQUET}
+export PYARROW_WITH_PARQUET_ENCRYPTION=${PARQUET_REQUIRE_ENCRYPTION}
 export PYARROW_WITH_PLASMA=${ARROW_PLASMA}
 export PYARROW_WITH_S3=${ARROW_S3}
 export PYARROW_CMAKE_OPTIONS="-DCMAKE_OSX_ARCHITECTURES=${CMAKE_OSX_ARCHITECTURES} -DARROW_SIMD_LEVEL=${ARROW_SIMD_LEVEL}"

ci/scripts/python_wheel_manylinux_build.sh

@@ -57,6 +57,7 @@ echo "=== (${PYTHON_VERSION}) Building Arrow C++ libraries ==="
 : ${ARROW_MIMALLOC:=ON}
 : ${ARROW_ORC:=ON}
 : ${ARROW_PARQUET:=ON}
+: ${PARQUET_REQUIRE_ENCRYPTION:=ON}
 : ${ARROW_PLASMA:=ON}
 : ${ARROW_S3:=ON}
 : ${ARROW_TENSORFLOW:=ON}
@@ -100,6 +101,7 @@ cmake \
     -DARROW_ORC=${ARROW_ORC} \
     -DARROW_PACKAGE_KIND="python-wheel-manylinux${MANYLINUX_VERSION}" \
     -DARROW_PARQUET=${ARROW_PARQUET} \
+    -DPARQUET_REQUIRE_ENCRYPTION=${PARQUET_REQUIRE_ENCRYPTION} \
     -DARROW_PLASMA=${ARROW_PLASMA} \
     -DARROW_PYTHON=ON \
     -DARROW_RPATH_ORIGIN=ON \
@@ -143,6 +145,7 @@ export PYARROW_WITH_GANDIVA=${ARROW_GANDIVA}
 export PYARROW_WITH_HDFS=${ARROW_HDFS}
 export PYARROW_WITH_ORC=${ARROW_ORC}
 export PYARROW_WITH_PARQUET=${ARROW_PARQUET}
+export PYARROW_WITH_PARQUET_ENCRYPTION=${PARQUET_REQUIRE_ENCRYPTION}
 export PYARROW_WITH_PLASMA=${ARROW_PLASMA}
 export PYARROW_WITH_S3=${ARROW_S3}
 # PyArrow build configuration

ci/scripts/python_wheel_windows_build.bat

@@ -36,6 +36,7 @@ set ARROW_GANDIVA=OFF
 set ARROW_HDFS=ON
 set ARROW_ORC=OFF
 set ARROW_PARQUET=ON
+set PARQUET_REQUIRE_ENCRYPTION=ON
 set ARROW_MIMALLOC=ON
 set ARROW_S3=ON
 set ARROW_TENSORFLOW=ON
@@ -68,6 +69,7 @@ cmake ^
     -DARROW_ORC=%ARROW_ORC% ^
     -DARROW_PACKAGE_KIND="python-wheel-windows" ^
     -DARROW_PARQUET=%ARROW_PARQUET% ^
+    -DPARQUET_REQUIRE_ENCRYPTION=%PARQUET_REQUIRE_ENCRYPTION% ^
     -DARROW_PYTHON=ON ^
     -DARROW_S3=%ARROW_S3% ^
     -DARROW_TENSORFLOW=%ARROW_TENSORFLOW% ^
@@ -102,6 +104,7 @@ set PYARROW_WITH_GANDIVA=%ARROW_GANDIVA%
 set PYARROW_WITH_HDFS=%ARROW_HDFS%
 set PYARROW_WITH_ORC=%ARROW_ORC%
 set PYARROW_WITH_PARQUET=%ARROW_PARQUET%
+set PYARROW_WITH_PARQUET_ENCRYPTION=%PARQUET_REQUIRE_ENCRYPTION%
 set PYARROW_WITH_S3=%ARROW_S3%
 set ARROW_HOME=C:\arrow-dist
 

cpp/src/arrow/python/CMakeLists.txt

@@ -57,6 +57,10 @@ if(ARROW_FILESYSTEM)
   list(APPEND ARROW_PYTHON_SRCS filesystem.cc)
 endif()
 
+if(PARQUET_REQUIRE_ENCRYPTION)
+  list(APPEND ARROW_PYTHON_SRCS parquet_encryption.cc)
+endif()
+
 set(ARROW_PYTHON_DEPENDENCIES arrow_dependencies)
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang" OR CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
@@ -70,6 +74,10 @@ if(WIN32)
   list(APPEND ARROW_PYTHON_SHARED_LINK_LIBS ${PYTHON_LIBRARIES} ${PYTHON_OTHER_LIBS})
 endif()
 
+if(PARQUET_REQUIRE_ENCRYPTION)
+  list(APPEND ARROW_PYTHON_SHARED_LINK_LIBS parquet_shared)
+endif()
+
 set(ARROW_PYTHON_INCLUDES ${NUMPY_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS})
 
 add_arrow_lib(arrow_python

cpp/src/arrow/python/parquet_encryption.cc

@@ -0,0 +1,98 @@
+// 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 "arrow/python/parquet_encryption.h"
+#include "parquet/exception.h"
+
+namespace arrow {
+namespace py {
+namespace parquet {
+namespace encryption {
+
+PyKmsClient::PyKmsClient(PyObject* handler, PyKmsClientVtable vtable)
+    : handler_(handler), vtable_(std::move(vtable)) {
+  Py_INCREF(handler);
+}
+
+PyKmsClient::~PyKmsClient() {}
+
+std::string PyKmsClient::WrapKey(const std::string& key_bytes,
+                                 const std::string& master_key_identifier) {
+  std::string wrapped;
+  auto st = SafeCallIntoPython([&]() -> Status {
+    vtable_.wrap_key(handler_.obj(), key_bytes, master_key_identifier, &wrapped);
+    return CheckPyError();
+  });
+  if (!st.ok()) {
+    throw ::parquet::ParquetStatusException(st);
+  }
+  return wrapped;
+}
+
+std::string PyKmsClient::UnwrapKey(const std::string& wrapped_key,
+                                   const std::string& master_key_identifier) {
+  std::string unwrapped;
+  auto st = SafeCallIntoPython([&]() -> Status {
+    vtable_.unwrap_key(handler_.obj(), wrapped_key, master_key_identifier, &unwrapped);
+    return CheckPyError();
+  });
+  if (!st.ok()) {
+    throw ::parquet::ParquetStatusException(st);
+  }
+  return unwrapped;
+}
+
+PyKmsClientFactory::PyKmsClientFactory(PyObject* handler, PyKmsClientFactoryVtable vtable)
+    : handler_(handler), vtable_(std::move(vtable)) {
+  Py_INCREF(handler);
+}
+
+PyKmsClientFactory::~PyKmsClientFactory() {}
+
+std::shared_ptr<::parquet::encryption::KmsClient> PyKmsClientFactory::CreateKmsClient(
+    const ::parquet::encryption::KmsConnectionConfig& kms_connection_config) {
+  std::shared_ptr<::parquet::encryption::KmsClient> kms_client;
+  auto st = SafeCallIntoPython([&]() -> Status {
+    vtable_.create_kms_client(handler_.obj(), kms_connection_config, &kms_client);
+    return CheckPyError();
+  });
+  if (!st.ok()) {
+    throw ::parquet::ParquetStatusException(st);
+  }
+  return kms_client;
+}
+
+arrow::Result<std::shared_ptr<::parquet::FileEncryptionProperties>>
+PyCryptoFactory::SafeGetFileEncryptionProperties(
+    const ::parquet::encryption::KmsConnectionConfig& kms_connection_config,
+    const ::parquet::encryption::EncryptionConfiguration& encryption_config) {
+  PARQUET_CATCH_AND_RETURN(
+      this->GetFileEncryptionProperties(kms_connection_config, encryption_config));
+}
+
+arrow::Result<std::shared_ptr<::parquet::FileDecryptionProperties>>
+PyCryptoFactory::SafeGetFileDecryptionProperties(
+    const ::parquet::encryption::KmsConnectionConfig& kms_connection_config,
+    const ::parquet::encryption::DecryptionConfiguration& decryption_config) {
+  PARQUET_CATCH_AND_RETURN(
+      this->GetFileDecryptionProperties(kms_connection_config, decryption_config));
+}
+
+}  // namespace encryption
+}  // namespace parquet
+}  // namespace py
+}  // namespace arrow

cpp/src/arrow/python/parquet_encryption.h

@@ -0,0 +1,109 @@
+// 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 <string>
+
+#include "arrow/python/common.h"
+#include "arrow/python/visibility.h"
+#include "arrow/util/macros.h"
+#include "parquet/encryption/crypto_factory.h"
+#include "parquet/encryption/kms_client.h"
+#include "parquet/encryption/kms_client_factory.h"
+
+namespace arrow {
+namespace py {
+namespace parquet {
+namespace encryption {
+
+/// \brief A table of function pointers for calling from C++ into
+/// Python.
+class ARROW_PYTHON_EXPORT PyKmsClientVtable {
+ public:
+  std::function<void(PyObject*, const std::string& key_bytes,
+                     const std::string& master_key_identifier, std::string* out)>
+      wrap_key;
+  std::function<void(PyObject*, const std::string& wrapped_key,
+                     const std::string& master_key_identifier, std::string* out)>
+      unwrap_key;
+};
+
+/// \brief A helper for KmsClient implementation in Python.
+class ARROW_PYTHON_EXPORT PyKmsClient : public ::parquet::encryption::KmsClient {
+ public:
+  PyKmsClient(PyObject* handler, PyKmsClientVtable vtable);
+  ~PyKmsClient() override;
+
+  std::string WrapKey(const std::string& key_bytes,
+                      const std::string& master_key_identifier) override;
+
+  std::string UnwrapKey(const std::string& wrapped_key,
+                        const std::string& master_key_identifier) override;
+
+ private:
+  OwnedRefNoGIL handler_;
+  PyKmsClientVtable vtable_;
+};
+
+/// \brief A table of function pointers for calling from C++ into
+/// Python.
+class ARROW_PYTHON_EXPORT PyKmsClientFactoryVtable {
+ public:
+  std::function<void(
+      PyObject*, const ::parquet::encryption::KmsConnectionConfig& kms_connection_config,
+      std::shared_ptr<::parquet::encryption::KmsClient>* out)>
+      create_kms_client;
+};
+
+/// \brief A helper for KmsClientFactory implementation in Python.
+class ARROW_PYTHON_EXPORT PyKmsClientFactory
+    : public ::parquet::encryption::KmsClientFactory {
+ public:
+  PyKmsClientFactory(PyObject* handler, PyKmsClientFactoryVtable vtable);
+  ~PyKmsClientFactory() override;
+
+  std::shared_ptr<::parquet::encryption::KmsClient> CreateKmsClient(
+      const ::parquet::encryption::KmsConnectionConfig& kms_connection_config) override;
+
+ private:
+  OwnedRefNoGIL handler_;
+  PyKmsClientFactoryVtable vtable_;
+};
+
+/// \brief A CryptoFactory that returns Results instead of throwing exceptions.
+class ARROW_PYTHON_EXPORT PyCryptoFactory : public ::parquet::encryption::CryptoFactory {
+ public:
+  arrow::Result<std::shared_ptr<::parquet::FileEncryptionProperties>>
+  SafeGetFileEncryptionProperties(
+      const ::parquet::encryption::KmsConnectionConfig& kms_connection_config,
+      const ::parquet::encryption::EncryptionConfiguration& encryption_config);
+
+  /// The returned FileDecryptionProperties object will use the cache inside this
+  /// CryptoFactory object, so please keep this
+  /// CryptoFactory object alive along with the returned
+  /// FileDecryptionProperties object.
+  arrow::Result<std::shared_ptr<::parquet::FileDecryptionProperties>>
+  SafeGetFileDecryptionProperties(
+      const ::parquet::encryption::KmsConnectionConfig& kms_connection_config,
+      const ::parquet::encryption::DecryptionConfiguration& decryption_config);
+};
+
+}  // namespace encryption
+}  // namespace parquet
+}  // namespace py
+}  // namespace arrow