mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-02 20:09:57 +08:00
03b5d32569
Signed-off-by: XuanYang-cn <xuan.yang@zilliz.com>
1136 lines
40 KiB
Go
1136 lines
40 KiB
Go
package master
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"math/rand"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/util/tsoutil"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/stretchr/testify/assert"
|
|
ms "github.com/zilliztech/milvus-distributed/internal/msgstream"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/masterpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/servicepb"
|
|
"go.etcd.io/etcd/clientv3"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
var testPORT = 53200
|
|
|
|
func genMasterTestPort() int64 {
|
|
testPORT++
|
|
return int64(testPORT)
|
|
}
|
|
|
|
func refreshMasterAddress() {
|
|
masterPort := genMasterTestPort()
|
|
Params.Port = int(masterPort)
|
|
masterAddr := makeMasterAddress(masterPort)
|
|
Params.Address = masterAddr
|
|
}
|
|
|
|
func makeMasterAddress(port int64) string {
|
|
masterAddr := "127.0.0.1:" + strconv.FormatInt(port, 10)
|
|
return masterAddr
|
|
}
|
|
|
|
func makeNewChannalNames(names []string, suffix string) []string {
|
|
var ret []string
|
|
for _, name := range names {
|
|
ret = append(ret, name+suffix)
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func refreshChannelNames() {
|
|
suffix := "_test" + strconv.FormatInt(rand.Int63n(100), 10)
|
|
Params.DDChannelNames = makeNewChannalNames(Params.DDChannelNames, suffix)
|
|
Params.WriteNodeTimeTickChannelNames = makeNewChannalNames(Params.WriteNodeTimeTickChannelNames, suffix)
|
|
Params.InsertChannelNames = makeNewChannalNames(Params.InsertChannelNames, suffix)
|
|
Params.K2SChannelNames = makeNewChannalNames(Params.K2SChannelNames, suffix)
|
|
Params.ProxyTimeTickChannelNames = makeNewChannalNames(Params.ProxyTimeTickChannelNames, suffix)
|
|
Params.QueryNodeStatsChannelName = Params.QueryNodeStatsChannelName + suffix
|
|
Params.MetaRootPath = "/test" + strconv.FormatInt(rand.Int63n(100), 10) + "/root/kv"
|
|
}
|
|
|
|
func receiveTimeTickMsg(stream *ms.MsgStream) bool {
|
|
for {
|
|
result := (*stream).Consume()
|
|
if len(result.Msgs) > 0 {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
func getTimeTickMsgPack(ttmsgs [][2]uint64) *ms.MsgPack {
|
|
msgPack := ms.MsgPack{}
|
|
for _, vi := range ttmsgs {
|
|
msgPack.Msgs = append(msgPack.Msgs, getTtMsg(internalpb.MsgType_kTimeTick, UniqueID(vi[0]), Timestamp(vi[1])))
|
|
}
|
|
return &msgPack
|
|
}
|
|
|
|
func TestMaster(t *testing.T) {
|
|
Init()
|
|
refreshMasterAddress()
|
|
refreshChannelNames()
|
|
etcdAddr := Params.EtcdAddress
|
|
etcdCli, err := clientv3.New(clientv3.Config{Endpoints: []string{etcdAddr}})
|
|
assert.Nil(t, err)
|
|
_, err = etcdCli.Delete(context.Background(), Params.MetaRootPath, clientv3.WithPrefix())
|
|
assert.Nil(t, err)
|
|
|
|
gTestTsoAllocator = NewGlobalTSOAllocator("timestamp", tsoutil.NewTSOKVBase([]string{etcdAddr}, Params.MetaRootPath, "tso"))
|
|
gTestIDAllocator = NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase([]string{etcdAddr}, Params.MetaRootPath, "gid"))
|
|
pulsarAddr := Params.PulsarAddress
|
|
Params.ProxyIDList = []UniqueID{0}
|
|
//Param
|
|
// Creates server.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
svr, err := CreateServer(ctx)
|
|
|
|
if err != nil {
|
|
log.Print("create server failed", zap.Error(err))
|
|
}
|
|
|
|
if err := svr.Run(int64(Params.Port)); err != nil {
|
|
log.Fatal("run server failed", zap.Error(err))
|
|
}
|
|
|
|
conn, err := grpc.DialContext(ctx, Params.Address, grpc.WithInsecure(), grpc.WithBlock())
|
|
require.Nil(t, err)
|
|
cli := masterpb.NewMasterClient(conn)
|
|
|
|
t.Run("TestConfigTask", func(t *testing.T) {
|
|
testKeys := []string{
|
|
"/etcd/address",
|
|
"/master/port",
|
|
"/master/proxyidlist",
|
|
"/master/segmentthresholdfactor",
|
|
"/pulsar/token",
|
|
"/reader/stopflag",
|
|
"/proxy/timezone",
|
|
"/proxy/network/address",
|
|
"/proxy/storage/path",
|
|
"/storage/accesskey",
|
|
}
|
|
|
|
testVals := []string{
|
|
"localhost",
|
|
"53100",
|
|
"[1 2]",
|
|
"0.75",
|
|
"eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJKb2UifQ.ipevRNuRP6HflG8cFKnmUPtypruRC4fb1DWtoLL62SY",
|
|
"-1",
|
|
"UTC+8",
|
|
"0.0.0.0",
|
|
"/var/lib/milvus",
|
|
"",
|
|
}
|
|
|
|
sc := SysConfig{kv: svr.kvBase}
|
|
sc.InitFromFile(".")
|
|
|
|
configRequest := &internalpb.SysConfigRequest{
|
|
MsgType: internalpb.MsgType_kGetSysConfigs,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Keys: testKeys,
|
|
KeyPrefixes: []string{},
|
|
}
|
|
|
|
response, err := cli.GetSysConfigs(ctx, configRequest)
|
|
assert.Nil(t, err)
|
|
assert.ElementsMatch(t, testKeys, response.Keys)
|
|
assert.ElementsMatch(t, testVals, response.Values)
|
|
assert.Equal(t, len(response.GetKeys()), len(response.GetValues()))
|
|
|
|
configRequest = &internalpb.SysConfigRequest{
|
|
MsgType: internalpb.MsgType_kGetSysConfigs,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Keys: []string{},
|
|
KeyPrefixes: []string{"/master"},
|
|
}
|
|
|
|
response, err = cli.GetSysConfigs(ctx, configRequest)
|
|
assert.Nil(t, err)
|
|
for i := range response.GetKeys() {
|
|
assert.True(t, strings.HasPrefix(response.GetKeys()[i], "/master"))
|
|
}
|
|
assert.Equal(t, len(response.GetKeys()), len(response.GetValues()))
|
|
|
|
})
|
|
|
|
t.Run("TestConfigDuplicateKeysAndKeyPrefix", func(t *testing.T) {
|
|
configRequest := &internalpb.SysConfigRequest{}
|
|
configRequest.Keys = []string{}
|
|
configRequest.KeyPrefixes = []string{"/master"}
|
|
|
|
configRequest.Timestamp = uint64(time.Now().Unix())
|
|
resp, err := cli.GetSysConfigs(ctx, configRequest)
|
|
require.Nil(t, err)
|
|
assert.Equal(t, len(resp.GetKeys()), len(resp.GetValues()))
|
|
assert.NotEqual(t, 0, len(resp.GetKeys()))
|
|
|
|
configRequest.Keys = []string{"/master/port"}
|
|
configRequest.KeyPrefixes = []string{"/master"}
|
|
|
|
configRequest.Timestamp = uint64(time.Now().Unix())
|
|
respDup, err := cli.GetSysConfigs(ctx, configRequest)
|
|
require.Nil(t, err)
|
|
assert.Equal(t, len(respDup.GetKeys()), len(respDup.GetValues()))
|
|
assert.NotEqual(t, 0, len(respDup.GetKeys()))
|
|
assert.Equal(t, len(respDup.GetKeys()), len(resp.GetKeys()))
|
|
})
|
|
|
|
t.Run("TestCollectionTask", func(t *testing.T) {
|
|
sch := schemapb.CollectionSchema{
|
|
Name: "col1",
|
|
Description: "test collection",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
Name: "col1_f1",
|
|
Description: "test collection filed 1",
|
|
DataType: schemapb.DataType_VECTOR_FLOAT,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f1_tk1",
|
|
Value: "col1_f1_tv1",
|
|
},
|
|
{
|
|
Key: "col1_f1_tk2",
|
|
Value: "col1_f1_tv2",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f1_ik1",
|
|
Value: "col1_f1_iv1",
|
|
},
|
|
{
|
|
Key: "col1_f1_ik2",
|
|
Value: "col1_f1_iv2",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: "col1_f2",
|
|
Description: "test collection filed 2",
|
|
DataType: schemapb.DataType_VECTOR_BINARY,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f2_tk1",
|
|
Value: "col1_f2_tv1",
|
|
},
|
|
{
|
|
Key: "col1_f2_tk2",
|
|
Value: "col1_f2_tv2",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f2_ik1",
|
|
Value: "col1_f2_iv1",
|
|
},
|
|
{
|
|
Key: "col1_f2_ik2",
|
|
Value: "col1_f2_iv2",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
schemaBytes, err := proto.Marshal(&sch)
|
|
assert.Nil(t, err)
|
|
|
|
createCollectionReq := internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreateCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
}
|
|
log.Printf("... [Create] collection col1\n")
|
|
st, err := cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// HasCollection
|
|
reqHasCollection := internalpb.HasCollectionRequest{
|
|
MsgType: internalpb.MsgType_kHasCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{
|
|
CollectionName: "col1",
|
|
},
|
|
}
|
|
|
|
// "col1" is true
|
|
log.Printf("... [Has] collection col1\n")
|
|
boolResp, err := cli.HasCollection(ctx, &reqHasCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, true, boolResp.Value)
|
|
assert.Equal(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// "colNotExist" is false
|
|
reqHasCollection.CollectionName.CollectionName = "colNotExist"
|
|
boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, boolResp.Value, false)
|
|
assert.Equal(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// error
|
|
reqHasCollection.Timestamp = Timestamp(0)
|
|
reqHasCollection.CollectionName.CollectionName = "col1"
|
|
boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// ShowCollection
|
|
reqShowCollection := internalpb.ShowCollectionRequest{
|
|
MsgType: internalpb.MsgType_kShowCollections,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
}
|
|
|
|
listResp, err := cli.ShowCollections(ctx, &reqShowCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, listResp.Status.ErrorCode)
|
|
assert.Equal(t, 1, len(listResp.Values))
|
|
assert.Equal(t, "col1", listResp.Values[0])
|
|
|
|
reqShowCollection.Timestamp = Timestamp(0)
|
|
listResp, err = cli.ShowCollections(ctx, &reqShowCollection)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, commonpb.ErrorCode_SUCCESS, listResp.Status.ErrorCode)
|
|
|
|
// CreateCollection Test
|
|
collMeta, err := svr.metaTable.GetCollectionByName(sch.Name)
|
|
assert.Nil(t, err)
|
|
t.Logf("collection id = %d", collMeta.ID)
|
|
assert.Equal(t, collMeta.Schema.Name, "col1")
|
|
assert.Equal(t, collMeta.Schema.AutoID, false)
|
|
assert.Equal(t, len(collMeta.Schema.Fields), 4)
|
|
assert.Equal(t, collMeta.Schema.Fields[0].Name, "col1_f1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].Name, "col1_f2")
|
|
assert.Equal(t, collMeta.Schema.Fields[2].Name, "RowID")
|
|
assert.Equal(t, collMeta.Schema.Fields[3].Name, "Timestamp")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].DataType, schemapb.DataType_VECTOR_FLOAT)
|
|
assert.Equal(t, collMeta.Schema.Fields[1].DataType, schemapb.DataType_VECTOR_BINARY)
|
|
assert.Equal(t, collMeta.Schema.Fields[2].DataType, schemapb.DataType_INT64)
|
|
assert.Equal(t, collMeta.Schema.Fields[3].DataType, schemapb.DataType_INT64)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[0].TypeParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[0].IndexParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[1].TypeParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[1].IndexParams), 2)
|
|
assert.Equal(t, int64(100), collMeta.Schema.Fields[0].FieldID)
|
|
assert.Equal(t, int64(101), collMeta.Schema.Fields[1].FieldID)
|
|
assert.Equal(t, int64(0), collMeta.Schema.Fields[2].FieldID)
|
|
assert.Equal(t, int64(1), collMeta.Schema.Fields[3].FieldID)
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[0].Key, "col1_f1_tk1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[1].Key, "col1_f1_tk2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[0].Value, "col1_f1_tv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[1].Value, "col1_f1_tv2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[0].Key, "col1_f1_ik1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[1].Key, "col1_f1_ik2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[0].Value, "col1_f1_iv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[1].Value, "col1_f1_iv2")
|
|
|
|
assert.Equal(t, int64(101), collMeta.Schema.Fields[1].FieldID)
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[0].Key, "col1_f2_tk1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[1].Key, "col1_f2_tk2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[0].Value, "col1_f2_tv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[1].Value, "col1_f2_tv2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[0].Key, "col1_f2_ik1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[1].Key, "col1_f2_ik2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[0].Value, "col1_f2_iv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[1].Value, "col1_f2_iv2")
|
|
|
|
createCollectionReq.Timestamp = Timestamp(0)
|
|
st, err = cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// DescribeCollection Test
|
|
reqDescribe := &internalpb.DescribeCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDescribeCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{
|
|
CollectionName: "col1",
|
|
},
|
|
}
|
|
des, err := cli.DescribeCollection(ctx, reqDescribe)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, des.Status.ErrorCode)
|
|
|
|
assert.Equal(t, "col1", des.Schema.Name)
|
|
assert.Equal(t, false, des.Schema.AutoID)
|
|
assert.Equal(t, 2, len(des.Schema.Fields))
|
|
assert.Equal(t, "col1_f1", des.Schema.Fields[0].Name)
|
|
assert.Equal(t, "col1_f2", des.Schema.Fields[1].Name)
|
|
assert.Equal(t, schemapb.DataType_VECTOR_FLOAT, des.Schema.Fields[0].DataType)
|
|
assert.Equal(t, schemapb.DataType_VECTOR_BINARY, des.Schema.Fields[1].DataType)
|
|
assert.Equal(t, 2, len(des.Schema.Fields[0].TypeParams))
|
|
assert.Equal(t, 2, len(des.Schema.Fields[0].IndexParams))
|
|
assert.Equal(t, 2, len(des.Schema.Fields[1].TypeParams))
|
|
assert.Equal(t, 2, len(des.Schema.Fields[1].IndexParams))
|
|
assert.Equal(t, int64(100), des.Schema.Fields[0].FieldID)
|
|
assert.Equal(t, "col1_f1_tk1", des.Schema.Fields[0].TypeParams[0].Key)
|
|
assert.Equal(t, "col1_f1_tv1", des.Schema.Fields[0].TypeParams[0].Value)
|
|
assert.Equal(t, "col1_f1_ik1", des.Schema.Fields[0].IndexParams[0].Key)
|
|
assert.Equal(t, "col1_f1_iv1", des.Schema.Fields[0].IndexParams[0].Value)
|
|
assert.Equal(t, "col1_f1_tk2", des.Schema.Fields[0].TypeParams[1].Key)
|
|
assert.Equal(t, "col1_f1_tv2", des.Schema.Fields[0].TypeParams[1].Value)
|
|
assert.Equal(t, "col1_f1_ik2", des.Schema.Fields[0].IndexParams[1].Key)
|
|
assert.Equal(t, "col1_f1_iv2", des.Schema.Fields[0].IndexParams[1].Value)
|
|
|
|
assert.Equal(t, int64(101), des.Schema.Fields[1].FieldID)
|
|
assert.Equal(t, "col1_f2_tk1", des.Schema.Fields[1].TypeParams[0].Key)
|
|
assert.Equal(t, "col1_f2_tv1", des.Schema.Fields[1].TypeParams[0].Value)
|
|
assert.Equal(t, "col1_f2_ik1", des.Schema.Fields[1].IndexParams[0].Key)
|
|
assert.Equal(t, "col1_f2_iv1", des.Schema.Fields[1].IndexParams[0].Value)
|
|
assert.Equal(t, "col1_f2_tk2", des.Schema.Fields[1].TypeParams[1].Key)
|
|
assert.Equal(t, "col1_f2_tv2", des.Schema.Fields[1].TypeParams[1].Value)
|
|
assert.Equal(t, "col1_f2_ik2", des.Schema.Fields[1].IndexParams[1].Key)
|
|
assert.Equal(t, "col1_f2_iv2", des.Schema.Fields[1].IndexParams[1].Value)
|
|
|
|
reqDescribe.CollectionName.CollectionName = "colNotExist"
|
|
des, err = cli.DescribeCollection(ctx, reqDescribe)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, commonpb.ErrorCode_SUCCESS, des.Status.ErrorCode)
|
|
log.Printf(des.Status.Reason)
|
|
|
|
reqDescribe.CollectionName.CollectionName = "col1"
|
|
reqDescribe.Timestamp = Timestamp(0)
|
|
des, err = cli.DescribeCollection(ctx, reqDescribe)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, commonpb.ErrorCode_SUCCESS, des.Status.ErrorCode)
|
|
log.Printf(des.Status.Reason)
|
|
|
|
// ------------------------------DropCollectionTask---------------------------
|
|
log.Printf("... [Drop] collection col1\n")
|
|
ser := servicepb.CollectionName{CollectionName: "col1"}
|
|
|
|
reqDrop := internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &ser,
|
|
}
|
|
|
|
// DropCollection
|
|
st, err = cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
collMeta, err = svr.metaTable.GetCollectionByName(sch.Name)
|
|
assert.NotNil(t, err)
|
|
|
|
// HasCollection "col1" is false
|
|
reqHasCollection.Timestamp = uint64(time.Now().Unix())
|
|
reqHasCollection.CollectionName.CollectionName = "col1"
|
|
boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, false, boolResp.Value)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, boolResp.Status.ErrorCode)
|
|
|
|
// ShowCollections
|
|
reqShowCollection.Timestamp = uint64(time.Now().Unix())
|
|
listResp, err = cli.ShowCollections(ctx, &reqShowCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, listResp.Status.ErrorCode)
|
|
assert.Equal(t, 0, len(listResp.Values))
|
|
|
|
// Drop again
|
|
st, err = cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.NotEqual(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
// Create "col1"
|
|
createCollectionReq.Timestamp = uint64(time.Now().Unix())
|
|
|
|
st, err = cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
reqHasCollection.Timestamp = uint64(time.Now().Unix())
|
|
boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, true, boolResp.Value)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, boolResp.Status.ErrorCode)
|
|
|
|
// Create "col2"
|
|
sch.Name = "col2"
|
|
schemaBytes, err = proto.Marshal(&sch)
|
|
assert.Nil(t, err)
|
|
|
|
createCollectionReq = internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreateCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
}
|
|
st, err = cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, st.ErrorCode)
|
|
|
|
// Show Collections
|
|
reqShowCollection.Timestamp = uint64(time.Now().Unix())
|
|
listResp, err = cli.ShowCollections(ctx, &reqShowCollection)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, listResp.Status.ErrorCode)
|
|
assert.Equal(t, 2, len(listResp.Values))
|
|
assert.ElementsMatch(t, []string{"col1", "col2"}, listResp.Values)
|
|
|
|
// Drop Collection
|
|
ser = servicepb.CollectionName{CollectionName: "col1"}
|
|
|
|
reqDrop = internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &ser,
|
|
}
|
|
|
|
// DropCollection
|
|
st, err = cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
ser = servicepb.CollectionName{CollectionName: "col2"}
|
|
reqDrop = internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &ser,
|
|
}
|
|
|
|
// DropCollection
|
|
st, err = cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
//consume msg
|
|
ddMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
ddMs.SetPulsarClient(pulsarAddr)
|
|
ddMs.CreatePulsarConsumers(Params.DDChannelNames, "DDStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
ddMs.Start()
|
|
|
|
var consumeMsg ms.MsgStream = ddMs
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
break
|
|
}
|
|
}
|
|
ddMs.Close()
|
|
})
|
|
|
|
t.Run("TestPartitionTask", func(t *testing.T) {
|
|
sch := schemapb.CollectionSchema{
|
|
Name: "col1",
|
|
Description: "test collection",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
Name: "col1_f1",
|
|
Description: "test collection filed 1",
|
|
DataType: schemapb.DataType_VECTOR_FLOAT,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f1_tk1",
|
|
Value: "col1_f1_tv1",
|
|
},
|
|
{
|
|
Key: "col1_f1_tk2",
|
|
Value: "col1_f1_tv2",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f1_ik1",
|
|
Value: "col1_f1_iv1",
|
|
},
|
|
{
|
|
Key: "col1_f1_ik2",
|
|
Value: "col1_f1_iv2",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: "col1_f2",
|
|
Description: "test collection filed 2",
|
|
DataType: schemapb.DataType_VECTOR_BINARY,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f2_tk1",
|
|
Value: "col1_f2_tv1",
|
|
},
|
|
{
|
|
Key: "col1_f2_tk2",
|
|
Value: "col1_f2_tv2",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "col1_f2_ik1",
|
|
Value: "col1_f2_iv1",
|
|
},
|
|
{
|
|
Key: "col1_f2_ik2",
|
|
Value: "col1_f2_iv2",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
schemaBytes, err := proto.Marshal(&sch)
|
|
assert.Nil(t, err)
|
|
|
|
createCollectionReq := internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
}
|
|
st, _ := cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.NotNil(t, st)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, st.ErrorCode)
|
|
|
|
createPartitionReq := internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
st, _ = cli.CreatePartition(ctx, &createPartitionReq)
|
|
assert.NotNil(t, st)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, st.ErrorCode)
|
|
|
|
createPartitionReq = internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
st, _ = cli.CreatePartition(ctx, &createPartitionReq)
|
|
assert.NotNil(t, st)
|
|
assert.Equal(t, commonpb.ErrorCode_UNEXPECTED_ERROR, st.ErrorCode)
|
|
|
|
createPartitionReq = internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition2"},
|
|
}
|
|
st, _ = cli.CreatePartition(ctx, &createPartitionReq)
|
|
assert.NotNil(t, st)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, st.ErrorCode)
|
|
|
|
collMeta, err := svr.metaTable.GetCollectionByName(sch.Name)
|
|
assert.Nil(t, err)
|
|
t.Logf("collection id = %d", collMeta.ID)
|
|
assert.Equal(t, collMeta.Schema.Name, "col1")
|
|
assert.Equal(t, collMeta.Schema.AutoID, false)
|
|
assert.Equal(t, len(collMeta.Schema.Fields), 4)
|
|
assert.Equal(t, collMeta.Schema.Fields[0].Name, "col1_f1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].Name, "col1_f2")
|
|
assert.Equal(t, collMeta.Schema.Fields[2].Name, "RowID")
|
|
assert.Equal(t, collMeta.Schema.Fields[3].Name, "Timestamp")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].DataType, schemapb.DataType_VECTOR_FLOAT)
|
|
assert.Equal(t, collMeta.Schema.Fields[1].DataType, schemapb.DataType_VECTOR_BINARY)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[0].TypeParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[0].IndexParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[1].TypeParams), 2)
|
|
assert.Equal(t, len(collMeta.Schema.Fields[1].IndexParams), 2)
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[0].Key, "col1_f1_tk1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[1].Key, "col1_f1_tk2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[0].Value, "col1_f1_tv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].TypeParams[1].Value, "col1_f1_tv2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[0].Key, "col1_f1_ik1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[1].Key, "col1_f1_ik2")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[0].Value, "col1_f1_iv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[0].IndexParams[1].Value, "col1_f1_iv2")
|
|
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[0].Key, "col1_f2_tk1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[1].Key, "col1_f2_tk2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[0].Value, "col1_f2_tv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].TypeParams[1].Value, "col1_f2_tv2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[0].Key, "col1_f2_ik1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[1].Key, "col1_f2_ik2")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[0].Value, "col1_f2_iv1")
|
|
assert.Equal(t, collMeta.Schema.Fields[1].IndexParams[1].Value, "col1_f2_iv2")
|
|
assert.ElementsMatch(t, []string{"_default", "partition1", "partition2"}, collMeta.PartitionTags)
|
|
|
|
showPartitionReq := internalpb.ShowPartitionRequest{
|
|
MsgType: internalpb.MsgType_kShowPartitions,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{CollectionName: "col1"},
|
|
}
|
|
|
|
stringList, err := cli.ShowPartitions(ctx, &showPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.ElementsMatch(t, []string{"_default", "partition1", "partition2"}, stringList.Values)
|
|
|
|
showPartitionReq = internalpb.ShowPartitionRequest{
|
|
MsgType: internalpb.MsgType_kShowPartitions,
|
|
ReqID: 1,
|
|
Timestamp: 0,
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{CollectionName: "col1"},
|
|
}
|
|
|
|
stringList, _ = cli.ShowPartitions(ctx, &showPartitionReq)
|
|
assert.NotNil(t, stringList)
|
|
assert.Equal(t, commonpb.ErrorCode_UNEXPECTED_ERROR, stringList.Status.ErrorCode)
|
|
|
|
hasPartitionReq := internalpb.HasPartitionRequest{
|
|
MsgType: internalpb.MsgType_kHasPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
|
|
hasPartition, err := cli.HasPartition(ctx, &hasPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.True(t, hasPartition.Value)
|
|
|
|
hasPartitionReq = internalpb.HasPartitionRequest{
|
|
MsgType: internalpb.MsgType_kHasPartition,
|
|
ReqID: 1,
|
|
Timestamp: 0,
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
|
|
hasPartition, _ = cli.HasPartition(ctx, &hasPartitionReq)
|
|
assert.NotNil(t, hasPartition)
|
|
assert.Equal(t, commonpb.ErrorCode_UNEXPECTED_ERROR, stringList.Status.ErrorCode)
|
|
|
|
hasPartitionReq = internalpb.HasPartitionRequest{
|
|
MsgType: internalpb.MsgType_kHasPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition3"},
|
|
}
|
|
|
|
hasPartition, err = cli.HasPartition(ctx, &hasPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.False(t, hasPartition.Value)
|
|
|
|
deletePartitionReq := internalpb.DropPartitionRequest{
|
|
MsgType: internalpb.MsgType_kDropPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition2"},
|
|
}
|
|
|
|
st, err = cli.DropPartition(ctx, &deletePartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_SUCCESS, st.ErrorCode)
|
|
|
|
deletePartitionReq = internalpb.DropPartitionRequest{
|
|
MsgType: internalpb.MsgType_kDropPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition2"},
|
|
}
|
|
|
|
st, _ = cli.DropPartition(ctx, &deletePartitionReq)
|
|
assert.NotNil(t, st)
|
|
assert.Equal(t, commonpb.ErrorCode_UNEXPECTED_ERROR, st.ErrorCode)
|
|
|
|
hasPartitionReq = internalpb.HasPartitionRequest{
|
|
MsgType: internalpb.MsgType_kHasPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition2"},
|
|
}
|
|
|
|
hasPartition, err = cli.HasPartition(ctx, &hasPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.False(t, hasPartition.Value)
|
|
|
|
describePartitionReq := internalpb.DescribePartitionRequest{
|
|
MsgType: internalpb.MsgType_kDescribePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
|
|
describePartition, err := cli.DescribePartition(ctx, &describePartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"}, describePartition.Name)
|
|
|
|
describePartitionReq = internalpb.DescribePartitionRequest{
|
|
MsgType: internalpb.MsgType_kDescribePartition,
|
|
ReqID: 1,
|
|
Timestamp: 0,
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{CollectionName: "col1", Tag: "partition1"},
|
|
}
|
|
|
|
describePartition, _ = cli.DescribePartition(ctx, &describePartitionReq)
|
|
assert.Equal(t, commonpb.ErrorCode_UNEXPECTED_ERROR, describePartition.Status.ErrorCode)
|
|
|
|
// DropCollection
|
|
ser := servicepb.CollectionName{CollectionName: "col1"}
|
|
reqDrop := internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &ser,
|
|
}
|
|
st, err = cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
//consume msg
|
|
ddMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
ddMs.SetPulsarClient(pulsarAddr)
|
|
ddMs.CreatePulsarConsumers(Params.DDChannelNames, "DDStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
ddMs.Start()
|
|
|
|
var consumeMsg ms.MsgStream = ddMs
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
break
|
|
}
|
|
}
|
|
ddMs.Close()
|
|
})
|
|
|
|
t.Run("TestBroadCastRequest", func(t *testing.T) {
|
|
|
|
proxyTimeTickStream := ms.NewPulsarMsgStream(ctx, 1024) //input stream
|
|
proxyTimeTickStream.SetPulsarClient(pulsarAddr)
|
|
proxyTimeTickStream.CreatePulsarProducers(Params.ProxyTimeTickChannelNames)
|
|
proxyTimeTickStream.Start()
|
|
|
|
writeNodeStream := ms.NewPulsarMsgStream(ctx, 1024) //input stream
|
|
writeNodeStream.SetPulsarClient(pulsarAddr)
|
|
writeNodeStream.CreatePulsarProducers(Params.WriteNodeTimeTickChannelNames)
|
|
writeNodeStream.Start()
|
|
|
|
ddMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
ddMs.SetPulsarClient(pulsarAddr)
|
|
ddMs.CreatePulsarConsumers(Params.DDChannelNames, "DDStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
ddMs.Start()
|
|
|
|
dMMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
dMMs.SetPulsarClient(pulsarAddr)
|
|
dMMs.CreatePulsarConsumers(Params.InsertChannelNames, "DMStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
dMMs.Start()
|
|
|
|
k2sMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
k2sMs.SetPulsarClient(pulsarAddr)
|
|
k2sMs.CreatePulsarConsumers(Params.K2SChannelNames, "K2SStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
k2sMs.Start()
|
|
|
|
ttsoftmsgs := [][2]uint64{
|
|
{0, 10},
|
|
}
|
|
msgSoftPackAddr := getTimeTickMsgPack(ttsoftmsgs)
|
|
|
|
err := proxyTimeTickStream.Produce(msgSoftPackAddr)
|
|
assert.Nil(t, err)
|
|
var dMMsgstream ms.MsgStream = dMMs
|
|
assert.True(t, receiveTimeTickMsg(&dMMsgstream))
|
|
var ddMsgstream ms.MsgStream = ddMs
|
|
assert.True(t, receiveTimeTickMsg(&ddMsgstream))
|
|
|
|
tthardmsgs := [][2]int{
|
|
{3, 10},
|
|
}
|
|
|
|
msghardPackAddr := getMsgPack(tthardmsgs)
|
|
err = writeNodeStream.Produce(msghardPackAddr)
|
|
assert.Nil(t, err)
|
|
var k2sMsgstream ms.MsgStream = k2sMs
|
|
assert.True(t, receiveTimeTickMsg(&k2sMsgstream))
|
|
|
|
conn, err := grpc.DialContext(ctx, Params.Address, grpc.WithInsecure(), grpc.WithBlock())
|
|
assert.Nil(t, err)
|
|
defer conn.Close()
|
|
|
|
cli := masterpb.NewMasterClient(conn)
|
|
|
|
sch := schemapb.CollectionSchema{
|
|
Name: "name" + strconv.FormatUint(rand.Uint64(), 10),
|
|
Description: "test collection",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{},
|
|
}
|
|
|
|
schemaBytes, err := proto.Marshal(&sch)
|
|
assert.Nil(t, err)
|
|
|
|
createCollectionReq := internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreateCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
}
|
|
st, err := cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var consumeMsg ms.MsgStream = ddMs
|
|
var createCollectionMsg *ms.CreateCollectionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
createCollectionMsg = v.(*ms.CreateCollectionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, createCollectionReq.MsgType, createCollectionMsg.CreateCollectionRequest.MsgType)
|
|
assert.Equal(t, createCollectionReq.ReqID, createCollectionMsg.CreateCollectionRequest.ReqID)
|
|
assert.Equal(t, createCollectionReq.Timestamp, createCollectionMsg.CreateCollectionRequest.Timestamp)
|
|
assert.Equal(t, createCollectionReq.ProxyID, createCollectionMsg.CreateCollectionRequest.ProxyID)
|
|
|
|
////////////////////////////CreatePartition////////////////////////
|
|
partitionName := "partitionName" + strconv.FormatUint(rand.Uint64(), 10)
|
|
createPartitionReq := internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{
|
|
CollectionName: sch.Name,
|
|
Tag: partitionName,
|
|
},
|
|
}
|
|
|
|
st, err = cli.CreatePartition(ctx, &createPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var createPartitionMsg *ms.CreatePartitionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
createPartitionMsg = v.(*ms.CreatePartitionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, createPartitionReq.MsgType, createPartitionMsg.CreatePartitionRequest.MsgType)
|
|
assert.Equal(t, createPartitionReq.ReqID, createPartitionMsg.CreatePartitionRequest.ReqID)
|
|
assert.Equal(t, createPartitionReq.Timestamp, createPartitionMsg.CreatePartitionRequest.Timestamp)
|
|
assert.Equal(t, createPartitionReq.ProxyID, createPartitionMsg.CreatePartitionRequest.ProxyID)
|
|
assert.Equal(t, createPartitionReq.PartitionName.CollectionName, createPartitionMsg.CreatePartitionRequest.PartitionName.CollectionName)
|
|
assert.Equal(t, createPartitionReq.PartitionName.Tag, createPartitionMsg.CreatePartitionRequest.PartitionName.Tag)
|
|
|
|
////////////////////////////DropPartition////////////////////////
|
|
dropPartitionReq := internalpb.DropPartitionRequest{
|
|
MsgType: internalpb.MsgType_kDropPartition,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{
|
|
CollectionName: sch.Name,
|
|
Tag: partitionName,
|
|
},
|
|
}
|
|
|
|
st, err = cli.DropPartition(ctx, &dropPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var dropPartitionMsg *ms.DropPartitionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
dropPartitionMsg = v.(*ms.DropPartitionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, dropPartitionReq.MsgType, dropPartitionMsg.DropPartitionRequest.MsgType)
|
|
assert.Equal(t, dropPartitionReq.ReqID, dropPartitionMsg.DropPartitionRequest.ReqID)
|
|
assert.Equal(t, dropPartitionReq.Timestamp, dropPartitionMsg.DropPartitionRequest.Timestamp)
|
|
assert.Equal(t, dropPartitionReq.ProxyID, dropPartitionMsg.DropPartitionRequest.ProxyID)
|
|
assert.Equal(t, dropPartitionReq.PartitionName.CollectionName, dropPartitionMsg.DropPartitionRequest.PartitionName.CollectionName)
|
|
|
|
////////////////////////////DropCollection////////////////////////
|
|
dropCollectionReq := internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{CollectionName: sch.Name},
|
|
}
|
|
|
|
st, err = cli.DropCollection(ctx, &dropCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var dropCollectionMsg *ms.DropCollectionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
dropCollectionMsg = v.(*ms.DropCollectionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, dropCollectionReq.MsgType, dropCollectionMsg.DropCollectionRequest.MsgType)
|
|
assert.Equal(t, dropCollectionReq.ReqID, dropCollectionMsg.DropCollectionRequest.ReqID)
|
|
assert.Equal(t, dropCollectionReq.Timestamp, dropCollectionMsg.DropCollectionRequest.Timestamp)
|
|
assert.Equal(t, dropCollectionReq.ProxyID, dropCollectionMsg.DropCollectionRequest.ProxyID)
|
|
assert.Equal(t, dropCollectionReq.CollectionName.CollectionName, dropCollectionMsg.DropCollectionRequest.CollectionName.CollectionName)
|
|
})
|
|
|
|
t.Run("TestSegmentManager_RPC", func(t *testing.T) {
|
|
collName := "test_coll"
|
|
partitionTag := "test_part"
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: collName,
|
|
Description: "test coll",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, Name: "f1", IsPrimaryKey: false, DataType: schemapb.DataType_INT32},
|
|
{FieldID: 1, Name: "f1", IsPrimaryKey: false, DataType: schemapb.DataType_VECTOR_FLOAT, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
|
|
},
|
|
}
|
|
schemaBytes, err := proto.Marshal(schema)
|
|
assert.Nil(t, err)
|
|
_, err = cli.CreateCollection(ctx, &internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreateCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
})
|
|
assert.Nil(t, err)
|
|
_, err = cli.CreatePartition(ctx, &internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 2,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{
|
|
CollectionName: collName,
|
|
Tag: partitionTag,
|
|
},
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
resp, err := cli.AssignSegmentID(ctx, &internalpb.AssignSegIDRequest{
|
|
PeerID: 1,
|
|
Role: internalpb.PeerRole_Proxy,
|
|
PerChannelReq: []*internalpb.SegIDRequest{
|
|
{Count: 10000, ChannelID: 0, CollName: collName, PartitionTag: partitionTag},
|
|
},
|
|
})
|
|
assert.Nil(t, err)
|
|
assignments := resp.GetPerChannelAssignment()
|
|
assert.EqualValues(t, 1, len(assignments))
|
|
assert.EqualValues(t, commonpb.ErrorCode_SUCCESS, assignments[0].Status.ErrorCode)
|
|
assert.EqualValues(t, collName, assignments[0].CollName)
|
|
assert.EqualValues(t, partitionTag, assignments[0].PartitionTag)
|
|
assert.EqualValues(t, int32(0), assignments[0].ChannelID)
|
|
assert.EqualValues(t, uint32(10000), assignments[0].Count)
|
|
|
|
// test stats
|
|
segID := assignments[0].SegID
|
|
pulsarAddress := Params.PulsarAddress
|
|
msgStream := ms.NewPulsarMsgStream(ctx, 1024)
|
|
msgStream.SetPulsarClient(pulsarAddress)
|
|
msgStream.CreatePulsarProducers([]string{Params.QueryNodeStatsChannelName})
|
|
msgStream.Start()
|
|
defer msgStream.Close()
|
|
|
|
err = msgStream.Produce(&ms.MsgPack{
|
|
BeginTs: 102,
|
|
EndTs: 104,
|
|
Msgs: []ms.TsMsg{
|
|
&ms.QueryNodeStatsMsg{
|
|
QueryNodeStats: internalpb.QueryNodeStats{
|
|
MsgType: internalpb.MsgType_kQueryNodeStats,
|
|
PeerID: 1,
|
|
SegStats: []*internalpb.SegmentStats{
|
|
{SegmentID: segID, MemorySize: 600000000, NumRows: 1000000, RecentlyModified: true},
|
|
},
|
|
},
|
|
BaseMsg: ms.BaseMsg{
|
|
HashValues: []uint32{0},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
segMeta, err := svr.metaTable.GetSegmentByID(segID)
|
|
assert.Nil(t, err)
|
|
assert.EqualValues(t, 1000000, segMeta.GetNumRows())
|
|
assert.EqualValues(t, int64(600000000), segMeta.GetMemSize())
|
|
|
|
ser := servicepb.CollectionName{CollectionName: collName}
|
|
reqDrop := internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProxyID: 1,
|
|
CollectionName: &ser,
|
|
}
|
|
|
|
// DropCollection
|
|
st, err := cli.DropCollection(ctx, &reqDrop)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
})
|
|
|
|
cancel()
|
|
conn.Close()
|
|
svr.Close()
|
|
}
|