add notNull annotation for request (#773)

Signed-off-by: Nian Liu <nian.liu@zilliz.com>

src/main/java/io/milvus/v2/client/ConnectConfig.java

@@ -2,6 +2,7 @@ package io.milvus.v2.client;
 
 import lombok.Builder;
 import lombok.Data;
+import lombok.NonNull;
 import lombok.experimental.SuperBuilder;
 
 import java.net.URI;
@@ -10,6 +11,7 @@ import java.util.concurrent.TimeUnit;
 @Data
 @SuperBuilder
 public class ConnectConfig {
+    @NonNull
     private String uri;
     private String token;
     private String username;

src/main/java/io/milvus/v2/service/collection/CollectionService.java

@@ -2,6 +2,8 @@ package io.milvus.v2.service.collection;
 
 import io.milvus.grpc.*;
 import io.milvus.v2.common.IndexParam;
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
 import io.milvus.v2.service.BaseService;
 import io.milvus.v2.service.collection.request.*;
 import io.milvus.v2.service.collection.response.DescribeCollectionResp;
@@ -125,6 +127,10 @@ public class CollectionService extends BaseService {
                 .build();
         Status status = milvusServiceBlockingStub.dropCollection(dropCollectionRequest);
         rpcUtils.handleResponse(title, status);
+
+        if (request.getAsync()) {
+            WaitForDropCollection(milvusServiceBlockingStub, request);
+        }
     }
 
     public Boolean hasCollection(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, HasCollectionReq request) {
@@ -175,9 +181,13 @@ public class CollectionService extends BaseService {
         String title = String.format("LoadCollectionRequest collectionName:%s", request.getCollectionName());
         LoadCollectionRequest loadCollectionRequest = LoadCollectionRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
+                .setReplicaNumber(request.getReplicaNum())
                 .build();
         Status status = milvusServiceBlockingStub.loadCollection(loadCollectionRequest);
         rpcUtils.handleResponse(title, status);
+        if (request.getAsync()) {
+            WaitForLoadCollection(milvusServiceBlockingStub, request);
+        }
     }
 
     public void releaseCollection(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, ReleaseCollectionReq request) {
@@ -187,6 +197,9 @@ public class CollectionService extends BaseService {
                 .build();
         Status status = milvusServiceBlockingStub.releaseCollection(releaseCollectionRequest);
         rpcUtils.handleResponse(title, status);
+        if (request.getAsync()) {
+            waitForCollectionRelease(milvusServiceBlockingStub, request);
+        }
     }
 
     public Boolean getLoadState(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, GetLoadStateReq request) {
@@ -223,4 +236,76 @@ public class CollectionService extends BaseService {
                 .description(description)
                 .build();
     }
+
+    public void waitForCollectionRelease(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, ReleaseCollectionReq request) {
+        boolean isLoaded = true;
+        long startTime = System.currentTimeMillis(); // Capture start time/ Timeout in milliseconds (60 seconds)
+
+        while (isLoaded) {
+            // Call the getLoadState method
+            isLoaded = getLoadState(milvusServiceBlockingStub, GetLoadStateReq.builder().collectionName(request.getCollectionName()).build());
+            if (isLoaded) {
+                // Check if timeout is exceeded
+                if (System.currentTimeMillis() - startTime > request.getTimeout()) {
+                    throw new MilvusClientException(ErrorCode.SERVER_ERROR, "Load collection timeout");
+                }
+                // Wait for a certain period before checking again
+                try {
+                    Thread.sleep(500); // Sleep for 0.5 second. Adjust this value as needed.
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    System.out.println("Thread was interrupted, Failed to complete operation");
+                    return; // or handle interruption appropriately
+                }
+            }
+        }
+    }
+
+    private void WaitForLoadCollection(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, LoadCollectionReq request) {
+        boolean isLoaded = false;
+        long startTime = System.currentTimeMillis(); // Capture start time/ Timeout in milliseconds (60 seconds)
+
+        while (!isLoaded) {
+            // Call the getLoadState method
+            isLoaded = getLoadState(milvusServiceBlockingStub, GetLoadStateReq.builder().collectionName(request.getCollectionName()).build());
+            if (!isLoaded) {
+                // Check if timeout is exceeded
+                if (System.currentTimeMillis() - startTime > request.getTimeout()) {
+                    throw new MilvusClientException(ErrorCode.SERVER_ERROR, "Load collection timeout");
+                }
+                // Wait for a certain period before checking again
+                try {
+                    Thread.sleep(500); // Sleep for 0.5 second. Adjust this value as needed.
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    System.out.println("Thread was interrupted, Failed to complete operation");
+                    return; // or handle interruption appropriately
+                }
+            }
+        }
+    }
+
+    private void WaitForDropCollection(MilvusServiceGrpc.MilvusServiceBlockingStub milvusServiceBlockingStub, DropCollectionReq request) {
+        boolean hasCollection = true;
+        long startTime = System.currentTimeMillis(); // Capture start time/ Timeout in milliseconds (60 seconds)
+
+        while (hasCollection) {
+            // Call the getLoadState method
+            hasCollection = hasCollection(milvusServiceBlockingStub, HasCollectionReq.builder().collectionName(request.getCollectionName()).build());
+            if (hasCollection) {
+                // Check if timeout is exceeded
+                if (System.currentTimeMillis() - startTime > request.getTimeout()) {
+                    throw new MilvusClientException(ErrorCode.SERVER_ERROR, "drop collection timeout");
+                }
+                // Wait for a certain period before checking again
+                try {
+                    Thread.sleep(500); // Sleep for 0.5 second. Adjust this value as needed.
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    System.out.println("Thread was interrupted, Failed to complete operation");
+                    return; // or handle interruption appropriately
+                }
+            }
+        }
+    }
 }

src/main/java/io/milvus/v2/service/collection/request/CreateCollectionReq.java

@@ -4,6 +4,7 @@ import io.milvus.v2.common.DataType;
 import io.milvus.v2.common.IndexParam;
 import lombok.Builder;
 import lombok.Data;
+import lombok.NonNull;
 import lombok.experimental.SuperBuilder;
 
 import java.util.ArrayList;
@@ -12,6 +13,7 @@ import java.util.List;
 @Data
 @SuperBuilder
 public class CreateCollectionReq {
+    @NonNull
     private String collectionName;
     private Integer dimension;
 
@@ -42,6 +44,7 @@ public class CreateCollectionReq {
         private List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<>();
         @Builder.Default
         private String description = "";
+        @NonNull
         private Boolean enableDynamicField;
 
         public CreateCollectionReq.FieldSchema getField(String fieldName) {

src/main/java/io/milvus/v2/service/collection/request/DropCollectionReq.java

@@ -1,5 +1,6 @@
 package io.milvus.v2.service.collection.request;
 
+import lombok.Builder;
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
@@ -7,4 +8,8 @@ import lombok.experimental.SuperBuilder;
 @SuperBuilder
 public class DropCollectionReq {
     private String collectionName;
+    @Builder.Default
+    private Boolean async = Boolean.TRUE;
+    @Builder.Default
+    private Long timeout = 60000L;
 }

src/main/java/io/milvus/v2/service/collection/request/LoadCollectionReq.java

@@ -1,5 +1,6 @@
 package io.milvus.v2.service.collection.request;
 
+import lombok.Builder;
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
@@ -7,4 +8,10 @@ import lombok.experimental.SuperBuilder;
 @SuperBuilder
 public class LoadCollectionReq {
     private String collectionName;
+    @Builder.Default
+    private Integer replicaNum = 1;
+    @Builder.Default
+    private Boolean async = Boolean.TRUE;
+    @Builder.Default
+    private Long timeout = 60000L;
 }

src/main/java/io/milvus/v2/service/collection/request/ReleaseCollectionReq.java

@@ -1,5 +1,6 @@
 package io.milvus.v2.service.collection.request;
 
+import lombok.Builder;
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
@@ -7,4 +8,8 @@ import lombok.experimental.SuperBuilder;
 @SuperBuilder
 public class ReleaseCollectionReq {
     private String collectionName;
+    @Builder.Default
+    private Boolean async = Boolean.TRUE;
+    @Builder.Default
+    private Long timeout = 60000L;
 }

src/main/java/io/milvus/v2/service/index/IndexService.java

@@ -57,8 +57,8 @@ public class IndexService extends BaseService {
                 request.getCollectionName(), request.getFieldName(), request.getIndexName());
         DropIndexRequest dropIndexRequest = DropIndexRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
-                .setFieldName(request.getFieldName())
-                .setIndexName(request.getIndexName())
+                .setFieldName(request.getFieldName() == null ? "" : request.getFieldName())
+                .setIndexName(request.getIndexName() == null ? "" : request.getIndexName())
                 .build();
 
         Status status = milvusServiceBlockingStub.dropIndex(dropIndexRequest);
@@ -70,18 +70,9 @@ public class IndexService extends BaseService {
                 request.getCollectionName(), request.getFieldName(), request.getIndexName());
         DescribeIndexRequest describeIndexRequest = DescribeIndexRequest.newBuilder()
                 .setCollectionName(request.getCollectionName())
-//                .setFieldName(request.getFieldName())
-//                .setIndexName(request.getIndexName())
+                .setFieldName(request.getFieldName() == null ? "" : request.getFieldName())
+                .setIndexName(request.getIndexName() == null ? "" : request.getIndexName())
                 .build();
-        if (request.getFieldName() != null) {
-            describeIndexRequest = describeIndexRequest.toBuilder()
-                    .setFieldName(request.getFieldName())
-                    .build();
-        } else if (request.getIndexName() != null) {
-            describeIndexRequest = describeIndexRequest.toBuilder()
-                    .setIndexName(request.getIndexName())
-                    .build();
-        }
 
         DescribeIndexResponse response = milvusServiceBlockingStub.describeIndex(describeIndexRequest);
         rpcUtils.handleResponse(title, response.getStatus());

src/main/java/io/milvus/v2/service/index/request/CreateIndexReq.java

@@ -2,6 +2,7 @@ package io.milvus.v2.service.index.request;
 
 import io.milvus.v2.common.IndexParam;
 import lombok.Data;
+import lombok.NonNull;
 import lombok.experimental.SuperBuilder;
 
 import java.util.List;
@@ -9,6 +10,7 @@ import java.util.List;
 @Data
 @SuperBuilder
 public class CreateIndexReq {
+    @NonNull
     private String collectionName;
     private List<IndexParam> indexParams;
 }

src/main/java/io/milvus/v2/service/index/request/DescribeIndexReq.java

@@ -1,11 +1,13 @@
 package io.milvus.v2.service.index.request;
 
 import lombok.Data;
+import lombok.NonNull;
 import lombok.experimental.SuperBuilder;
 
 @Data
 @SuperBuilder
 public class DescribeIndexReq {
+    @NonNull
     private String collectionName;
     private String fieldName;
     private String indexName;

src/main/java/io/milvus/v2/service/index/request/DropIndexReq.java

@@ -1,14 +1,14 @@
 package io.milvus.v2.service.index.request;
 
-import lombok.Builder;
 import lombok.Data;
+import lombok.NonNull;
 import lombok.experimental.SuperBuilder;
 
 @Data
 @SuperBuilder
 public class DropIndexReq {
+    @NonNull
     private String collectionName;
     private String fieldName;
-    @Builder.Default
-    private String indexName = "";
+    private String indexName;
 }