feat(storage): support presigned url with put method (#2197)

star-whale · May 10, 2023 · 642df1d · 642df1d
1 parent 0943fb7
commit 642df1d
Show file tree

Hide file tree

Showing 13 changed files with 119 additions and 16 deletions.
diff --git a/...r/storage-access-layer/src/main/java/ai/starwhale/mlops/storage/StorageAccessService.java b/...r/storage-access-layer/src/main/java/ai/starwhale/mlops/storage/StorageAccessService.java
@@ -50,4 +50,6 @@ public interface StorageAccessService {
      * @throws IOException any possible IO exception
      */
     String signedUrl(String path, Long expTimeMillis) throws IOException;
+
+    String signedPutUrl(String path, Long expTimeMillis) throws IOException;
 }
diff --git a/...ess-layer/src/main/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyun.java b/...ess-layer/src/main/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyun.java
@@ -23,6 +23,7 @@
 import ai.starwhale.mlops.storage.s3.S3Config;
 import ai.starwhale.mlops.storage.util.MetaHelper;
 import com.aliyun.oss.ClientBuilderConfiguration;
+import com.aliyun.oss.HttpMethod;
 import com.aliyun.oss.OSS;
 import com.aliyun.oss.OSSClientBuilder;
 import com.aliyun.oss.OSSErrorCode;
@@ -166,4 +167,13 @@ public String signedUrl(String path, Long expTimeMillis) throws IOException {
         return ossClient.generatePresignedUrl(this.bucket, path, new Date(System.currentTimeMillis() + expTimeMillis))
                 .toString();
     }
+
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return ossClient.generatePresignedUrl(this.bucket,
+                        path,
+                        new Date(System.currentTimeMillis() + expTimeMillis),
+                        HttpMethod.PUT)
+                .toString();
+    }
 }
diff --git a/...ayer/src/main/java/ai/starwhale/mlops/storage/autofit/CompatibleStorageAccessService.java b/...ayer/src/main/java/ai/starwhale/mlops/storage/autofit/CompatibleStorageAccessService.java
@@ -88,5 +88,8 @@ public String signedUrl(String path, Long expTimeMillis) throws IOException {
         return storageAccessService.signedUrl(path, expTimeMillis);
     }
 
-
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return storageAccessService.signedPutUrl(path, expTimeMillis);
+    }
 }
diff --git a/...c/main/java/ai/starwhale/mlops/storage/autofit/fs/CompatibleStorageAccessServiceFile.java b/...c/main/java/ai/starwhale/mlops/storage/autofit/fs/CompatibleStorageAccessServiceFile.java
@@ -92,5 +92,8 @@ public String signedUrl(String path, Long expTimeMillis) throws IOException {
         return storageAccessService.signedUrl(path, expTimeMillis);
     }
 
-
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return storageAccessService.signedPutUrl(path, expTimeMillis);
+    }
 }
diff --git a/...main/java/ai/starwhale/mlops/storage/autofit/s3/CompatibleStorageAccessServiceS3Like.java b/...main/java/ai/starwhale/mlops/storage/autofit/s3/CompatibleStorageAccessServiceS3Like.java
@@ -127,5 +127,8 @@ public String signedUrl(String path, Long expTimeMillis) throws IOException {
         return storageAccessService.signedUrl(path, expTimeMillis);
     }
 
-
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return storageAccessService.signedPutUrl(path, expTimeMillis);
+    }
 }
diff --git a/...ge-access-layer/src/main/java/ai/starwhale/mlops/storage/fs/StorageAccessServiceFile.java b/...ge-access-layer/src/main/java/ai/starwhale/mlops/storage/fs/StorageAccessServiceFile.java
@@ -146,4 +146,9 @@ public void delete(String path) throws IOException {
     public String signedUrl(String path, Long expTimeMillis) throws IOException {
         return serviceProvider + "/" + path + "/" + (System.currentTimeMillis() + expTimeMillis);
     }
+
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return serviceProvider + "/" + path + "/" + (System.currentTimeMillis() + expTimeMillis);
+    }
 }
