support stageManager & stageFileManager

examples/src/main/java/io/milvus/v2/StageFileManagerExample.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package io.milvus.v2;
+
+import com.google.gson.Gson;
+import io.milvus.bulkwriter.StageFileManager;
+import io.milvus.bulkwriter.StageFileManagerParam;
+import io.milvus.bulkwriter.model.UploadFilesResult;
+import io.milvus.bulkwriter.request.stage.UploadFilesRequest;
+
+
+/**
+ * This is currently a private preview feature. If you need to use it, please submit a request and contact us.
+ */
+public class StageFileManagerExample {
+    private static final StageFileManager stageFileManager;
+    static {
+        StageFileManagerParam stageFileManagerParam = StageFileManagerParam.newBuilder()
+                .withCloudEndpoint("https://api.cloud.zilliz.com")
+                .withApiKey("_api_key_for_cluster_org_")
+                .withStageName("_stage_name_for_project_")
+                .build();
+        stageFileManager = new StageFileManager(stageFileManagerParam);
+    }
+
+    public static void main(String[] args) throws Exception {
+        uploadFiles();
+        shutdown();
+    }
+
+    private static void uploadFiles() throws Exception {
+        UploadFilesRequest request = UploadFilesRequest.builder()
+                .sourceFilePath("/Users/zilliz/data/")
+                .targetStagePath("data/")
+                .build();
+        UploadFilesResult result = stageFileManager.uploadFilesAsync(request).get();
+        System.out.println("\nuploadFiles results: " + new Gson().toJson(result));
+    }
+
+    private static void shutdown() {
+        stageFileManager.shutdownGracefully();
+    }
+}

examples/src/main/java/io/milvus/v2/StageManagerExample.java

@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package io.milvus.v2;
+
+import com.google.gson.Gson;
+import io.milvus.bulkwriter.StageManager;
+import io.milvus.bulkwriter.StageManagerParam;
+import io.milvus.bulkwriter.request.stage.CreateStageRequest;
+import io.milvus.bulkwriter.request.stage.DeleteStageRequest;
+import io.milvus.bulkwriter.request.stage.ListStagesRequest;
+import io.milvus.bulkwriter.response.stage.ListStagesResponse;
+
+
+/**
+ * This is currently a private preview feature. If you need to use it, please submit a request and contact us.
+ */
+public class StageManagerExample {
+    private static final StageManager stageManager;
+    static {
+        StageManagerParam stageManagerParam = StageManagerParam.newBuilder()
+                .withCloudEndpoint("https://api.cloud.zilliz.com")
+                .withApiKey("_api_key_for_cluster_org_")
+                .build();
+        stageManager = new StageManager(stageManagerParam);
+    }
+
+    private static final String PROJECT_ID = "_id_for_project_";
+    private static final String REGION_ID = "_id_for_region_";
+    private static final String STAGE_NAME = "_stage_name_for_project_";
+
+    public static void main(String[] args) throws Exception {
+        createStage();
+        listStages();
+        deleteStage();
+    }
+
+    private static void createStage() {
+        CreateStageRequest request = CreateStageRequest.builder()
+                .projectId(PROJECT_ID).regionId(REGION_ID).stageName(STAGE_NAME)
+                .build();
+        stageManager.createStage(request);
+        System.out.printf("\nStage %s created%n", STAGE_NAME);
+    }
+
+    private static void listStages() {
+        ListStagesRequest request = ListStagesRequest.builder()
+                .projectId(PROJECT_ID).currentPage(1).pageSize(10)
+                .build();
+        ListStagesResponse listStagesResponse = stageManager.listStages(request);
+        System.out.println("\nlistStages results: " + new Gson().toJson(listStagesResponse));
+    }
+
+    private static void deleteStage() {
+        DeleteStageRequest request = DeleteStageRequest.builder()
+                .stageName(STAGE_NAME)
+                .build();
+        stageManager.deleteStage(request);
+        System.out.printf("\nStage %s deleted%n", STAGE_NAME);
+    }
+}

examples/src/main/java/io/milvus/v2/bulkwriter/BulkWriterStageExample.java

@@ -27,7 +27,7 @@
 import io.milvus.bulkwriter.StageBulkWriterParam;
 import io.milvus.bulkwriter.common.clientenum.BulkFileType;
 import io.milvus.bulkwriter.common.utils.GeneratorUtils;
-import io.milvus.bulkwriter.model.StageUploadResult;
+import io.milvus.bulkwriter.model.UploadFilesResult;
 import io.milvus.bulkwriter.request.describe.CloudDescribeImportRequest;
 import io.milvus.bulkwriter.request.import_.StageImportRequest;
 import io.milvus.bulkwriter.request.list.CloudListImportJobsRequest;
@@ -118,7 +118,7 @@ private static void exampleCollectionRemoteStage(BulkFileType fileType) throws E
         CreateCollectionReq.CollectionSchema collectionSchema = buildAllTypesSchema();
         createCollection(COLLECTION_NAME, collectionSchema, false);
 
