Fix bulkWriter example (milvus-io#834)

Signed-off-by: lentitude2tk <xushuang.hu@zilliz.com>
yongpengli-z · Mar 29, 2024 · b9abebd · b9abebd
1 parent 5ab3c8b
commit b9abebd
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 7 deletions.
diff --git a/examples/main/java/io/milvus/BulkWriterExample.java b/examples/main/java/io/milvus/BulkWriterExample.java
@@ -89,12 +89,12 @@ public static class MilvusConsts {
      */
     public static class StorageConsts {
         public static final CloudStorage cloudStorage = CloudStorage.AWS;
-        public static final String STORAGE_ENDPOINT = cloudStorage.getEndpoint();
 
         /**
          * If using remote storage such as AWS S3, GCP GCS, Aliyun OSS, Tencent Cloud TOS,
          * please configure the following parameters.
          */
+        public static final String STORAGE_ENDPOINT = cloudStorage.getEndpoint();
         public static final String STORAGE_BUCKET = "storage.bucket";
         public static final String STORAGE_ACCESS_KEY = "storage.access.key";
         public static final String STORAGE_SECRET_KEY = "storage.secret.key";
@@ -531,7 +531,9 @@ private void callBulkInsert(CollectionSchemaParam collectionSchema, List<List<St
     private void callCloudImport(List<List<String>> batchFiles, String collectionName) throws InterruptedException, MalformedURLException {
         System.out.println("\n===================== call cloudImport ====================");
 
-        String objectUrl = StorageConsts.cloudStorage.getObjectUrl(StorageConsts.STORAGE_BUCKET, ImportUtils.getCommonPrefix(batchFiles), StorageConsts.STORAGE_REGION);
+        String objectUrl = StorageConsts.cloudStorage == CloudStorage.AZURE
+                ? StorageConsts.cloudStorage.getAzureObjectUrl(StorageConsts.AZURE_ACCOUNT_NAME, StorageConsts.AZURE_CONTAINER_NAME, ImportUtils.getCommonPrefix(batchFiles))
+                : StorageConsts.cloudStorage.getS3ObjectUrl(StorageConsts.STORAGE_BUCKET, ImportUtils.getCommonPrefix(batchFiles), StorageConsts.STORAGE_REGION);
         String accessKey = StorageConsts.cloudStorage == CloudStorage.AZURE ? StorageConsts.AZURE_ACCOUNT_NAME : StorageConsts.STORAGE_ACCESS_KEY;
         String secretKey = StorageConsts.cloudStorage == CloudStorage.AZURE ? StorageConsts.AZURE_ACCOUNT_KEY : StorageConsts.STORAGE_SECRET_KEY;
 

diff --git a/pom.xml b/pom.xml
@@ -322,6 +322,10 @@
                     <artifactId>netty</artifactId>
                     <groupId>io.netty</groupId>
                 </exclusion>
+                <exclusion>
+                    <artifactId>netty-all</artifactId>
+                    <groupId>io.netty</groupId>
+                </exclusion>
             </exclusions>
         </dependency>
         <dependency>

diff --git a/src/main/java/io/milvus/bulkwriter/common/clientenum/CloudStorage.java b/src/main/java/io/milvus/bulkwriter/common/clientenum/CloudStorage.java
@@ -30,14 +30,12 @@ public String getEndpoint(String... replaceParams) {
         return String.format(endpoint, replaceParams[0]);
     }
 
-    public String getObjectUrl(String bucketName, String commonPrefix, String region) {
+    public String getS3ObjectUrl(String bucketName, String commonPrefix, String region) {
         switch (this) {
             case AWS:
                 return String.format("https://s3.%s.amazonaws.com/%s/%s", region, bucketName, commonPrefix);
             case GCP:
                 return String.format("https://storage.cloud.google.com/%s/%s", bucketName, commonPrefix);
-            case AZURE:
-                return String.format("https://zillizvdctest.blob.core.windows.net/%s/%s", bucketName, commonPrefix);
             case TC:
                 return String.format("https://%s.cos.%s.myqcloud.com/%s", bucketName, region, commonPrefix);
             case ALI:
@@ -46,4 +44,11 @@ public String getObjectUrl(String bucketName, String commonPrefix, String region
                 throw new ParamException("no support others storage address");
         }
     }
+
+    public String getAzureObjectUrl(String accountName, String containerName, String commonPrefix) {
+        if (this == CloudStorage.AZURE) {
+            return String.format("https://%s.blob.core.windows.net/%s/%s", accountName, containerName, commonPrefix);
+        }
+        throw new ParamException("no support others storage address");
+    }
 }
diff --git a/tests/milvustest/src/test/java/com/zilliz/milvustest/bulkimport/BulkImportTest.java b/tests/milvustest/src/test/java/com/zilliz/milvustest/bulkimport/BulkImportTest.java
@@ -309,8 +309,9 @@ public void bulkRemoteImportTest() throws Exception {
         if (Objects.equals(PropertyFilesUtil.getRunValue("storageType"), "aws")) {
             cloudStorage = CloudStorage.AWS;
         }
-        String objectUrl = cloudStorage.getObjectUrl(PropertyFilesUtil.getRunValue("storageBucket"),
-                ImportUtils.getCommonPrefix(batchFiles), PropertyFilesUtil.getRunValue("storageRegion"));
+        String objectUrl = cloudStorage == CloudStorage.AZURE
+                ? cloudStorage.getAzureObjectUrl(PropertyFilesUtil.getRunValue("azureAccountName"), PropertyFilesUtil.getRunValue("azureContainerName"), ImportUtils.getCommonPrefix(batchFiles))
+                : cloudStorage.getS3ObjectUrl(PropertyFilesUtil.getRunValue("storageBucket"), ImportUtils.getCommonPrefix(batchFiles), PropertyFilesUtil.getRunValue("storageRegion"));
         String accessKey = cloudStorage == CloudStorage.AZURE ? PropertyFilesUtil.getRunValue("azureAccountName") : PropertyFilesUtil.getRunValue("storageAccessKey");
         String secretKey = cloudStorage == CloudStorage.AZURE ? PropertyFilesUtil.getRunValue("azureAccountKey") : PropertyFilesUtil.getRunValue("storageSecretKey");