Database interfaces for V2 (#1015)

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.*;
@@ -64,6 +67,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();
@@ -227,18 +231,43 @@ 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));
+    }
+
     //Collection Operations
     /**
      * Creates a collection in Milvus.

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

@@ -0,0 +1,62 @@
+/*
+ * 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());
+        Status response = blockingStub.createDatabase(CreateDatabaseRequest.newBuilder()
+                .setDbName(request.getDatabaseName()).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;
+    }
+}

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

@@ -0,0 +1,30 @@
+/*
+ * 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 CreateDatabaseReq {
+    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/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

@@ -29,19 +29,16 @@ import io.milvus.param.Constant;
 import io.milvus.response.QueryResultsWrapper;
 import io.milvus.v2.common.ConsistencyLevel;
 import io.milvus.v2.common.DataType;
-import io.milvus.v2.common.IndexBuildState;
 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.index.request.CreateIndexReq;
-import io.milvus.v2.service.index.request.DescribeIndexReq;
-import io.milvus.v2.service.index.request.DropIndexReq;
-import io.milvus.v2.service.index.response.DescribeIndexResp;
+import io.milvus.v2.service.collection.response.*;
+import io.milvus.v2.service.database.request.*;
+import io.milvus.v2.service.database.response.*;
+import io.milvus.v2.service.index.request.*;
+import io.milvus.v2.service.index.response.*;
 import io.milvus.v2.service.partition.request.*;
-import io.milvus.v2.service.utility.request.AlterAliasReq;
-import io.milvus.v2.service.utility.request.CreateAliasReq;
-import io.milvus.v2.service.utility.request.DropAliasReq;
+import io.milvus.v2.service.utility.request.*;
 import io.milvus.v2.service.vector.request.*;
 import io.milvus.v2.service.vector.request.data.*;
 import io.milvus.v2.service.vector.response.*;
@@ -57,7 +54,6 @@ import org.testcontainers.milvus.MilvusContainer;
 
 import java.nio.ByteBuffer;
 import java.util.*;
-import java.util.concurrent.TimeUnit;
 
 @Testcontainers(disabledWithoutDocker = true)
 class MilvusClientV2DockerTest {
@@ -1072,4 +1068,73 @@ 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";
+        CreateDatabaseReq createDatabaseReq = CreateDatabaseReq.builder()
+                .databaseName(tempDatabaseName)
+                .build();
+        client.createDatabase(createDatabaseReq);
+
+        listDatabasesResp = client.listDatabases();
+        dbNames = listDatabasesResp.getDatabaseNames();
+        Assertions.assertTrue(dbNames.contains(tempDatabaseName));
+
+        // 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));
+    }
 }