From 9a79802046b6e3acf2afaceb4870cc13ef25ee64 Mon Sep 17 00:00:00 2001
From: Gaius <gaius.qi@gmail.com>
Date: Mon, 20 Nov 2023 19:51:25 +0800
Subject: [PATCH] feat: add endpoint for s3

Signed-off-by: Gaius <gaius.qi@gmail.com>
---
 .../dragonfly/config/ObjectStorageConfig.java  |  6 +++---
 .../plugins/dragonfly/objectstorage/OSS.java   |  4 ++--
 .../plugins/dragonfly/objectstorage/S3.java    | 18 ++++++++++++------
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java
index 4c021c0..8e2bf74 100644
--- a/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java
+++ b/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java
@@ -49,7 +49,7 @@ public class ObjectStorageConfig {
 
   // Endpoint of OSS.
   @SerializedName("endpoint")
-  private String endPoint;
+  private String endpoint;
 
   // Access key ID of OSS.
   @SerializedName("access_key_id")
@@ -123,8 +123,8 @@ public String getContainerName() {
     return containerName;
   }
 
-  public String getEendPoint() {
-    return endPoint;
+  public String getEndpoint() {
+    return endpoint;
   }
 
   public String getAccessKeyId() {
diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java
index 939bde3..e4ee8f4 100644
--- a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java
+++ b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java
@@ -27,11 +27,11 @@ public class OSS implements ObjectStorage {
 
   public OSS(ObjectStorageConfig objectStorageConfig) {
 
-    String endPoint = objectStorageConfig.getEendPoint();
+    String endpoint = objectStorageConfig.getEndpoint();
     String accessKeyId = objectStorageConfig.getAccessKeyId();
     String accessKeySecret = objectStorageConfig.getAccessKeySecret();
 
-    ossClient = new OSSClientBuilder().build(endPoint, accessKeyId, accessKeySecret);
+    ossClient = new OSSClientBuilder().build(endpoint, accessKeyId, accessKeySecret);
   }
 
   @Override
diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java
index 0c26634..0ade442 100644
--- a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java
+++ b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java
@@ -18,6 +18,7 @@
 
 import java.net.URL;
 import java.time.Duration;
+import org.apache.commons.lang3.StringUtils;
 import org.pytorch.serve.plugins.dragonfly.config.ObjectStorageConfig;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
@@ -31,16 +32,21 @@ public class S3 implements ObjectStorage {
   private S3Presigner presigner;
 
   public S3(ObjectStorageConfig objectStorageConfig) {
+    String endpoint = objectStorageConfig.getEndpoint();
     String accessKey = objectStorageConfig.getAccessKey();
     String secretKey = objectStorageConfig.getSecretKey();
     Region region = Region.of(objectStorageConfig.getRegion());
-
     AwsBasicCredentials awsCreds = AwsBasicCredentials.create(accessKey, secretKey);
-    presigner =
-        S3Presigner.builder()
-            .region(region)
-            .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
-            .build();
+
+    val builder = S3Presigner.builder();
+    builder.region(region);
+    builder.credentialsProvider(StaticCredentialsProvider.create(awsCreds));
+
+    if (StringUtils.isNotBlank(endpoint)) {
+      builder.endpointOverride(new URI(endpoint));
+    }
+
+    presigner = builder.build();
   }
 
   @Override