enhance: enable AlterDatabase and DescribeDatabase API (#923)

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

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

@@ -462,6 +462,59 @@ public abstract class AbstractMilvusGrpcClient implements MilvusClient {
         }
     }
 
+    @Override
+    public R<RpcStatus> alterDatabase(AlterDatabaseParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logDebug(requestParam.toString());
+        String title = String.format("AlterDatabaseRequest databaseName:%s", requestParam.getDatabaseName());
+
+        try {
+            List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(requestParam.getProperties());
+            AlterDatabaseRequest alterDatabaseRequest = AlterDatabaseRequest.newBuilder()
+                .setDbName(requestParam.getDatabaseName())
+                .addAllProperties(propertiesList)
+                .build();
+
+            Status response = blockingStub().alterDatabase(alterDatabaseRequest);
+            handleResponse(title, response);
+            return R.success(new RpcStatus(RpcStatus.SUCCESS_MSG));
+        } catch (StatusRuntimeException e) {
+            logError("{} RPC failed! Exception:{}", title, e);
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("{} failed! Exception:{}", title, e);
+            return R.failed(e);
+        }
+    }
+
+    @Override
+    public R<DescribeDatabaseResponse> describeDatabase(DescribeDatabaseParam requestParam) {
+        if (!clientIsReady()) {
+            return R.failed(new ClientNotConnectedException("Client rpc channel is not ready"));
+        }
+
+        logDebug(requestParam.toString());
+        String title = String.format("DescribeDatabaseRequest databaseName:%s", requestParam.getDatabaseName());
+        try {
+            DescribeDatabaseRequest describeDatabaseRequest = DescribeDatabaseRequest.newBuilder()
+                .setDbName(requestParam.getDatabaseName())
+                .build();
+
+            DescribeDatabaseResponse response = blockingStub().describeDatabase(describeDatabaseRequest);
+            handleResponse(title, response.getStatus());
+            return R.success(response);
+        } catch (StatusRuntimeException e) {
+            logError("{} RPC failed! Exception:{}", title, e);
+            return R.failed(e);
+        } catch (Exception e) {
+            logError("{} failed! Exception:{}", title, e);
+            return R.failed(e);
+        }
+    }
+
     @Override
     public R<RpcStatus> createCollection(@NonNull CreateCollectionParam requestParam) {
         if (!clientIsReady()) {

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

@@ -137,6 +137,20 @@ public interface MilvusClient {
      */
     R<ListDatabasesResponse> listDatabases();
 
+    /**
+     * Alter database with key value pair
+     * @param requestParam {@link AlterDatabaseParam}
+     * @return  {status:result code, data:RpcStatus{msg: result message}}
+     */
+    R<RpcStatus> alterDatabase(AlterDatabaseParam requestParam);
+
+    /**
+     * show detail of database base, such as replica number and resource groups
+     * @param requestParam {@link DescribeDatabaseParam}
+     * @return {status:result code, data:DescribeDatabaseResponse{replica_number,resource_groups}}
+     */
+    R<DescribeDatabaseResponse> describeDatabase(DescribeDatabaseParam requestParam);
+
     /**
      * Creates a collection in Milvus.
      *

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

@@ -163,6 +163,22 @@ public class MilvusMultiServiceClient implements MilvusClient {
         return handleResponse(response);
     }
 
+    @Override
+    public R<RpcStatus> alterDatabase(AlterDatabaseParam requestParam) {
+        List<R<RpcStatus>> response = this.clusterFactory.getAvailableServerSettings().stream()
+                .map(serverSetting -> serverSetting.getClient().alterDatabase(requestParam))
+                .collect(Collectors.toList());
+        return handleResponse(response);
+    }
+
+    @Override
+    public R<DescribeDatabaseResponse> describeDatabase(DescribeDatabaseParam requestParam) {
+        List<R<DescribeDatabaseResponse>> response = this.clusterFactory.getAvailableServerSettings().stream()
+                .map(serverSetting -> serverSetting.getClient().describeDatabase(requestParam))
+                .collect(Collectors.toList());
+    return handleResponse(response);
+    }
+
     @Override
     public R<RpcStatus> createCollection(CreateCollectionParam requestParam) {
         List<R<RpcStatus>> response = this.clusterFactory.getAvailableServerSettings().stream()

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

@@ -439,6 +439,16 @@ public class MilvusServiceClient extends AbstractMilvusGrpcClient {
         return retry(super::listDatabases);
     }
 
+    @Override
+    public R<RpcStatus> alterDatabase(AlterDatabaseParam requestParam) {
+        return retry(()-> super.alterDatabase(requestParam));
+    }
+
+    @Override
+    public R<DescribeDatabaseResponse> describeDatabase(DescribeDatabaseParam requestParam) {
+        return retry(()-> super.describeDatabase(requestParam));
+    }
+
     @Override
     public R<RpcStatus> createCollection(CreateCollectionParam requestParam) {
         return retry(()-> super.createCollection(requestParam));

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

@@ -50,6 +50,8 @@ public class Constant {
     // constant values for general
     public static final String TTL_SECONDS = "collection.ttl.seconds";
     public static final String MMAP_ENABLED = "mmap.enabled";
+    public static final String DATABASE_REPLICA_NUMBER = "database.replica.number";
+    public static final String DATABASE_RESOURCE_GROUPS = "database.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

@@ -0,0 +1,117 @@
+/*
+ * 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.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>alterDatabase</code> interface.
+ */
+@Getter
+@ToString
+public class AlterDatabaseParam {
+    private final String databaseName;
+    private final Map<String, String> properties = new HashMap<>();
+
+    private AlterDatabaseParam(@NonNull Builder builder) {
+        this.databaseName = builder.databaseName;
+        this.properties.putAll(builder.properties);
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Builder for {@link AlterDatabaseParam} class.
+     */
+    public static final class Builder {
+        private String databaseName;
+
+        private final Map<String, String> properties = new HashMap<>();
+
+
+        private Builder() {
+        }
+
+        /**
+         * Sets the database name. database name can be nil.
+         *
+         * @param databaseName database name
+         * @return <code>Builder</code>
+         */
+        public Builder withDatabaseName(@NonNull String databaseName) {
+            this.databaseName = databaseName;
+            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
+         */
+        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 AlterDatabaseParam} instance.
+         *
+         * @return {@link AlterDatabaseParam}
+         */
+        public AlterDatabaseParam build() throws ParamException {
+            ParamUtils.CheckNullEmptyString(databaseName, "Database name");
+
+            return new AlterDatabaseParam(this);
+        }
+    }
+
+}

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

@@ -0,0 +1,77 @@
+/*
+ * 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.param.collection;
+
+ import io.milvus.exception.ParamException;
+ import io.milvus.param.ParamUtils;
+ 
+ import lombok.Getter;
+ import lombok.NonNull;
+ import lombok.ToString;
+ 
+ /**
+  * Parameters for <code>describeDatabase</code> interface.
+  */
+ @Getter
+ @ToString
+ public class DescribeDatabaseParam {
+     private final String databaseName;
+ 
+     private DescribeDatabaseParam(@NonNull Builder builder) {
+         this.databaseName = builder.databaseName;
+     }
+ 
+     public static Builder newBuilder() {
+         return new Builder();
+     }
+ 
+     /**
+      * Builder for {@link DescribeDatabaseParam} class.
+      */
+     public static final class Builder {
+         private String databaseName;
+ 
+         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;
+         }
+ 
+         /**
+          * Verifies parameters and creates a new {@link DescribeDatabaseParam} instance.
+          *
+          * @return {@link DescribeDatabaseParam}
+          */
+         public DescribeDatabaseParam build() throws ParamException {
+            ParamUtils.CheckNullEmptyString(databaseName, "Database name");
+ 
+             return new DescribeDatabaseParam(this);
+         }
+     }
+ 
+ }

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

@@ -0,0 +1,103 @@
+/*
+ * 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.response;
+
+import io.milvus.grpc.*;
+import io.milvus.param.Constant;
+import lombok.NonNull;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.testcontainers.shaded.org.bouncycastle.util.Strings;
+
+/**
+ * Util class to wrap response of <code>describeCollection</code> interface.
+ */
+public class DescDBResponseWrapper {
+    private final DescribeDatabaseResponse response;
+    Map<String, String> pairs = new HashMap<>();
+
+    public DescDBResponseWrapper(@NonNull DescribeDatabaseResponse response) {
+        this.response = response;
+        response.getPropertiesList().forEach((prop) -> pairs.put(prop.getKey(), prop.getValue()));
+    }
+
+
+    /**
+     * get database name
+     *
+     * @return database name
+     */
+    public String getDatabaseName() {
+        return response.getDbName();
+    }
+
+    /**
+     * Get properties of the collection.
+     *
+     * @return map of key value pair
+     */
+    public Map<String, String> getProperties() {
+        return pairs;
+    }
+
+    /**
+     * return database resource groups
+     *
+     * @return resource group names
+     */
+    public List<String> getResourceGroups() {
+        String value = pairs.get(Constant.DATABASE_RESOURCE_GROUPS);
+        if (value == null) {
+            return new ArrayList<>();
+        }
+        return Arrays.asList(Strings.split(value, ','));
+    }
+
+    /**
+     * return database replica number
+     *
+     * @return database replica number
+     */
+    public int getReplicaNumber() {
+        String value = pairs.get(Constant.DATABASE_REPLICA_NUMBER);
+        if (value == null) {
+            return 0;
+        }
+        return Integer.parseInt(pairs.get(Constant.DATABASE_REPLICA_NUMBER));
+    }
+
+    /**
+     * Construct a <code>String</code> by {@link DescCollResponseWrapper} instance.
+     *
+     * @return <code>String</code>
+     */
+    @Override
+    public String toString() {
+        return "Database Description{" +
+            "name:'" + getDatabaseName() + '\'' +
+            ", properties:" + getProperties() +
+            '}';
+    }
+}

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

@@ -79,7 +79,7 @@ class MilvusClientDockerTest {
     protected static final Gson GSON_INSTANCE = new Gson();
 
     @Container
-    private static final MilvusContainer milvus = new MilvusContainer("milvusdb/milvus:v2.4.1");
+    private static final MilvusContainer milvus = new MilvusContainer("milvusdb/milvus:v2.4.4");
 
     @BeforeAll
     public static void setUp() {
@@ -2817,4 +2817,39 @@ class MilvusClientDockerTest {
         }
         Assertions.assertEquals(50, counter);
     }
+
+    @Test
+    void testDatabase() {
+        String dbName = "test_database";
+        CreateDatabaseParam createDatabaseParam = CreateDatabaseParam.newBuilder().withDatabaseName(dbName).build();
+        R<RpcStatus> createResponse = client.createDatabase(createDatabaseParam);
+        Assertions.assertEquals(R.Status.Success.getCode(), createResponse.getStatus().intValue());
+
+        // check database props
+        DescribeDatabaseParam describeDBParam = DescribeDatabaseParam.newBuilder().withDatabaseName(dbName).build();
+        R<DescribeDatabaseResponse> describeResponse = client.describeDatabase(describeDBParam);
+        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());
+
+        // alter database props
+        AlterDatabaseParam alterDatabaseParam = AlterDatabaseParam.newBuilder().withDatabaseName(dbName).withReplicaNumber(3).WithResourceGroups(Arrays.asList("rg1", "rg2", "rg3")).build();
+        R<RpcStatus> alterDatabaseResponse = client.alterDatabase(alterDatabaseParam);
+        Assertions.assertEquals(R.Status.Success.getCode(), alterDatabaseResponse.getStatus().intValue());
+
+        // check database props
+        describeResponse = client.describeDatabase(describeDBParam);
+        Assertions.assertEquals(R.Status.Success.getCode(), describeResponse.getStatus().intValue());
+        describeDBWrapper = new DescDBResponseWrapper(describeResponse.getData());
+        Assertions.assertEquals(dbName, describeDBWrapper.getDatabaseName());
+        Assertions.assertEquals(3, describeDBWrapper.getReplicaNumber());
+        Assertions.assertEquals(3, describeDBWrapper.getResourceGroups().size());
+
+
+        DropDatabaseParam dropDatabaseParam = DropDatabaseParam.newBuilder().withDatabaseName(dbName).build();
+        R<RpcStatus> dropResponse = client.dropDatabase(dropDatabaseParam);
+        Assertions.assertEquals(R.Status.Success.getCode(), dropResponse.getStatus().intValue());
+    }
 }