Database interfaces for V2 (#1013)

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

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

@@ -28,6 +28,9 @@ import io.milvus.orm.iterator.SearchIterator;
 
 import io.milvus.v2.exception.ErrorCode;
 import io.milvus.v2.exception.MilvusClientException;
+import io.milvus.v2.service.database.DatabaseService;
+import io.milvus.v2.service.database.request.*;
+import io.milvus.v2.service.database.response.*;
 import io.milvus.v2.service.collection.CollectionService;
 import io.milvus.v2.service.collection.request.*;
 import io.milvus.v2.service.collection.response.*;
@@ -62,6 +65,7 @@ public class MilvusClientV2 {
     @Setter
     private MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub;
     private final ClientUtils clientUtils = new ClientUtils();
+    private final DatabaseService databaseService = new DatabaseService();
     private final CollectionService collectionService = new CollectionService();
     private final IndexService indexService = new IndexService();
     private final VectorService vectorService = new VectorService();
@@ -225,18 +229,61 @@ public class MilvusClientV2 {
      * use Database
      * @param dbName databaseName
      */
-    public void useDatabase(@NonNull String dbName) {
+    public void useDatabase(@NonNull String dbName) throws InterruptedException {
         // check if database exists
         clientUtils.checkDatabaseExist(this.blockingStub, dbName);
         try {
             this.connectConfig.setDbName(dbName);
             this.close(3);
             this.connect(this.connectConfig);
-        }catch (InterruptedException e){
+        } catch (InterruptedException e){
             logger.error("close connect error");
+            throw new RuntimeException(e);
         }
     }
 
+    /**
+     * Creates a database in Milvus.
+     * @param request create database request
+     */
+    public void createDatabase(CreateDatabaseReq request) {
+        retry(()-> databaseService.createDatabase(this.blockingStub, request));
+    }
+
+    /**
+     * Drops a database. Note that this method drops all data in the database.
+     * @param request drop database request
+     */
+    public void dropDatabase(DropDatabaseReq request) {
+        retry(()-> databaseService.dropDatabase(this.blockingStub, request));
+    }
+
+    /**
+     * List all databases.
+     * @return List of String database names
+     */
+    public ListDatabasesResp listDatabases() {
+        return retry(()-> databaseService.listDatabases(this.blockingStub));
+    }
+
+    /**
+     * Alter database with key value pair. (Available from Milvus v2.4.4)
+     * @param request alter database request
+     */
+    public void alterDatabase(AlterDatabaseReq request) {
+        retry(()-> databaseService.alterDatabase(this.blockingStub, request));
+    }
+
+    /**
+     * Show detail of database base, such as replica number and resource groups. (Available from Milvus v2.4.4)
+     * @param request describe database request
+     *
+     * @return DescribeDatabaseResp
+     */
+    public DescribeDatabaseResp describeDatabase(DescribeDatabaseReq request) {
+        return retry(()-> databaseService.describeDatabase(this.blockingStub, request));
+    }
+
     //Collection Operations
     /**
      * Creates a collection in Milvus.

src/main/java/io/milvus/v2/service/database/DatabaseService.java

@@ -0,0 +1,102 @@
+/*
+ * 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.v2.service.database;
+
+import io.milvus.grpc.*;
+import io.milvus.param.ParamUtils;
+import io.milvus.v2.service.BaseService;
+import io.milvus.v2.service.collection.response.ListCollectionsResp;
+import io.milvus.v2.service.database.request.*;
+import io.milvus.v2.service.database.response.*;
+import org.apache.commons.collections4.CollectionUtils;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class DatabaseService extends BaseService {
+    public Void createDatabase(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, CreateDatabaseReq request) {
+        String title = String.format("CreateDatabaseRequest databaseName:%s", request.getDatabaseName());
+        CreateDatabaseRequest.Builder builder = CreateDatabaseRequest.newBuilder()
+                .setDbName(request.getDatabaseName());
+        List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(request.getProperties());
+        if (CollectionUtils.isNotEmpty(propertiesList)) {
+            propertiesList.forEach(builder::addProperties);
+        }
+
+        Status response = blockingStub.createDatabase(builder.build());
+        rpcUtils.handleResponse(title, response);
+        return null;
+    }
+
+    public Void dropDatabase(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DropDatabaseReq request) {
+        String title = String.format("DropDatabaseRequest databaseName:%s", request.getDatabaseName());
+        DropDatabaseRequest rpcRequest = DropDatabaseRequest.newBuilder()
+                .setDbName(request.getDatabaseName())
+                .build();
+
+        Status response = blockingStub.dropDatabase(rpcRequest);
+        rpcUtils.handleResponse(title, response);
+        return null;
+    }
+
+    public ListDatabasesResp listDatabases(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub) {
+        ListDatabasesResponse response = blockingStub.listDatabases(ListDatabasesRequest.newBuilder().build());
+        ListDatabasesResp listDatabasesResp = ListDatabasesResp.builder()
+                .databaseNames(response.getDbNamesList())
+                .build();
+
+        return listDatabasesResp;
+    }
+
+    public Void alterDatabase(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, AlterDatabaseReq request) {
+        String title = String.format("AlterDatabaseRequest databaseName:%s", request.getDatabaseName());
+        AlterDatabaseRequest.Builder builder = AlterDatabaseRequest.newBuilder()
+                .setDbName(request.getDatabaseName());
+        List<KeyValuePair> propertiesList = ParamUtils.AssembleKvPair(request.getProperties());
+        if (CollectionUtils.isNotEmpty(propertiesList)) {
+            propertiesList.forEach(builder::addProperties);
+        }
+
+        Status response = blockingStub.alterDatabase(builder.build());
+        rpcUtils.handleResponse(title, response);
+        return null;
+    }
+
+    public DescribeDatabaseResp describeDatabase(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, DescribeDatabaseReq request) {
+        String title = String.format("DescribeDatabaseRequest databaseName:%s", request.getDatabaseName());
+        DescribeDatabaseRequest rpcRequest = DescribeDatabaseRequest.newBuilder()
+                .setDbName(request.getDatabaseName())
+                .build();
+
+        DescribeDatabaseResponse response = blockingStub.describeDatabase(rpcRequest);
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        Map<String, String> properties = new HashMap<>();
+        response.getPropertiesList().forEach((prop) -> properties.put(prop.getKey(), prop.getValue()));
+
+        DescribeDatabaseResp describeDatabaseResp = DescribeDatabaseResp.builder()
+                .databaseName(response.getDbName())
+                .properties(properties)
+                .build();
+
+        return describeDatabaseResp;
+    }
+}

src/main/java/io/milvus/v2/service/database/request/AlterDatabaseReq.java

@@ -0,0 +1,35 @@
+/*
+ * 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.v2.service.database.request;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.HashMap;
+import java.util.Map;
+
+@Data
+@SuperBuilder
+public class AlterDatabaseReq {
+    private String databaseName;
+    @Builder.Default
+    private Map<String, String> properties = new HashMap<>();
+}

src/main/java/io/milvus/v2/service/database/request/CreateDatabaseReq.java

@@ -0,0 +1,35 @@
+/*
+ * 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.v2.service.database.request;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.HashMap;
+import java.util.Map;
+
+@Data
+@SuperBuilder
+public class CreateDatabaseReq {
+    private String databaseName;
+    @Builder.Default
+    private Map<String, String> properties = new HashMap<>();
+}

src/main/java/io/milvus/v2/service/database/request/DescribeDatabaseReq.java

@@ -0,0 +1,29 @@
+/*
+ * 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.v2.service.database.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class DescribeDatabaseReq {
+    private String databaseName;
+}

src/main/java/io/milvus/v2/service/database/request/DropDatabaseReq.java

@@ -0,0 +1,29 @@
+/*
+ * 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.v2.service.database.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class DropDatabaseReq {
+    private String databaseName;
+}

src/main/java/io/milvus/v2/service/database/response/DescribeDatabaseResp.java

@@ -0,0 +1,35 @@
+/*
+ * 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.v2.service.database.response;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.HashMap;
+import java.util.Map;
+
+@Data
+@SuperBuilder
+public class DescribeDatabaseResp {
+    private String databaseName;
+    @Builder.Default
+    private Map<String, String> properties = new HashMap<>();
+}

src/main/java/io/milvus/v2/service/database/response/ListDatabasesResp.java

@@ -0,0 +1,31 @@
+/*
+ * 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.v2.service.database.response;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class ListDatabasesResp {
+    private List<String> databaseNames;
+}

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

@@ -36,6 +36,12 @@ import io.milvus.v2.common.IndexParam;
 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.CreateDatabaseReq;
+import io.milvus.v2.service.database.request.DescribeDatabaseReq;
+import io.milvus.v2.service.database.request.DropDatabaseReq;
+import io.milvus.v2.service.database.response.DescribeDatabaseResp;
+import io.milvus.v2.service.database.response.ListDatabasesResp;
 import io.milvus.v2.service.index.request.AlterIndexReq;
 import io.milvus.v2.service.index.request.CreateIndexReq;
 import io.milvus.v2.service.index.request.DescribeIndexReq;
@@ -1502,4 +1508,83 @@ class MilvusClientV2DockerTest {
 
         client.dropCollection(DropCollectionReq.builder().collectionName(randomCollectionName).build());
     }
+
+    @Test
+    void testDatabase() {
+        // get current database
+        ListDatabasesResp listDatabasesResp = client.listDatabases();
+        List<String> dbNames = listDatabasesResp.getDatabaseNames();
+        Assertions.assertEquals(1, dbNames.size());
+        String currentDbName = dbNames.get(0);
+
+        // create a new database
+        String tempDatabaseName = "db_temp";
+        Map<String, String> properties = new HashMap<>();
+        properties.put(Constant.DATABASE_REPLICA_NUMBER, "5");
+        CreateDatabaseReq createDatabaseReq = CreateDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .properties(properties)
+                .build();
+        client.createDatabase(createDatabaseReq);
+
+        listDatabasesResp = client.listDatabases();
+        dbNames = listDatabasesResp.getDatabaseNames();
+        Assertions.assertTrue(dbNames.contains(tempDatabaseName));
+
+        DescribeDatabaseResp descDBResp = client.describeDatabase(DescribeDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .build());
+        Map<String,String> propertiesResp = descDBResp.getProperties();
+        Assertions.assertTrue(propertiesResp.containsKey(Constant.DATABASE_REPLICA_NUMBER));
+        Assertions.assertEquals("5", propertiesResp.get(Constant.DATABASE_REPLICA_NUMBER));
+
+        // switch to the new database
+        Assertions.assertDoesNotThrow(()->client.useDatabase(tempDatabaseName));
+
+        // create a collection in the new database
+        String randomCollectionName = generator.generate(10);
+        String vectorFieldName = "float_vector";
+        CreateCollectionReq.CollectionSchema collectionSchema = baseSchema();
+        collectionSchema.addField(AddFieldReq.builder()
+                .fieldName(vectorFieldName)
+                .dataType(DataType.FloatVector)
+                .dimension(dimension)
+                .build());
+
+        IndexParam indexParam = IndexParam.builder()
+                .fieldName(vectorFieldName)
+                .indexType(IndexParam.IndexType.FLAT)
+                .metricType(IndexParam.MetricType.COSINE)
+                .build();
+
+        CreateCollectionReq requestCreate = CreateCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .collectionSchema(collectionSchema)
+                .indexParams(Collections.singletonList(indexParam))
+                .build();
+        client.createCollection(requestCreate);
+
+        ListCollectionsResp listCollectionsResp = client.listCollections();
+        List<String> collectionNames = listCollectionsResp.getCollectionNames();
+        Assertions.assertEquals(1, collectionNames.size());
+        Assertions.assertTrue(collectionNames.contains(randomCollectionName));
+
+        // drop the collection so that we can drop the database later
+        client.dropCollection(DropCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+
+        // switch to the old database
+        Assertions.assertDoesNotThrow(()->client.useDatabase(currentDbName));
+
+        // drop the new database
+        client.dropDatabase(DropDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .build());
+
+        // check the new database is deleted
+        listDatabasesResp = client.listDatabases();
+        dbNames = listDatabasesResp.getDatabaseNames();
+        Assertions.assertFalse(dbNames.contains(tempDatabaseName));
+    }
 }