[cherry-pick]Add cases about resource group and database

Signed-off-by: ThreadDao <[email protected]>

client/options.go

@@ -111,6 +111,13 @@ func WithReplicaNumber(rn int32) LoadCollectionOption {
 	}
 }
 
+// WithResourceGroups specifies some specific ResourceGroup(s) to load the replica(s), rather than using the default ResourceGroup.
+func WithResourceGroups(rgs []string) LoadCollectionOption {
+	return func(req *server.LoadCollectionRequest) {
+		req.ResourceGroups = rgs
+	}
+}
+
 // SearchQueryOption is an option of search/query request
 type SearchQueryOption struct {
 	// Consistency Level & Time travel

test/base/milvus_client.go

@@ -35,6 +35,7 @@ func NewMilvusClient(ctx context.Context, cfg client.Config) (*MilvusClient, err
 	}, err
 }
 
+// !!!Deprecated in future, use `NewMilvusClient` first.
 func NewDefaultMilvusClient(ctx context.Context, addr string) (*MilvusClient, error) {
 	preRequest("NewDefaultGrpcClient", ctx, addr)
 	mClient, err := client.NewDefaultGrpcClient(ctx, addr)
@@ -86,9 +87,9 @@ func (mc *MilvusClient) Close() error {
 // 3. goroutine A access DB2 after 2.
 func (mc *MilvusClient) UsingDatabase(ctx context.Context, dbName string) error {
 	preRequest("UsingDatabase", ctx, dbName)
-	mc.mClient.UsingDatabase(ctx, dbName)
-	postResponse("UsingDatabase", nil)
-	return nil
+	err := mc.mClient.UsingDatabase(ctx, dbName)
+	postResponse("UsingDatabase", err)
+	return err
 }
 
 // -- database --
@@ -103,15 +104,15 @@ func (mc *MilvusClient) ListDatabases(ctx context.Context) ([]entity.Database, e
 
 // CreateDatabase create database with the given name.
 func (mc *MilvusClient) CreateDatabase(ctx context.Context, dbName string) error {
-	preRequest("CreateDatabase", ctx)
+	preRequest("CreateDatabase", ctx, dbName)
 	err := mc.mClient.CreateDatabase(ctx, dbName)
 	postResponse("CreateDatabase", err)
 	return err
 }
 
 // DropDatabase drop database with the given db name.
 func (mc *MilvusClient) DropDatabase(ctx context.Context, dbName string) error {
-	preRequest("DropDatabase", ctx)
+	preRequest("DropDatabase", ctx, dbName)
 	err := mc.mClient.DropDatabase(ctx, dbName)
 	postResponse("DropDatabase", err)
 	return err
@@ -408,11 +409,11 @@ func (mc *MilvusClient) CreateCollectionByRow(ctx context.Context, row entity.Ro
 	return err
 }
 
-// InsertByRows insert by rows
-func (mc *MilvusClient) InsertByRows(ctx context.Context, collName string, paritionName string, rows []entity.Row) (entity.Column, error) {
-	preRequest("InsertByRows", ctx, collName, paritionName, rows)
-	column, err := mc.mClient.InsertByRows(ctx, collName, paritionName, rows)
-	postResponse("InsertByRows", err, column)
+// InsertRows InsertByRows insert by rows
+func (mc *MilvusClient) InsertRows(ctx context.Context, collName string, partitionName string, rows []interface{}) (entity.Column, error) {
+	preRequest("InsertRows", ctx, collName, partitionName, len(rows))
+	column, err := mc.mClient.InsertRows(ctx, collName, partitionName, rows)
+	postResponse("InsertRows", err, column)
 	return column, err
 }
 
@@ -477,3 +478,51 @@ func (mc *MilvusClient) WaitForCompactionCompleted(ctx context.Context, compacti
 		}
 	}
 }
+
+// List Resource Groups
+func (mc *MilvusClient) ListResourceGroups(ctx context.Context) ([]string, error) {
+	preRequest("ListResourceGroups", ctx)
+	rgs, err := mc.mClient.ListResourceGroups(ctx)
+	postResponse("ListResourceGroups", err, rgs)
+	return rgs, err
+}
+
+// CreateResourceGroup
+func (mc *MilvusClient) CreateResourceGroup(ctx context.Context, rgName string) error {
+	preRequest("CreateResourceGroup", ctx, rgName)
+	err := mc.mClient.CreateResourceGroup(ctx, rgName)
+	postResponse("CreateResourceGroup", err)
+	return err
+}
+
+// DescribeResourceGroup
+func (mc *MilvusClient) DescribeResourceGroup(ctx context.Context, rgName string) (*entity.ResourceGroup, error) {
+	preRequest("DescribeResourceGroup", ctx, rgName)
+	rg, err := mc.mClient.DescribeResourceGroup(ctx, rgName)
+	postResponse("DescribeResourceGroup", err, rg)
+	return rg, err
+}
+
+// DropResourceGroup
+func (mc *MilvusClient) DropResourceGroup(ctx context.Context, rgName string) error {
+	preRequest("DropResourceGroup", ctx, rgName)
+	err := mc.mClient.DropResourceGroup(ctx, rgName)
+	postResponse("DropResourceGroup", err)
+	return err
+}
+
+// TransferNode
+func (mc *MilvusClient) TransferNode(ctx context.Context, sourceRg, targetRg string, nodesNum int32) error {
+	preRequest("TransferNode", ctx, sourceRg, targetRg, nodesNum)
+	err := mc.mClient.TransferNode(ctx, sourceRg, targetRg, nodesNum)
+	postResponse("TransferNode", err)
+	return err
+}
+
+// TransferReplica
+func (mc *MilvusClient) TransferReplica(ctx context.Context, sourceRg, targetRg string, collectionName string, replicaNum int64) error {
+	preRequest("TransferReplica", ctx, sourceRg, targetRg, collectionName, replicaNum)
+	err := mc.mClient.TransferReplica(ctx, sourceRg, targetRg, collectionName, replicaNum)
+	postResponse("TransferReplica", err)
+	return err
+}

test/common/response_check.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"fmt"
 	"log"
 	"testing"
 
@@ -72,7 +73,16 @@ func CheckContainsCollection(t *testing.T, collections []*entity.Collection, col
 	for _, collection := range collections {
 		allCollNames = append(allCollNames, collection.Name)
 	}
-	require.Contains(t, allCollNames, collName)
+	require.Containsf(t, allCollNames, collName, fmt.Sprintf("The collection %s not in: %v", collName, allCollNames))
+}
+
+// check collections not contains collName
+func CheckNotContainsCollection(t *testing.T, collections []*entity.Collection, collName string) {
+	allCollNames := make([]string, 0, len(collections))
+	for _, collection := range collections {
+		allCollNames = append(allCollNames, collection.Name)
+	}
+	require.NotContainsf(t, allCollNames, collName, fmt.Sprintf("The collection %s should not be in: %v", collName, allCollNames))
 }
 
 // check insert result, ids len (insert count), ids data (pks, but no auto ids)
@@ -180,3 +190,27 @@ func CheckPersistentSegments(t *testing.T, actualSegments []*entity.Segment, exp
 	}
 	require.Equal(t, actualNb, expNb)
 }
+
+func CheckResourceGroup(t *testing.T, actualRg *entity.ResourceGroup, expRg *entity.ResourceGroup) {
+	require.EqualValues(t, expRg, actualRg)
+}
+
+func getDbNames(dbs []entity.Database) []string {
+	allDbNames := make([]string, 0, len(dbs))
+	for _, db := range dbs {
+		allDbNames = append(allDbNames, db.Name)
+	}
+	return allDbNames
+}
+
+// check collections contains collName
+func CheckContainsDb(t *testing.T, dbs []entity.Database, dbName string) {
+	allDbNames := getDbNames(dbs)
+	require.Containsf(t, allDbNames, dbName, fmt.Sprintf("%s db not in dbs: %v", dbName, dbs))
+}
+
+// check collections contains collName
+func CheckNotContainsDb(t *testing.T, dbs []entity.Database, dbName string) {
+	allDbNames := getDbNames(dbs)
+	require.NotContainsf(t, allDbNames, dbName, fmt.Sprintf("%s db should not be in dbs: %v", dbName, dbs))
+}

test/common/utils.go

@@ -20,6 +20,7 @@ const (
 	DefaultPartition          = "_default"
 	DefaultIndexName          = "_default_idx_102"
 	DefaultIndexNameBinary    = "_default_idx_100"
+	DefaultDb                 = "default"
 	DefaultDim                = 128
 	DefaultDimStr             = "128"
 	MaxDim                    = 32768
@@ -32,6 +33,8 @@ const (
 	MaxCollectionNameLen      = 255
 	RowCount                  = "row_count"
 	RetentionDuration         = 40 // common.retentionDuration
+	DefaultRgName             = "__default_resource_group"
+	DefaultRgCapacity         = 1000000
 )
 
 var IndexStateValue = map[string]int32{

test/testcases/connect_test.go

@@ -6,12 +6,36 @@ import (
 	"testing"
 	"time"
 
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+
 	"github.com/milvus-io/milvus-sdk-go/v2/test/base"
 	"github.com/milvus-io/milvus-sdk-go/v2/test/common"
 )
 
 // test connect and close, connect again
 func TestConnectClose(t *testing.T) {
+	// connect
+	ctx := createContext(t, time.Second*common.DefaultTimeout)
+	mc, errConnect := base.NewMilvusClient(ctx, client.Config{Address: *addr})
+	common.CheckErr(t, errConnect, true)
+
+	// connect success
+	_, errList := mc.ListCollections(ctx)
+	common.CheckErr(t, errList, true)
+
+	// close success
+	mc.Close()
+	_, errList2 := mc.ListCollections(ctx)
+	common.CheckErr(t, errList2, false, "the client connection is closing")
+
+	// connect again
+	mc, errConnect2 := base.NewMilvusClient(ctx, client.Config{Address: *addr})
+	common.CheckErr(t, errConnect2, true)
+	_, errList3 := mc.ListCollections(ctx)
+	common.CheckErr(t, errList3, true)
+}
+
+func TestConnectCloseDefault(t *testing.T) {
 	// connect
 	ctx := createContext(t, time.Second*common.DefaultTimeout)
 	mc, errConnect := base.NewDefaultMilvusClient(ctx, *addr)
@@ -39,7 +63,7 @@ func TestConnectInvalidAddr(t *testing.T) {
 	// connect
 	ctx := createContext(t, time.Second*10)
 
-	_, errConnect := base.NewDefaultMilvusClient(ctx, "aa")
+	_, errConnect := base.NewMilvusClient(ctx, client.Config{Address: "aa"})
 	common.CheckErr(t, errConnect, false, "xxx")
 }
 
@@ -48,11 +72,11 @@ func TestConnectRepeat(t *testing.T) {
 	// connect
 	ctx := createContext(t, time.Second*10)
 
-	_, errConnect := base.NewDefaultMilvusClient(ctx, *addr)
+	_, errConnect := base.NewMilvusClient(ctx, client.Config{Address: *addr})
 	common.CheckErr(t, errConnect, true)
 
 	// connect again
-	mc, errConnect2 := base.NewDefaultMilvusClient(ctx, *addr)
+	mc, errConnect2 := base.NewMilvusClient(ctx, client.Config{Address: *addr})
 	common.CheckErr(t, errConnect2, true)
 
 	_, err := mc.ListCollections(ctx)
@@ -63,7 +87,7 @@ func TestConnectRepeat(t *testing.T) {
 func TestCloseRepeat(t *testing.T) {
 	// connect
 	ctx := createContext(t, time.Second*10)
-	mc, errConnect2 := base.NewDefaultMilvusClient(ctx, *addr)
+	mc, errConnect2 := base.NewMilvusClient(ctx, client.Config{Address: *addr})
 	common.CheckErr(t, errConnect2, true)
 
 	// close and again