diff --git a/...ess-layer/src/main/java/ai/starwhale/mlops/storage/memory/StorageAccessServiceMemory.java b/...ess-layer/src/main/java/ai/starwhale/mlops/storage/memory/StorageAccessServiceMemory.java
@@ -87,7 +87,12 @@ public void delete(String path) throws IOException {
     }
 
     @Override
-    public String signedUrl(String path, Long expTimeMillis) throws IOException {
-        return null;
+    public String signedUrl(String path, Long expTimeMillis) {
+        return path;
+    }
+
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) {
+        return path;
     }
 }
diff --git a/...ccess-layer/src/main/java/ai/starwhale/mlops/storage/minio/StorageAccessServiceMinio.java b/...ccess-layer/src/main/java/ai/starwhale/mlops/storage/minio/StorageAccessServiceMinio.java
@@ -180,17 +180,20 @@ public void delete(String path) throws IOException {
 
     @Override
     public String signedUrl(String path, Long expTimeMillis) throws IOException {
+        return this.signUrl(path, expTimeMillis, Method.GET);
+    }
+
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) throws IOException {
+        return this.signUrl(path, expTimeMillis, Method.PUT);
+    }
+
+    private String signUrl(String path, Long expTimeMillis, Method method) throws IOException {
         GetPresignedObjectUrlArgs request = GetPresignedObjectUrlArgs.builder().expiry(expTimeMillis.intValue(),
-                TimeUnit.MILLISECONDS).bucket(this.bucket).object(path).method(Method.GET).build();
+                TimeUnit.MILLISECONDS).bucket(this.bucket).object(path).method(method).build();
         try {
             return minioClient.getPresignedObjectUrl(request);
-        } catch (MinioException e) {
-            log.error("sin url for object {} fails", path, e);
-            throw new IOException(e);
-        } catch (InvalidKeyException e) {
-            log.error("sin url for object {} fails", path, e);
-            throw new IOException(e);
-        } catch (NoSuchAlgorithmException e) {
+        } catch (MinioException | NoSuchAlgorithmException | InvalidKeyException e) {
             log.error("sin url for object {} fails", path, e);
             throw new IOException(e);
         }

diff --git a/...rage-access-layer/src/main/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3.java b/...rage-access-layer/src/main/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3.java
@@ -57,6 +57,7 @@
 import software.amazon.awssdk.services.s3.presigner.S3Presigner;
 import software.amazon.awssdk.services.s3.presigner.S3Presigner.Builder;
 import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest;
+import software.amazon.awssdk.services.s3.presigner.model.PutObjectPresignRequest;
 
 @Slf4j
 public class StorageAccessServiceS3 implements StorageAccessService {
@@ -283,10 +284,18 @@ public void delete(String path) throws IOException {
     }
 
     @Override
-    public String signedUrl(String path, Long expTimeMillis) throws IOException {
+    public String signedUrl(String path, Long expTimeMillis) {
         return s3Presigner.presignGetObject(GetObjectPresignRequest.builder()
                 .getObjectRequest(GetObjectRequest.builder().bucket(this.s3Config.getBucket()).key(path).build())
                 .signatureDuration(Duration.ofMillis(expTimeMillis))
                 .build()).url().toString();
     }
+
+    @Override
+    public String signedPutUrl(String path, Long expTimeMillis) {
+        return s3Presigner.presignPutObject(PutObjectPresignRequest.builder()
+                .putObjectRequest(PutObjectRequest.builder().bucket(this.s3Config.getBucket()).key(path).build())
+                .signatureDuration(Duration.ofMillis(expTimeMillis))
+                .build()).url().toString();
+    }
 }
diff --git a/...layer/src/test/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyunTest.java b/...layer/src/test/java/ai/starwhale/mlops/storage/aliyun/StorageAccessServiceAliyunTest.java
@@ -20,6 +20,7 @@
 import com.adobe.testing.s3mock.testcontainers.S3MockContainer;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.HttpURLConnection;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import org.junit.jupiter.api.Assertions;
@@ -30,6 +31,7 @@
 
 @Testcontainers
 public class StorageAccessServiceAliyunTest {
+
     @Container
     private static final S3MockContainer s3Mock =
             new S3MockContainer(System.getProperty("s3mock.version", "latest"))
@@ -60,5 +62,20 @@ public void testSignedUrl() throws IOException {
         }
     }
 
-
+    @Test
+    public void testSignedPutUrl() throws IOException {
+        String path = "unit_test/x";
+        String content = "testSignedPutUrl";
+        String signedUrl = aliyun.signedPutUrl(path, 1000 * 60L);
+        var conn = (HttpURLConnection) new URL(signedUrl).openConnection();
+        conn.setDoOutput(true);
+        conn.setRequestMethod("PUT");
+        try (var out = conn.getOutputStream()) {
+            out.write(content.getBytes(StandardCharsets.UTF_8));
+        }
+        try (var in = conn.getInputStream()) {
+            in.readAllBytes();
+        }
+        Assertions.assertEquals(content, new String(aliyun.get(path).readAllBytes()));
+    }
 }
