enhance: enable setting properties during create database/collection (#972)

Signed-off-by: Wei Liu <wei.liu@zilliz.com>

src/main/java/io/milvus/client/AbstractMilvusGrpcClient.java

@@ -445,11 +445,15 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         String title = String.format("CreateDatabaseRequest databaseName:%s", requestParam.getDatabaseName());
 
         try {
+            List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(requestParam.getProperties());
             // Construct CreateDatabaseRequest
             CreateDatabaseRequest createDatabaseRequest = CreateDatabaseRequest.newBuilder()
                     .setDbName(requestParam.getDatabaseName())
+                    .addAllProperties(propertiesList)
                     .build();
 
+            System.out.println(requestParam.getProperties());
+
             Status response = blockingStub().createDatabase(createDatabaseRequest);
             handleResponse(title, response);
             return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
@@ -588,12 +592,14 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
                 collectionSchemaBuilder.addFields(ParamUtils.ConvertField(fieldType));
             }
 
+            List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(requestParam.getProperties());
             // Construct CreateCollectionRequest
             CreateCollectionRequest.Builder builder = CreateCollectionRequest.newBuilder()
                     .setCollectionName(requestParam.getCollectionName())
                     .setShardsNum(requestParam.getShardsNum())
                     .setConsistencyLevelValue(requestParam.getConsistencyLevel().getCode())
-                    .setSchema(collectionSchemaBuilder.build().toByteString());
+                    .setSchema(collectionSchemaBuilder.build().toByteString())
+                    .addAllProperties(propertiesList);
             if (StringUtils.isNotEmpty(requestParam.getDatabaseName())) {
                 builder.setDbName(requestParam.getDatabaseName());
             }

src/main/java/io/milvus/param/Constant.java

@@ -53,6 +53,9 @@ public class Constant {
     public static final String DATABASE_REPLICA_NUMBER = "database.replica.number";
     public static final String DATABASE_RESOURCE_GROUPS = "database.resource_groups";
 
+    public static final String COLLECTION_REPLICA_NUMBER = "collection.replica.number";
+    public static final String COLLECTION_RESOURCE_GROUPS = "collection.resource_groups";
+
     // max value for waiting loading collection/partition interval, unit: millisecond
     public static final Long MAX_WAITING_LOADING_INTERVAL = 2000L;
 

src/main/java/io/milvus/param/collection/AlterDatabaseParam.java

@@ -83,7 +83,7 @@ public class AlterDatabaseParam {
         /**
          * Sets the resource groups in database level, then if load collection doesn't have resource groups, it will use this resource groups.
          * @param resourceGroups resource group names
-         * @return
+         * @return <code>Builder</code>
          */
         public Builder WithResourceGroups(@NonNull List<String> resourceGroups) {
             return this.withProperty(Constant.DATABASE_RESOURCE_GROUPS, String.join(",", resourceGroups));

src/main/java/io/milvus/param/collection/CreateCollectionParam.java

@@ -21,13 +21,16 @@ package io.milvus.param.collection;
 
 import io.milvus.common.clientenum.ConsistencyLevelEnum;
 import io.milvus.exception.ParamException;
+import io.milvus.param.Constant;
 import io.milvus.param.ParamUtils;
 import lombok.Getter;
 import lombok.NonNull;
 import lombok.ToString;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Parameters for <code>createCollection</code> interface.
@@ -44,6 +47,7 @@ public class CreateCollectionParam {
 
     private final List<FieldType> fieldTypes;
     private final boolean enableDynamicField;
+    private final Map<String, String> properties = new HashMap<>();
 
     private CreateCollectionParam(@NonNull Builder builder) {
         this.collectionName = builder.collectionName;
@@ -54,6 +58,7 @@ public class CreateCollectionParam {
         this.consistencyLevel = builder.consistencyLevel;
         this.databaseName = builder.databaseName;
         this.enableDynamicField = builder.enableDynamicField;
+        this.properties.putAll(builder.properties);
     }
 
     public static Builder newBuilder() {
@@ -74,6 +79,9 @@ public class CreateCollectionParam {
         private CollectionSchemaParam schema;
 
         private boolean enableDynamicField;
+
+        private final Map<String, String> properties = new HashMap<>();
+
         private Builder() {
         }
 
@@ -202,6 +210,37 @@ public class CreateCollectionParam {
             return this;
         }
 
+        /**
+         * Sets the replica number in collection level, then if load collection doesn't have replica number, it will use this replica number.
+         * @param replicaNumber replica number
+         * @return <code>Builder</code>
+         */
+        public Builder withReplicaNumber(int replicaNumber) {
+            return this.withProperty(Constant.COLLECTION_REPLICA_NUMBER, Integer.toString(replicaNumber));
+        }
+
+        /**
+         * Sets the resource groups in collection level, then if load collection doesn't have resource groups, it will use this resource groups.
+         * @param resourceGroups resource group names
+         * @return <code>Builder</code>
+         */
+        public Builder withResourceGroups(@NonNull List<String> resourceGroups) {
+            return this.withProperty(Constant.COLLECTION_RESOURCE_GROUPS, String.join(",", resourceGroups));
+
+        }
+
+        /**
+         * Basic method to set a key-value property.
+         *
+         * @param key the key
+         * @param value the value
+         * @return <code>Builder</code>
+         */
+        public Builder withProperty(@NonNull String key, @NonNull String value) {
+            this.properties.put(key, value);
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link CreateCollectionParam} instance.
          *

src/main/java/io/milvus/param/collection/CreateDatabaseParam.java

@@ -20,11 +20,16 @@
 package io.milvus.param.collection;
 
 import io.milvus.exception.ParamException;
+import io.milvus.param.Constant;
 import io.milvus.param.ParamUtils;
 import lombok.Getter;
 import lombok.NonNull;
 import lombok.ToString;
 
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 /**
  * Parameters for <code>createDatabase</code> interface.
  */
@@ -32,9 +37,11 @@ import lombok.ToString;
 @ToString
 public class CreateDatabaseParam {
     private final String databaseName;
+    private final Map<String, String> properties = new HashMap<>();
 
     private CreateDatabaseParam(@NonNull Builder builder) {
         this.databaseName = builder.databaseName;
+        this.properties.putAll(builder.properties);
     }
 
     public static Builder newBuilder() {
@@ -47,6 +54,8 @@ public class CreateDatabaseParam {
     public static final class Builder {
         private String databaseName;
 
+        private final Map<String, String> properties = new HashMap<>();
+
         private Builder() {
         }
 
@@ -61,6 +70,37 @@ public class CreateDatabaseParam {
             return this;
         }
 
+         /**
+         * Sets the replica number in database level, then if load collection doesn't have replica number, it will use this replica number.
+         * @param replicaNumber replica number
+         * @return <code>Builder</code>
+         */
+        public Builder withReplicaNumber(int replicaNumber) {
+            return this.withProperty(Constant.DATABASE_REPLICA_NUMBER, Integer.toString(replicaNumber));
+        }
+
+        /**
+         * Sets the resource groups in database level, then if load collection doesn't have resource groups, it will use this resource groups.
+         * @param resourceGroups resource group names
+         * @return <code>Builder</code>
+         */
+        public Builder withResourceGroups(@NonNull List<String> resourceGroups) {
+            return this.withProperty(Constant.DATABASE_RESOURCE_GROUPS, String.join(",", resourceGroups));
+
+        }
+
+        /**
+         * Basic method to set a key-value property.
+         *
+         * @param key the key
+         * @param value the value
+         * @return <code>Builder</code>
+         */
+        public Builder withProperty(@NonNull String key, @NonNull String value) {
+            this.properties.put(key, value);
+            return this;
+        }
+
         /**
          * Verifies parameters and creates a new {@link CreateDatabaseParam} instance.
          *

src/main/java/io/milvus/response/DescCollResponseWrapper.java

@@ -22,10 +22,12 @@ package io.milvus.response;
 import io.milvus.common.clientenum.ConsistencyLevelEnum;
 import io.milvus.exception.ParamException;
 import io.milvus.grpc.*;
+import io.milvus.param.Constant;
 import io.milvus.param.ParamUtils;
 import io.milvus.param.collection.CollectionSchemaParam;
 import io.milvus.param.collection.FieldType;
 import lombok.NonNull;
+import org.testcontainers.shaded.org.bouncycastle.util.Strings;
 
 import java.util.*;
 
@@ -35,8 +37,11 @@ import java.util.*;
 public class DescCollResponseWrapper {
     private final DescribeCollectionResponse response;
 
+    Map<String, String> pairs = new HashMap<>();
+
     public DescCollResponseWrapper(@NonNull DescribeCollectionResponse response) {
         this.response = response;
+        response.getPropertiesList().forEach((prop) -> pairs.put(prop.getKey(), prop.getValue()));
     }
 
     /**
@@ -249,10 +254,6 @@ public class DescCollResponseWrapper {
      * @return List of String, aliases of the collection
      */
     public Map<String, String> getProperties() {
-        Map<String, String> pairs = new HashMap<>();
-        List<KeyValuePair> props = response.getPropertiesList();
-        props.forEach((prop) -> pairs.put(prop.getKey(), prop.getValue()));
-
         return pairs;
     }
 
@@ -268,6 +269,34 @@ public class DescCollResponseWrapper {
                 .build();
     }
 
+
+
+    /**
+     * return collection resource groups
+     *
+     * @return resource group names
+     */
+    public List<String> getResourceGroups() {
+        String value = pairs.get(Constant.COLLECTION_RESOURCE_GROUPS);
+        if (value == null) {
+            return new ArrayList<>();
+        }
+        return Arrays.asList(Strings.split(value, ','));
+    }
+
+    /**
+     * return collection replica number
+     *
+     * @return replica number
+     */
+    public int getReplicaNumber() {
+        String value = pairs.get(Constant.COLLECTION_REPLICA_NUMBER);
+        if (value == null) {
+            return 0;
+        }
+        return Integer.parseInt(pairs.get(Constant.COLLECTION_REPLICA_NUMBER));
+    }
+
     /**
      * Construct a <code>String</code> by {@link DescCollResponseWrapper} instance.
      *

src/main/milvus-proto

@@ -1 +1 @@
-Subproject commit b5c6a727ed3cc665a62684da218e4f211c82f9f3
+Subproject commit 073886363680921ac3659c0ead9e7408219ce29b

src/test/java/io/milvus/client/MilvusClientDockerTest.java

@@ -82,7 +82,7 @@ class MilvusClientDockerTest {
     private static final Random RANDOM = new Random();
 
     @Container
-    private static final MilvusContainer milvus = new MilvusContainer("milvusdb/milvus:v2.4.4");
+    private static final MilvusContainer milvus = new MilvusContainer("milvusdb/milvus:master-20240709-314f4d99");
 
     @BeforeAll
     public static void setUp() {
@@ -519,6 +519,8 @@ class MilvusClientDockerTest {
                 .withSchema(schema)
                 .withShardsNum(3)
                 .withConsistencyLevel(ConsistencyLevelEnum.EVENTUALLY)
+                .withReplicaNumber(1)
+                .withResourceGroups(Arrays.asList("rg1"))
                 .build();
 
         R<RpcStatus> createR = client.createCollection(createParam);
@@ -540,6 +542,8 @@ class MilvusClientDockerTest {
         CollectionSchemaParam fetchSchema = collDescWrapper.getSchema();
         Assertions.assertFalse(fetchSchema.isEnableDynamicField());
         Assertions.assertEquals(ConsistencyLevelEnum.EVENTUALLY, collDescWrapper.getConsistencyLevel());
+        Assertions.assertEquals(1, collDescWrapper.getReplicaNumber());
+        Assertions.assertEquals(1, collDescWrapper.getResourceGroups().size());
         System.out.println(collDescWrapper);
 
         R<ShowPartitionsResponse> spResp = client.showPartitions(ShowPartitionsParam.newBuilder()
@@ -2991,7 +2995,7 @@ class MilvusClientDockerTest {
     @Test
     void testDatabase() {
         String dbName = "test_database";
-        CreateDatabaseParam createDatabaseParam = CreateDatabaseParam.newBuilder().withDatabaseName(dbName).build();
+        CreateDatabaseParam createDatabaseParam = CreateDatabaseParam.newBuilder().withDatabaseName(dbName).withReplicaNumber(1).withResourceGroups(Arrays.asList("rg1")).build();
         R<RpcStatus> createResponse = client.createDatabase(createDatabaseParam);
         Assertions.assertEquals(R.Status.Success.getCode(), createResponse.getStatus().intValue());
 
@@ -3001,8 +3005,8 @@ class MilvusClientDockerTest {
         Assertions.assertEquals(R.Status.Success.getCode(), describeResponse.getStatus().intValue());
         DescDBResponseWrapper describeDBWrapper = new DescDBResponseWrapper(describeResponse.getData());
         Assertions.assertEquals(dbName, describeDBWrapper.getDatabaseName());
-        Assertions.assertEquals(0, describeDBWrapper.getReplicaNumber());
-        Assertions.assertEquals(0, describeDBWrapper.getResourceGroups().size());
+        Assertions.assertEquals(1, describeDBWrapper.getReplicaNumber());
+        Assertions.assertEquals(1, describeDBWrapper.getResourceGroups().size());
 
         // alter database props
         AlterDatabaseParam alterDatabaseParam = AlterDatabaseParam.newBuilder().withDatabaseName(dbName).withReplicaNumber(3).WithResourceGroups(Arrays.asList("rg1", "rg2", "rg3")).build();