Add withDatabaseBame for releaseCollection

Signed-off-by: yhmo <[email protected]>
milvus-io · Mar 4, 2024 · 6a41d40 · 6a41d40
1 parent f8943d5
commit 6a41d40
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java b/src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java
@@ -559,10 +559,7 @@ public R<RpcStatus> loadCollection(@NonNull LoadCollectionParam requestParam) {
                 builder.setDbName(requestParam.getDatabaseName());
             }
 
-            LoadCollectionRequest loadCollectionRequest = builder
-                    .build();
-
-            Status response = blockingStub().loadCollection(loadCollectionRequest);
+            Status response = blockingStub().loadCollection(builder.build());
             handleResponse(title, response);
 
             // sync load, wait until collection finish loading
@@ -591,11 +588,13 @@ public R<RpcStatus> releaseCollection(@NonNull ReleaseCollectionParam requestPar
         String title = String.format("ReleaseCollectionRequest collectionName:%s", requestParam.getCollectionName());
 
         try {
-            ReleaseCollectionRequest releaseCollectionRequest = ReleaseCollectionRequest.newBuilder()
-                    .setCollectionName(requestParam.getCollectionName())
-                    .build();
+            ReleaseCollectionRequest.Builder builder = ReleaseCollectionRequest.newBuilder()
+                    .setCollectionName(requestParam.getCollectionName());
+            if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
+                builder.setDbName(requestParam.getDatabaseName());
+            }
 
-            Status response = blockingStub().releaseCollection(releaseCollectionRequest);
+            Status response = blockingStub().releaseCollection(builder.build());
             handleResponse(title, response);
             return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
         } catch (StatusRuntimeException e) {

diff --git a/src/main/java/io/milvus/param/collection/ReleaseCollectionParam.java b/src/main/java/io/milvus/param/collection/ReleaseCollectionParam.java
@@ -30,9 +30,11 @@
  */
 @Getter
 public class ReleaseCollectionParam {
+    private final String databaseName;
     private final String collectionName;
 
     private ReleaseCollectionParam(@NonNull Builder builder) {
+        this.databaseName = builder.databaseName;
         this.collectionName = builder.collectionName;
     }
 
@@ -44,11 +46,23 @@ public static Builder newBuilder() {
      * Builder for {@link ReleaseCollectionParam} class.
      */
     public static final class Builder {
+        private String databaseName;
         private String collectionName;
 
         private Builder() {
         }
 
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(String databaseName) {
+            this.databaseName = databaseName;
+            return this;
+        }
+
         /**
          * Sets the collection name. Collection name cannot be empty or null.
          *
@@ -80,6 +94,7 @@ public ReleaseCollectionParam build() throws ParamException {
     @Override
     public String toString() {
         return "ReleaseCollectionParam{" +
+                "databaseName='" + databaseName + '\'' +
                 "collectionName='" + collectionName + '\'' + '}';
     }
 }