Fix bug of describeCollection V2 (#1028)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

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

@@ -93,7 +93,7 @@ public interface MilvusClient {
     void setLogLevel(LogLevel level);
 
     /**
-     * Disconnects from a Milvus server with timeout of 1 minute
+     * Disconnects from a Milvus server with timeout of 1 second
      */
     default void close() {
         try {

src/main/java/io/milvus/pool/ClientPool.java

@@ -35,6 +35,15 @@ public class ClientPool<C, T> {
         this.clientPool = new GenericKeyedObjectPool<String, T>(clientFactory, poolConfig);
     }
 
+    /**
+     * Get a client object which is idle from the pool.
+     * Once the client is hold by the caller, it will be marked as active state and cannot be fetched by other caller.
+     * If the number of clients hits the MaxTotalPerKey value, this method will be blocked for MaxBlockWaitDuration.
+     * If no idle client available after MaxBlockWaitDuration, this method will return a null object to caller.
+     *
+     * @param key the key of a group where the client belong
+     * @return MilvusClient or MilvusClientV2
+     */
     public T getClient(String key) {
         try {
             return clientPool.borrowObject(key);
@@ -44,7 +53,14 @@ public class ClientPool<C, T> {
         }
     }
 
-
+    /**
+     * Return a client object. Once a client is returned, it becomes idle state and wait the next caller.
+     * The caller should ensure the client is returned. Otherwise, the client will keep in active state and cannot be used by the next caller.
+     * Throw exceptions if the key doesn't exist or the client is not belong to this key group.
+     *
+     * @param key the key of a group where the client belong
+     * @param grpcClient the client object to return
+     */
     public void returnClient(String key, T grpcClient) {
         try {
             clientPool.returnObject(key, grpcClient);
@@ -54,6 +70,10 @@ public class ClientPool<C, T> {
         }
     }
 
+    /**
+     * Release/disconnect all clients of all key groups, close the pool.
+     *
+     */
     public void close() {
         if (clientPool != null && !clientPool.isClosed()) {
             clientPool.close();
@@ -61,30 +81,57 @@ public class ClientPool<C, T> {
         }
     }
 
+    /**
+     * Release/disconnect idle clients of all key groups.
+     *
+     */
     public void clear() {
         if (clientPool != null && !clientPool.isClosed()) {
             clientPool.clear();
         }
     }
 
+    /**
+     * Release/disconnect idle clients of a key group.
+     *
+     * @param key the key of a group
+     */
     public void clear(String key) {
         if (clientPool != null && !clientPool.isClosed()) {
             clientPool.clear(key);
         }
     }
 
+    /**
+     * Return the number of idle clients of a key group
+     *
+     * @param key the key of a group
+     */
     public int getIdleClientNumber(String key) {
         return clientPool.getNumIdle(key);
     }
 
+    /**
+     * Return the number of active clients of a key group
+     *
+     * @param key the key of a group
+     */
     public int getActiveClientNumber(String key) {
         return clientPool.getNumActive(key);
     }
 
+    /**
+     * Return the number of idle clients of all key group
+     *
+     */
     public int getTotalIdleClientNumber() {
         return clientPool.getNumIdle();
     }
 
+    /**
+     * Return the number of active clients of all key group
+     *
+     */
     public int getTotalActiveClientNumber() {
         return clientPool.getNumActive();
     }

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

@@ -734,12 +734,12 @@ public class MilvusClientV2 {
      *
      * @return String
      */
-    public String getVersion() {
-        return retry(()->clientUtils.getVersion(this.blockingStub));
+    public String getServerVersion() {
+        return retry(()->clientUtils.getServerVersion(this.blockingStub));
     }
 
     /**
-     * close client
+     * Disconnects from a Milvus server with configurable timeout
      *
      * @param maxWaitSeconds max wait seconds
      * @throws InterruptedException throws InterruptedException if the client failed to close connection
@@ -751,6 +751,18 @@ public class MilvusClientV2 {
         }
     }
 
+    /**
+     * Disconnects from a Milvus server with timeout of 1 second
+     *
+     */
+    public void close() {
+        try {
+            close(TimeUnit.MINUTES.toSeconds(1));
+        } catch (InterruptedException e) {
+            System.out.println("Interrupted during shutdown Milvus client!");
+        }
+    }
+
     public boolean clientIsReady() {
         return channel != null && !channel.isShutdown() && !channel.isTerminated();
     }

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

@@ -34,7 +34,9 @@ import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 public class CollectionService extends BaseService {
     public IndexService indexService = new IndexService();
@@ -211,25 +213,7 @@ public class CollectionService extends BaseService {
                 .build();
         DescribeCollectionResponse response = blockingStub.describeCollection(describeCollectionRequest);
         rpcUtils.handleResponse(title, response.getStatus());
-        return convertDescCollectionResp(response);
-    }
-
-    public static DescribeCollectionResp convertDescCollectionResp(DescribeCollectionResponse response) {
-        DescribeCollectionResp describeCollectionResp = DescribeCollectionResp.builder()
-                .collectionName(response.getCollectionName())
-                .databaseName(response.getDbName())
-                .description(response.getSchema().getDescription())
-                .numOfPartitions(response.getNumPartitions())
-                .collectionSchema(SchemaUtils.convertFromGrpcCollectionSchema(response.getSchema()))
-                .autoID(response.getSchema().getFieldsList().stream().anyMatch(FieldSchema::getAutoID))
-                .enableDynamicField(response.getSchema().getEnableDynamicField())
-                .fieldNames(response.getSchema().getFieldsList().stream().map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()))
-                .vectorFieldNames(response.getSchema().getFieldsList().stream().filter(fieldSchema -> ParamUtils.isVectorDataType(fieldSchema.getDataType())).map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()))
-                .primaryFieldName(response.getSchema().getFieldsList().stream().filter(FieldSchema::getIsPrimaryKey).map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()).get(0))
-                .createTime(response.getCreatedTimestamp())
-                .consistencyLevel(io.milvus.v2.common.ConsistencyLevel.valueOf(response.getConsistencyLevel().name().toUpperCase()))
-                .build();
-        return describeCollectionResp;
+        return convertUtils.convertDescCollectionResp(response);
     }
 
     public Void renameCollection(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, RenameCollectionReq request) {

src/main/java/io/milvus/v2/service/collection/response/DescribeCollectionResp.java

@@ -21,10 +21,13 @@ package io.milvus.v2.service.collection.response;
 
 import io.milvus.v2.common.ConsistencyLevel;
 import io.milvus.v2.service.collection.request.CreateCollectionReq;
+import lombok.Builder;
 import lombok.Data;
 import lombok.experimental.SuperBuilder;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 @Data
 @SuperBuilder
@@ -43,4 +46,6 @@ public class DescribeCollectionResp {
     private CreateCollectionReq.CollectionSchema collectionSchema;
     private Long createTime;
     private ConsistencyLevel consistencyLevel;
+    @Builder.Default
+    private final Map<String, String> properties = new HashMap<>();
 }

src/main/java/io/milvus/v2/service/vector/VectorService.java

@@ -184,7 +184,7 @@ public class VectorService extends BaseService {
     public QueryIterator queryIterator(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                            QueryIteratorReq request) {
         DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, "", request.getCollectionName());
-        DescribeCollectionResp respR = CollectionService.convertDescCollectionResp(descResp);
+        DescribeCollectionResp respR = convertUtils.convertDescCollectionResp(descResp);
         CreateCollectionReq.FieldSchema pkField = respR.getCollectionSchema().getField(respR.getPrimaryFieldName());
         return new QueryIterator(request, blockingStub, pkField);
     }
@@ -192,7 +192,7 @@ public class VectorService extends BaseService {
     public SearchIterator searchIterator(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
                                             SearchIteratorReq request) {
         DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, "", request.getCollectionName());
-        DescribeCollectionResp respR = CollectionService.convertDescCollectionResp(descResp);
+        DescribeCollectionResp respR = convertUtils.convertDescCollectionResp(descResp);
         CreateCollectionReq.FieldSchema pkField = respR.getCollectionSchema().getField(respR.getPrimaryFieldName());
         return new SearchIterator(request, blockingStub, pkField);
     }
@@ -205,7 +205,7 @@ public class VectorService extends BaseService {
         }
 
         DescribeCollectionResponse descResp = getCollectionInfo(blockingStub, "", request.getCollectionName());
-        DescribeCollectionResp respR = CollectionService.convertDescCollectionResp(descResp);
+        DescribeCollectionResp respR = convertUtils.convertDescCollectionResp(descResp);
         if (request.getFilter() == null) {
             request.setFilter(vectorUtils.getExprById(respR.getPrimaryFieldName(), request.getIds()));
         }

src/main/java/io/milvus/v2/utils/ClientUtils.java

@@ -127,7 +127,7 @@ public class ClientUtils {
             throw new IllegalArgumentException("Database " + dbName + " not exist");
         }
     }
-    public String getVersion(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
+    public String getServerVersion(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
         GetVersionResponse response = blockingStub.getVersion(GetVersionRequest.newBuilder().build());
         rpcUtils.handleResponse("Get server version", response.getStatus());
         return response.getVersion();

src/main/java/io/milvus/v2/utils/ConvertUtils.java

@@ -21,10 +21,12 @@ package io.milvus.v2.utils;
 
 import io.milvus.grpc.*;
 import io.milvus.param.Constant;
+import io.milvus.param.ParamUtils;
 import io.milvus.response.QueryResultsWrapper;
 import io.milvus.response.SearchResultsWrapper;
 import io.milvus.v2.common.IndexBuildState;
 import io.milvus.v2.common.IndexParam;
+import io.milvus.v2.service.collection.response.DescribeCollectionResp;
 import io.milvus.v2.service.index.response.DescribeIndexResp;
 import io.milvus.v2.service.vector.response.QueryResp;
 import io.milvus.v2.service.vector.response.SearchResp;
@@ -112,4 +114,26 @@ public class ConvertUtils {
 
         return DescribeIndexResp.builder().indexDescriptions(descs).build();
     }
+
+    public DescribeCollectionResp convertDescCollectionResp(DescribeCollectionResponse response) {
+        Map<String, String> properties = new HashMap<>();
+        response.getPropertiesList().forEach(prop->properties.put(prop.getKey(), prop.getValue()));
+
+        DescribeCollectionResp describeCollectionResp = DescribeCollectionResp.builder()
+                .collectionName(response.getCollectionName())
+                .databaseName(response.getDbName())
+                .description(response.getSchema().getDescription())
+                .numOfPartitions(response.getNumPartitions())
+                .collectionSchema(SchemaUtils.convertFromGrpcCollectionSchema(response.getSchema()))
+                .autoID(response.getSchema().getFieldsList().stream().anyMatch(FieldSchema::getAutoID))
+                .enableDynamicField(response.getSchema().getEnableDynamicField())
+                .fieldNames(response.getSchema().getFieldsList().stream().map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()))
+                .vectorFieldNames(response.getSchema().getFieldsList().stream().filter(fieldSchema -> ParamUtils.isVectorDataType(fieldSchema.getDataType())).map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()))
+                .primaryFieldName(response.getSchema().getFieldsList().stream().filter(FieldSchema::getIsPrimaryKey).map(FieldSchema::getName).collect(java.util.stream.Collectors.toList()).get(0))
+                .createTime(response.getCreatedTimestamp())
+                .consistencyLevel(io.milvus.v2.common.ConsistencyLevel.valueOf(response.getConsistencyLevel().name().toUpperCase()))
+                .properties(properties)
+                .build();
+        return describeCollectionResp;
+    }
 }

src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java

@@ -37,6 +37,7 @@ import io.milvus.v2.exception.MilvusClientException;
 import io.milvus.v2.service.collection.request.*;
 import io.milvus.v2.service.collection.response.DescribeCollectionResp;
 import io.milvus.v2.service.collection.response.ListCollectionsResp;
+import io.milvus.v2.service.database.request.AlterDatabaseReq;
 import io.milvus.v2.service.database.request.CreateDatabaseReq;
 import io.milvus.v2.service.database.request.DescribeDatabaseReq;
 import io.milvus.v2.service.database.request.DropDatabaseReq;
@@ -1199,10 +1200,21 @@ class MilvusClientV2DockerTest {
         Map<String, String> properties = new HashMap<>();
         properties.put(Constant.TTL_SECONDS, "10");
         properties.put(Constant.MMAP_ENABLED, "true");
+        properties.put("prop", "val");
         client.alterCollection(AlterCollectionReq.builder()
                 .collectionName(randomCollectionName)
                 .properties(properties)
                 .build());
+        DescribeCollectionResp descCollResp = client.describeCollection(DescribeCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+        Map<String, String> collProps = descCollResp.getProperties();
+        Assertions.assertTrue(collProps.containsKey(Constant.TTL_SECONDS));
+        Assertions.assertTrue(collProps.containsKey(Constant.MMAP_ENABLED));
+        Assertions.assertTrue(collProps.containsKey("prop"));
+        Assertions.assertEquals("10", collProps.get(Constant.TTL_SECONDS));
+        Assertions.assertEquals("true", collProps.get(Constant.MMAP_ENABLED));
+        Assertions.assertEquals("val", collProps.get("prop"));
 
         DescribeIndexResp descResp = client.describeIndex(DescribeIndexReq.builder()
                 .collectionName(randomCollectionName)
@@ -1214,13 +1226,22 @@ class MilvusClientV2DockerTest {
         Assertions.assertEquals(IndexParam.IndexType.AUTOINDEX, desc.getIndexType());
 
         properties.clear();
-        properties.put(Constant.MMAP_ENABLED, "true");
+        properties.put(Constant.MMAP_ENABLED, "false");
         client.alterIndex(AlterIndexReq.builder()
                 .collectionName(randomCollectionName)
                 .indexName(desc.getIndexName())
                 .properties(properties)
                 .build());
 
+        descResp = client.describeIndex(DescribeIndexReq.builder()
+                .collectionName(randomCollectionName)
+                .fieldName("vector")
+                .build());
+        desc = descResp.getIndexDescByFieldName("vector");
+        Map<String, String> indexProps = desc.getExtraParams();
+        Assertions.assertTrue(indexProps.containsKey(Constant.MMAP_ENABLED));
+        Assertions.assertEquals("false", indexProps.get(Constant.MMAP_ENABLED));
+
         client.dropIndex(DropIndexReq.builder()
                 .collectionName(randomCollectionName)
                 .fieldName("vector")
@@ -1541,6 +1562,22 @@ class MilvusClientV2DockerTest {
         Assertions.assertTrue(propertiesResp.containsKey(Constant.DATABASE_REPLICA_NUMBER));
         Assertions.assertEquals("5", propertiesResp.get(Constant.DATABASE_REPLICA_NUMBER));
 
+        // alter the database
+        properties.put(Constant.DATABASE_REPLICA_NUMBER, "10");
+        properties.put("prop", "val");
+        client.alterDatabase(AlterDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .properties(properties)
+                .build());
+        descDBResp = client.describeDatabase(DescribeDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .build());
+        propertiesResp = descDBResp.getProperties();
+        Assertions.assertTrue(propertiesResp.containsKey(Constant.DATABASE_REPLICA_NUMBER));
+        Assertions.assertEquals("10", propertiesResp.get(Constant.DATABASE_REPLICA_NUMBER));
+        Assertions.assertTrue(propertiesResp.containsKey("prop"));
+        Assertions.assertEquals("val", propertiesResp.get("prop"));
+
         // switch to the new database
         Assertions.assertDoesNotThrow(()->client.useDatabase(tempDatabaseName));
 
@@ -1609,7 +1646,7 @@ class MilvusClientV2DockerTest {
                 Thread t = new Thread(() -> {
                     for (int i = 0; i < requestPerThread; i++) {
                         MilvusClientV2 client = pool.getClient(key);
-                        String version = client.getVersion();
+                        String version = client.getServerVersion();
 //                            System.out.printf("%d, %s%n", i, version);
                         System.out.printf("idle %d, active %d%n", pool.getIdleClientNumber(key), pool.getActiveClientNumber(key));
                         pool.returnClient(key, client);