refine the code

apache · Aug 8, 2024 · 7e1f200 · 7e1f200
1 parent c16ec3b
commit 7e1f200
Showing 1 changed file with 23 additions and 25 deletions.
diff --git 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 <typename S3RequestType>
+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<S3Model::GetObjectResult> 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<StringViewStream>(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<StringViewStream>(data, nbytes));