milvus/internal/master/master_test.go
XuanYang-cn 03b5d32569 Add default fields when create collection
Signed-off-by: XuanYang-cn <xuan.yang@zilliz.com>
2021-01-04 12:03:29 +08:00

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()
}