diff --git a/cpp/src/arrow/filesystem/s3fs.cc b/cpp/src/arrow/filesystem/s3fs.cc index e441607a138ad..e56955b6e2f45 100644 --- a/cpp/src/arrow/filesystem/s3fs.cc +++ b/cpp/src/arrow/filesystem/s3fs.cc @@ -170,6 +170,21 @@ static constexpr const char kAwsEndpointUrlEnvVar[] = "AWS_ENDPOINT_URL"; static constexpr const char kAwsEndpointUrlS3EnvVar[] = "AWS_ENDPOINT_URL_S3"; static constexpr const char kAwsDirectoryContentType[] = "application/x-directory"; +template +void SetSSECustomerKey(S3RequestType& request, const std::string& sse_customer_algorithm, + const std::string& sse_customer_key, + const std::string& sse_customer_key_md5) { + if (!sse_customer_algorithm.empty()) { + request.SetSSECustomerAlgorithm(sse_customer_algorithm); + } + if (!sse_customer_key.empty()) { + request.SetSSECustomerKey(sse_customer_key); + } + if (!sse_customer_key_md5.empty()) { + request.SetSSECustomerKeyMD5(sse_customer_key_md5); + } +} + // ----------------------------------------------------------------------- // S3ProxyOptions implementation @@ -1330,11 +1345,7 @@ Result GetObjectRange(Aws::S3::S3Client* client, S3Model::GetObjectRequest req; req.SetBucket(ToAwsString(path.bucket)); req.SetKey(ToAwsString(path.key)); - if (!sse_customer_key.empty()){ - req.SetSSECustomerAlgorithm(sse_customer_algorithm); - req.SetSSECustomerKey(sse_customer_key); - req.SetSSECustomerKeyMD5(sse_customer_key_md5); - } + SetSSECustomerKey(req, sse_customer_algorithm, sse_customer_key, sse_customer_key_md5); req.SetRange(ToAwsString(FormatRange(start, length))); req.SetResponseStreamFactory(AwsWriteableStreamFactory(out, length)); return OutcomeToResult("GetObject", client->GetObject(req)); @@ -1491,11 +1502,7 @@ class ObjectInputFile final : public io::RandomAccessFile { S3Model::HeadObjectRequest req; req.SetBucket(ToAwsString(path_.bucket)); req.SetKey(ToAwsString(path_.key)); - if (!sse_customer_key.empty()){ - req.SetSSECustomerAlgorithm(sse_customer_algorithm); - req.SetSSECustomerKey(sse_customer_key); - req.SetSSECustomerKeyMD5(sse_customer_key_md5); - } + SetSSECustomerKey(req, sse_customer_algorithm, sse_customer_key, sse_customer_key_md5); ARROW_ASSIGN_OR_RAISE(auto client_lock, holder_->Lock()); auto outcome = client_lock.Move()->HeadObject(req); @@ -1720,11 +1727,8 @@ class ObjectOutputStream final : public io::OutputStream { S3Model::CreateMultipartUploadRequest req; req.SetBucket(ToAwsString(path_.bucket)); req.SetKey(ToAwsString(path_.key)); - if (!sse_customer_key.empty()){ - req.SetSSECustomerAlgorithm(sse_customer_algorithm); - req.SetSSECustomerKey(sse_customer_key); - req.SetSSECustomerKeyMD5(sse_customer_key_md5); - } + SetSSECustomerKey(req, sse_customer_algorithm, sse_customer_key, sse_customer_key_md5); + RETURN_NOT_OK(SetMetadataInRequest(&req)); auto outcome = client_lock.Move()->CreateMultipartUpload(req); @@ -1826,11 +1830,8 @@ class ObjectOutputStream final : public io::OutputStream { S3Model::CompleteMultipartUploadRequest req; req.SetBucket(ToAwsString(path_.bucket)); req.SetKey(ToAwsString(path_.key)); - if (!sse_customer_key.empty()){ - req.SetSSECustomerAlgorithm(sse_customer_algorithm); - req.SetSSECustomerKey(sse_customer_key); - req.SetSSECustomerKeyMD5(sse_customer_key_md5); - } + SetSSECustomerKey(req, sse_customer_algorithm, sse_customer_key, sse_customer_key_md5); + req.SetUploadId(multipart_upload_id_); req.SetMultipartUpload(std::move(completed_upload)); @@ -2013,11 +2014,8 @@ class ObjectOutputStream final : public io::OutputStream { req.SetKey(ToAwsString(path_.key)); req.SetBody(std::make_shared(data, nbytes)); req.SetContentLength(nbytes); - if (!sse_customer_key.empty()){ - req.SetSSECustomerAlgorithm(sse_customer_algorithm); - req.SetSSECustomerKey(sse_customer_key); - req.SetSSECustomerKeyMD5(sse_customer_key_md5); - } + SetSSECustomerKey(req, sse_customer_algorithm, sse_customer_key, sse_customer_key_md5); + if (!background_writes_) { req.SetBody(std::make_shared(data, nbytes));