diff --git a/...ccess-layer/src/test/java/ai/starwhale/mlops/storage/fs/StorageAccessServiceFileTest.java b/...ccess-layer/src/test/java/ai/starwhale/mlops/storage/fs/StorageAccessServiceFileTest.java
@@ -74,6 +74,15 @@ public void testSignedUrl() throws IOException {
         Assertions.assertTrue(signedUrl.startsWith("http://localhost:8082/unit_test/x"));
     }
 
+    @Test
+    public void testSignedPutUrl() throws IOException {
+        String path = "unit_test/x";
+        String content = "hello word";
+        service.put(path, content.getBytes(StandardCharsets.UTF_8));
+        String signedUrl = service.signedPutUrl(path, 1000 * 60L);
+        Assertions.assertTrue(signedUrl.startsWith("http://localhost:8082/unit_test/x"));
+    }
+
     @Test
     public void testRange() throws IOException {
         String path = "unit_test/x";

diff --git a/...s-layer/src/test/java/ai/starwhale/mlops/storage/minio/StorageAccessServiceMinioTest.java b/...s-layer/src/test/java/ai/starwhale/mlops/storage/minio/StorageAccessServiceMinioTest.java
@@ -22,6 +22,7 @@
 import com.adobe.testing.s3mock.testcontainers.S3MockContainer;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.HttpURLConnection;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
@@ -111,5 +112,20 @@ public void testSignedUrl() throws IOException {
         }
     }
 
-
+    @Test
+    public void testSignedPutUrl() throws IOException {
+        String path = "unit_test/x";
+        String content = "testSignedPutUrl";
+        String signedUrl = minio.signedPutUrl(path, 1000 * 60L);
+        var conn = (HttpURLConnection) new URL(signedUrl).openConnection();
+        conn.setDoOutput(true);
+        conn.setRequestMethod("PUT");
+        try (var out = conn.getOutputStream()) {
+            out.write(content.getBytes(StandardCharsets.UTF_8));
+        }
+        try (var in = conn.getInputStream()) {
+            in.readAllBytes();
+        }
+        Assertions.assertEquals(content, new String(minio.get(path).readAllBytes()));
+    }
 }
diff --git a/...-access-layer/src/test/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3Test.java b/...-access-layer/src/test/java/ai/starwhale/mlops/storage/s3/StorageAccessServiceS3Test.java
@@ -26,6 +26,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.HttpURLConnection;
 import java.net.URI;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
@@ -204,4 +205,21 @@ public void testSignedUrl() throws IOException {
             Assertions.assertEquals("a", new String(content.readAllBytes()));
         }
     }
+
+    @Test
+    public void testSignedPutUrl() throws IOException {
+        String path = "x";
+        String content = "testSignedPutUrl";
+        String signedUrl = s3.signedPutUrl(path, 1000 * 60L);
+        var conn = (HttpURLConnection) new URL(signedUrl).openConnection();
+        conn.setDoOutput(true);
+        conn.setRequestMethod("PUT");
+        try (var out = conn.getOutputStream()) {
+            out.write(content.getBytes(StandardCharsets.UTF_8));
+        }
+        try (var in = conn.getInputStream()) {
+            in.readAllBytes();
+        }
+        Assertions.assertEquals(content, new String(s3.get(path).readAllBytes()));
+    }
 }