-        StageUploadResult stageUploadResult = stageRemoteWriter(collectionSchema, fileType, rows);
+        UploadFilesResult stageUploadResult = stageRemoteWriter(collectionSchema, fileType, rows);
         callStageImport(stageUploadResult.getStageName(), stageUploadResult.getPath());
         verifyImportData(collectionSchema, originalData);
     }
@@ -284,7 +284,7 @@ private static List<JsonObject> genImportData(List<Map<String, Object>> original
         return data;
     }
 
-    private static StageUploadResult stageRemoteWriter(CreateCollectionReq.CollectionSchema collectionSchema,
+    private static UploadFilesResult stageRemoteWriter(CreateCollectionReq.CollectionSchema collectionSchema,
                                                        BulkFileType fileType,
                                                        List<JsonObject> data) throws Exception {
         System.out.printf("\n===================== all field types (%s) ====================%n", fileType.name());
@@ -297,7 +297,7 @@ private static StageUploadResult stageRemoteWriter(CreateCollectionReq.Collectio
             System.out.println("Generate data files...");
             stageBulkWriter.commit(false);
 
-            StageUploadResult stageUploadResult = stageBulkWriter.getStageUploadResult();
+            UploadFilesResult stageUploadResult = stageBulkWriter.getStageUploadResult();
             System.out.printf("Data files have been uploaded: %s%n", stageUploadResult);
             return stageUploadResult;
         } catch (Exception e) {

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/StageBulkWriter.java

@@ -21,7 +21,8 @@
 
 import com.google.common.collect.Lists;
 import com.google.gson.JsonObject;
-import io.milvus.bulkwriter.model.StageUploadResult;
+import io.milvus.bulkwriter.model.UploadFilesResult;
+import io.milvus.bulkwriter.request.stage.UploadFilesRequest;
 import io.milvus.common.utils.ExceptionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -39,7 +40,7 @@ public class StageBulkWriter extends LocalBulkWriter {
 
     private String remotePath;
     private List<List<String>> remoteFiles;
-    private StageOperation stageWriter;
+    private StageFileManager stageFileManager;
     private StageBulkWriterParam stageBulkWriterParam;
 
     public StageBulkWriter(StageBulkWriterParam bulkWriterParam) throws IOException {
@@ -51,20 +52,20 @@ public StageBulkWriter(StageBulkWriterParam bulkWriterParam) throws IOException
         Path path = Paths.get(bulkWriterParam.getRemotePath());
         Path remoteDirPath = path.resolve(getUUID());
         this.remotePath = remoteDirPath + "/";
-        this.stageWriter = initStageWriterParams(bulkWriterParam);
+        this.stageFileManager = initStageFileManagerParams(bulkWriterParam);
         this.stageBulkWriterParam = bulkWriterParam;
 
         this.remoteFiles = Lists.newArrayList();
         logger.info("Remote buffer writer initialized, target path: {}", remotePath);
 
     }
 
-    private StageOperation initStageWriterParams(StageBulkWriterParam bulkWriterParam) throws IOException {
-        StageOperationParam stageWriterParam = StageOperationParam.newBuilder()
+    private StageFileManager initStageFileManagerParams(StageBulkWriterParam bulkWriterParam) throws IOException {
+        StageFileManagerParam stageFileManagerParam = StageFileManagerParam.newBuilder()
                 .withCloudEndpoint(bulkWriterParam.getCloudEndpoint()).withApiKey(bulkWriterParam.getApiKey())
-                .withStageName(bulkWriterParam.getStageName()).withPath(remotePath)
+                .withStageName(bulkWriterParam.getStageName())
                 .build();
-        return new StageOperation(stageWriterParam);
+        return new StageFileManager(stageFileManagerParam);
     }
 
     @Override
@@ -87,8 +88,8 @@ public List<List<String>> getBatchFiles() {
         return remoteFiles;
     }
 
-    public StageUploadResult getStageUploadResult() {
-        return StageUploadResult.builder()
+    public UploadFilesResult getStageUploadResult() {
+        return UploadFilesResult.builder()
                 .stageName(stageBulkWriterParam.getStageName())
                 .path(remotePath)
                 .build();
@@ -175,8 +176,13 @@ private void serialImportData(List<String> fileList) {
     private void uploadObject(String filePath, String objectName) throws Exception {
         logger.info(String.format("Prepare to upload %s to %s", filePath, objectName));
 
-        stageWriter.uploadFileToStage(filePath);
+        UploadFilesRequest uploadFilesRequest = UploadFilesRequest.builder()
+                .sourceFilePath(filePath).targetStagePath(remotePath)
+                .build();
+
+        stageFileManager.uploadFilesAsync(uploadFilesRequest).get();
         logger.info(String.format("Upload file %s to %s", filePath, objectName));
+
     }
 
     private static String generatorLocalPath() {