milvus/internal/proxy/rootcoord_mock_test.go
Jiquan Long ff2a68e65a
Fix collection not exist when tried to do recovery (#21471)
Signed-off-by: longjiquan <jiquan.long@zilliz.com>
2023-01-04 16:37:35 +08:00

1227 lines
39 KiB
Go

// Licensed to the LF AI & Data foundation 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 proxy
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"time"
"github.com/milvus-io/milvus/internal/common"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/funcutil"
"github.com/milvus-io/milvus/internal/util/uniquegenerator"
"github.com/golang/protobuf/proto"
"github.com/milvus-io/milvus/internal/util/milvuserrors"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/proxypb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
)
type collectionMeta struct {
name string
id typeutil.UniqueID
schema *schemapb.CollectionSchema
shardsNum int32
virtualChannelNames []string
physicalChannelNames []string
createdTimestamp uint64
createdUtcTimestamp uint64
}
type partitionMeta struct {
createdTimestamp uint64
createdUtcTimestamp uint64
}
type partitionMap struct {
collID typeutil.UniqueID
// naive inverted index
partitionName2ID map[string]typeutil.UniqueID
partitionID2Name map[typeutil.UniqueID]string
partitionID2Meta map[typeutil.UniqueID]partitionMeta
}
type RootCoordMockOption func(mock *RootCoordMock)
type describeCollectionFuncType func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error)
type showPartitionsFuncType func(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error)
type RootCoordMock struct {
nodeID typeutil.UniqueID
address string
state atomic.Value // internal.StateCode
statisticsChannel string
timeTickChannel string
// naive inverted index
collName2ID map[string]typeutil.UniqueID
collID2Meta map[typeutil.UniqueID]collectionMeta
collAlias2ID map[string]typeutil.UniqueID
collMtx sync.RWMutex
// TODO(dragondriver): need default partition?
collID2Partitions map[typeutil.UniqueID]partitionMap
partitionMtx sync.RWMutex
describeCollectionFunc describeCollectionFuncType
showPartitionsFunc showPartitionsFuncType
showConfigurationsFunc showConfigurationsFuncType
getMetricsFunc getMetricsFuncType
// TODO(dragondriver): index-related
// TODO(dragondriver): segment-related
// TODO(dragondriver): TimeTick-related
lastTs typeutil.Timestamp
lastTsMtx sync.Mutex
checkHealthFunc func(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error)
}
func (coord *RootCoordMock) CreateAlias(ctx context.Context, req *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.Lock()
defer coord.collMtx.Unlock()
_, exist := coord.collAlias2ID[req.Alias]
if exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("duplicate collection alias, alias = %s", req.Alias),
}, nil
}
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("aliased collection name does not exist, name = %s", req.CollectionName),
}, nil
}
coord.collAlias2ID[req.Alias] = collID
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) DropAlias(ctx context.Context, req *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.Lock()
defer coord.collMtx.Unlock()
_, exist := coord.collAlias2ID[req.Alias]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("alias does not exist, alias = %s", req.Alias),
}, nil
}
delete(coord.collAlias2ID, req.Alias)
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) AlterAlias(ctx context.Context, req *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.Lock()
defer coord.collMtx.Unlock()
_, exist := coord.collAlias2ID[req.Alias]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: fmt.Sprintf("alias does not exist, alias = %s", req.Alias),
}, nil
}
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: fmt.Sprintf("aliased collection name does not exist, name = %s", req.CollectionName),
}, nil
}
coord.collAlias2ID[req.Alias] = collID
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) updateState(state commonpb.StateCode) {
coord.state.Store(state)
}
func (coord *RootCoordMock) getState() commonpb.StateCode {
return coord.state.Load().(commonpb.StateCode)
}
func (coord *RootCoordMock) healthy() bool {
return coord.getState() == commonpb.StateCode_Healthy
}
func (coord *RootCoordMock) Init() error {
coord.updateState(commonpb.StateCode_Initializing)
return nil
}
func (coord *RootCoordMock) Start() error {
defer coord.updateState(commonpb.StateCode_Healthy)
return nil
}
func (coord *RootCoordMock) Stop() error {
defer coord.updateState(commonpb.StateCode_Abnormal)
return nil
}
func (coord *RootCoordMock) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
NodeID: coord.nodeID,
Role: typeutil.RootCoordRole,
StateCode: coord.getState(),
ExtraInfo: nil,
},
SubcomponentStates: nil,
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
}, nil
}
func (coord *RootCoordMock) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.StringResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
}, nil
}
return &milvuspb.StringResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Value: coord.statisticsChannel,
}, nil
}
func (coord *RootCoordMock) Register() error {
return nil
}
func (coord *RootCoordMock) GetTimeTickChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.StringResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
}, nil
}
return &milvuspb.StringResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Value: coord.timeTickChannel,
}, nil
}
func (coord *RootCoordMock) CreateCollection(ctx context.Context, req *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.Lock()
defer coord.collMtx.Unlock()
_, exist := coord.collName2ID[req.CollectionName]
if exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: milvuserrors.MsgCollectionAlreadyExist(req.CollectionName),
}, nil
}
var schema schemapb.CollectionSchema
err := proto.Unmarshal(req.Schema, &schema)
if err != nil {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("failed to parse schema, error: %v", err),
}, nil
}
for i := range schema.Fields {
schema.Fields[i].FieldID = int64(common.StartOfUserFieldID + i)
}
collID := typeutil.UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
coord.collName2ID[req.CollectionName] = collID
var shardsNum int32
if req.ShardsNum <= 0 {
shardsNum = common.DefaultShardsNum
} else {
shardsNum = req.ShardsNum
}
virtualChannelNames := make([]string, 0, shardsNum)
physicalChannelNames := make([]string, 0, shardsNum)
for i := 0; i < int(shardsNum); i++ {
virtualChannelNames = append(virtualChannelNames, funcutil.GenRandomStr())
physicalChannelNames = append(physicalChannelNames, funcutil.GenRandomStr())
}
ts := uint64(time.Now().Nanosecond())
coord.collID2Meta[collID] = collectionMeta{
name: req.CollectionName,
id: collID,
schema: &schema,
shardsNum: shardsNum,
virtualChannelNames: virtualChannelNames,
physicalChannelNames: physicalChannelNames,
createdTimestamp: ts,
createdUtcTimestamp: ts,
}
coord.partitionMtx.Lock()
defer coord.partitionMtx.Unlock()
coord.collID2Partitions[collID] = partitionMap{
collID: collID,
partitionName2ID: make(map[string]typeutil.UniqueID),
partitionID2Name: make(map[typeutil.UniqueID]string),
partitionID2Meta: make(map[typeutil.UniqueID]partitionMeta),
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) DropCollection(ctx context.Context, req *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.Lock()
defer coord.collMtx.Unlock()
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
}, nil
}
delete(coord.collName2ID, req.CollectionName)
delete(coord.collID2Meta, collID)
coord.partitionMtx.Lock()
defer coord.partitionMtx.Unlock()
delete(coord.collID2Partitions, collID)
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) HasCollection(ctx context.Context, req *milvuspb.HasCollectionRequest) (*milvuspb.BoolResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Value: false,
}, nil
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
_, exist := coord.collName2ID[req.CollectionName]
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Value: exist,
}, nil
}
func (coord *RootCoordMock) SetDescribeCollectionFunc(f describeCollectionFuncType) {
coord.describeCollectionFunc = f
}
func (coord *RootCoordMock) ResetDescribeCollectionFunc() {
coord.describeCollectionFunc = nil
}
func (coord *RootCoordMock) DescribeCollection(ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Schema: nil,
CollectionID: 0,
}, nil
}
if coord.describeCollectionFunc != nil {
return coord.describeCollectionFunc(ctx, req)
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
var collID UniqueID
usingID := false
if req.CollectionName == "" {
usingID = true
}
collID, exist := coord.collName2ID[req.CollectionName]
if !exist && !usingID {
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
},
}, nil
}
if usingID {
collID = req.CollectionID
}
meta := coord.collID2Meta[collID]
if meta.shardsNum == 0 {
meta.shardsNum = int32(len(meta.virtualChannelNames))
}
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Schema: meta.schema,
CollectionID: collID,
ShardsNum: meta.shardsNum,
VirtualChannelNames: meta.virtualChannelNames,
PhysicalChannelNames: meta.physicalChannelNames,
CreatedTimestamp: meta.createdUtcTimestamp,
CreatedUtcTimestamp: meta.createdUtcTimestamp,
}, nil
}
func (coord *RootCoordMock) DescribeCollectionInternal(ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return coord.DescribeCollection(ctx, req)
}
func (coord *RootCoordMock) ShowCollections(ctx context.Context, req *milvuspb.ShowCollectionsRequest) (*milvuspb.ShowCollectionsResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
CollectionNames: nil,
}, nil
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
names := make([]string, 0, len(coord.collName2ID))
ids := make([]int64, 0, len(coord.collName2ID))
createdTimestamps := make([]uint64, 0, len(coord.collName2ID))
createdUtcTimestamps := make([]uint64, 0, len(coord.collName2ID))
for name, id := range coord.collName2ID {
names = append(names, name)
ids = append(ids, id)
meta := coord.collID2Meta[id]
createdTimestamps = append(createdTimestamps, meta.createdTimestamp)
createdUtcTimestamps = append(createdUtcTimestamps, meta.createdUtcTimestamp)
}
return &milvuspb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
CollectionNames: names,
CollectionIds: ids,
CreatedTimestamps: createdTimestamps,
CreatedUtcTimestamps: createdUtcTimestamps,
InMemoryPercentages: nil,
}, nil
}
func (coord *RootCoordMock) CreatePartition(ctx context.Context, req *milvuspb.CreatePartitionRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
}, nil
}
coord.partitionMtx.Lock()
defer coord.partitionMtx.Unlock()
_, partitionExist := coord.collID2Partitions[collID].partitionName2ID[req.PartitionName]
if partitionExist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: milvuserrors.MsgPartitionAlreadyExist(req.PartitionName),
}, nil
}
ts := uint64(time.Now().Nanosecond())
partitionID := typeutil.UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
coord.collID2Partitions[collID].partitionName2ID[req.PartitionName] = partitionID
coord.collID2Partitions[collID].partitionID2Name[partitionID] = req.PartitionName
coord.collID2Partitions[collID].partitionID2Meta[partitionID] = partitionMeta{
createdTimestamp: ts,
createdUtcTimestamp: ts,
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) DropPartition(ctx context.Context, req *milvuspb.DropPartitionRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
}, nil
}
coord.partitionMtx.Lock()
defer coord.partitionMtx.Unlock()
partitionID, partitionExist := coord.collID2Partitions[collID].partitionName2ID[req.PartitionName]
if !partitionExist {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: milvuserrors.MsgPartitionNotExist(req.PartitionName),
}, nil
}
delete(coord.collID2Partitions[collID].partitionName2ID, req.PartitionName)
delete(coord.collID2Partitions[collID].partitionID2Name, partitionID)
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) HasPartition(ctx context.Context, req *milvuspb.HasPartitionRequest) (*milvuspb.BoolResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Value: false,
}, nil
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
},
Value: false,
}, nil
}
coord.partitionMtx.RLock()
defer coord.partitionMtx.RUnlock()
_, partitionExist := coord.collID2Partitions[collID].partitionName2ID[req.PartitionName]
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Value: partitionExist,
}, nil
}
func (coord *RootCoordMock) ShowPartitions(ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("rootcoord is not healthy, state code = %s", commonpb.StateCode_name[int32(code)]),
},
PartitionNames: nil,
PartitionIDs: nil,
}, nil
}
if coord.showPartitionsFunc != nil {
return coord.showPartitionsFunc(ctx, req)
}
coord.collMtx.RLock()
defer coord.collMtx.RUnlock()
collID, exist := coord.collName2ID[req.CollectionName]
if !exist {
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: milvuserrors.MsgCollectionNotExist(req.CollectionName),
},
}, nil
}
coord.partitionMtx.RLock()
defer coord.partitionMtx.RUnlock()
l := len(coord.collID2Partitions[collID].partitionName2ID)
names := make([]string, 0, l)
ids := make([]int64, 0, l)
createdTimestamps := make([]uint64, 0, l)
createdUtcTimestamps := make([]uint64, 0, l)
for name, id := range coord.collID2Partitions[collID].partitionName2ID {
names = append(names, name)
ids = append(ids, id)
meta := coord.collID2Partitions[collID].partitionID2Meta[id]
createdTimestamps = append(createdTimestamps, meta.createdTimestamp)
createdUtcTimestamps = append(createdUtcTimestamps, meta.createdUtcTimestamp)
}
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
PartitionNames: names,
PartitionIDs: ids,
CreatedTimestamps: createdTimestamps,
CreatedUtcTimestamps: createdUtcTimestamps,
InMemoryPercentages: nil,
}, nil
}
func (coord *RootCoordMock) ShowPartitionsInternal(ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
return coord.ShowPartitions(ctx, req)
}
//func (coord *RootCoordMock) CreateIndex(ctx context.Context, req *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
// code := coord.state.Load().(commonpb.StateCode)
// if code != commonpb.StateCode_Healthy {
// return &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_UnexpectedError,
// Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
// }, nil
// }
// return &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_Success,
// Reason: "",
// }, nil
//}
//func (coord *RootCoordMock) DescribeIndex(ctx context.Context, req *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
// code := coord.state.Load().(commonpb.StateCode)
// if code != commonpb.StateCode_Healthy {
// return &milvuspb.DescribeIndexResponse{
// Status: &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_UnexpectedError,
// Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
// },
// IndexDescriptions: nil,
// }, nil
// }
// return &milvuspb.DescribeIndexResponse{
// Status: &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_Success,
// Reason: "",
// },
// IndexDescriptions: nil,
// }, nil
//}
//func (coord *RootCoordMock) GetIndexState(ctx context.Context, req *milvuspb.GetIndexStateRequest) (*indexpb.GetIndexStatesResponse, error) {
// code := coord.state.Load().(commonpb.StateCode)
// if code != commonpb.StateCode_Healthy {
// return &indexpb.GetIndexStatesResponse{
// Status: &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_UnexpectedError,
// Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
// },
// }, nil
// }
// return &indexpb.GetIndexStatesResponse{
// Status: &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_Success,
// Reason: "",
// },
// }, nil
//}
//func (coord *RootCoordMock) DropIndex(ctx context.Context, req *milvuspb.DropIndexRequest) (*commonpb.Status, error) {
// code := coord.state.Load().(commonpb.StateCode)
// if code != commonpb.StateCode_Healthy {
// return &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_UnexpectedError,
// Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
// }, nil
// }
// return &commonpb.Status{
// ErrorCode: commonpb.ErrorCode_Success,
// Reason: "",
// }, nil
//}
func (coord *RootCoordMock) AllocTimestamp(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &rootcoordpb.AllocTimestampResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Timestamp: 0,
Count: 0,
}, nil
}
coord.lastTsMtx.Lock()
defer coord.lastTsMtx.Unlock()
ts := uint64(time.Now().UnixNano())
if ts < coord.lastTs+typeutil.Timestamp(req.Count) {
ts = coord.lastTs + typeutil.Timestamp(req.Count)
}
coord.lastTs = ts
return &rootcoordpb.AllocTimestampResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Timestamp: ts,
Count: req.Count,
}, nil
}
func (coord *RootCoordMock) AllocID(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &rootcoordpb.AllocIDResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
ID: 0,
Count: 0,
}, nil
}
begin, _ := uniquegenerator.GetUniqueIntGeneratorIns().GetInts(int(req.Count))
return &rootcoordpb.AllocIDResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
ID: int64(begin),
Count: req.Count,
}, nil
}
func (coord *RootCoordMock) UpdateChannelTimeTick(ctx context.Context, req *internalpb.ChannelTimeTickMsg) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) DescribeSegment(ctx context.Context, req *milvuspb.DescribeSegmentRequest) (*milvuspb.DescribeSegmentResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.DescribeSegmentResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
IndexID: 0,
}, nil
}
return &milvuspb.DescribeSegmentResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
IndexID: 0,
BuildID: 0,
EnableIndex: false,
}, nil
}
func (coord *RootCoordMock) ShowSegments(ctx context.Context, req *milvuspb.ShowSegmentsRequest) (*milvuspb.ShowSegmentsResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.ShowSegmentsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
SegmentIDs: nil,
}, nil
}
return &milvuspb.ShowSegmentsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
SegmentIDs: nil,
}, nil
}
func (coord *RootCoordMock) DescribeSegments(ctx context.Context, req *rootcoordpb.DescribeSegmentsRequest) (*rootcoordpb.DescribeSegmentsResponse, error) {
panic("implement me")
}
func (coord *RootCoordMock) InvalidateCollectionMetaCache(ctx context.Context, in *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) SegmentFlushCompleted(ctx context.Context, in *datapb.SegmentFlushCompletedMsg) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func (coord *RootCoordMock) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
if !coord.healthy() {
return &internalpb.ShowConfigurationsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "unhealthy",
},
}, nil
}
if coord.getMetricsFunc != nil {
return coord.showConfigurationsFunc(ctx, req)
}
return &internalpb.ShowConfigurationsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "not implemented",
},
}, nil
}
func (coord *RootCoordMock) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
if !coord.healthy() {
return &milvuspb.GetMetricsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "unhealthy",
},
}, nil
}
if coord.getMetricsFunc != nil {
return coord.getMetricsFunc(ctx, req)
}
return &milvuspb.GetMetricsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "not implemented",
},
Response: "",
ComponentName: "",
}, nil
}
func (coord *RootCoordMock) Import(ctx context.Context, req *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.ImportResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Tasks: make([]int64, 0),
}, nil
}
return &milvuspb.ImportResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Tasks: make([]int64, 3),
}, nil
}
func (coord *RootCoordMock) GetImportState(ctx context.Context, req *milvuspb.GetImportStateRequest) (*milvuspb.GetImportStateResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.GetImportStateResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
RowCount: 0,
IdList: make([]int64, 0),
}, nil
}
return &milvuspb.GetImportStateResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
RowCount: 10,
IdList: make([]int64, 3),
}, nil
}
func (coord *RootCoordMock) ListImportTasks(ctx context.Context, in *milvuspb.ListImportTasksRequest) (*milvuspb.ListImportTasksResponse, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &milvuspb.ListImportTasksResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
},
Tasks: make([]*milvuspb.GetImportStateResponse, 0),
}, nil
}
return &milvuspb.ListImportTasksResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
},
Tasks: make([]*milvuspb.GetImportStateResponse, 3),
}, nil
}
func (coord *RootCoordMock) ReportImport(ctx context.Context, req *rootcoordpb.ImportResult) (*commonpb.Status, error) {
code := coord.state.Load().(commonpb.StateCode)
if code != commonpb.StateCode_Healthy {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: fmt.Sprintf("state code = %s", commonpb.StateCode_name[int32(code)]),
}, nil
}
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}, nil
}
func NewRootCoordMock(opts ...RootCoordMockOption) *RootCoordMock {
rc := &RootCoordMock{
nodeID: typeutil.UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
address: funcutil.GenRandomStr(), // TODO(dragondriver): random address
statisticsChannel: funcutil.GenRandomStr(),
timeTickChannel: funcutil.GenRandomStr(),
collName2ID: make(map[string]typeutil.UniqueID),
collID2Meta: make(map[typeutil.UniqueID]collectionMeta),
collID2Partitions: make(map[typeutil.UniqueID]partitionMap),
lastTs: typeutil.Timestamp(time.Now().UnixNano()),
}
for _, opt := range opts {
opt(rc)
}
return rc
}
func (coord *RootCoordMock) CreateCredential(ctx context.Context, req *internalpb.CredentialInfo) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) UpdateCredential(ctx context.Context, req *internalpb.CredentialInfo) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) DeleteCredential(ctx context.Context, req *milvuspb.DeleteCredentialRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) ListCredUsers(ctx context.Context, req *milvuspb.ListCredUsersRequest) (*milvuspb.ListCredUsersResponse, error) {
return &milvuspb.ListCredUsersResponse{}, nil
}
func (coord *RootCoordMock) GetCredential(ctx context.Context, req *rootcoordpb.GetCredentialRequest) (*rootcoordpb.GetCredentialResponse, error) {
return &rootcoordpb.GetCredentialResponse{}, nil
}
func (coord *RootCoordMock) CreateRole(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) DropRole(ctx context.Context, req *milvuspb.DropRoleRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) OperateUserRole(ctx context.Context, req *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) SelectRole(ctx context.Context, req *milvuspb.SelectRoleRequest) (*milvuspb.SelectRoleResponse, error) {
return &milvuspb.SelectRoleResponse{}, nil
}
func (coord *RootCoordMock) SelectUser(ctx context.Context, req *milvuspb.SelectUserRequest) (*milvuspb.SelectUserResponse, error) {
return &milvuspb.SelectUserResponse{}, nil
}
func (coord *RootCoordMock) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantRequest) (*milvuspb.SelectGrantResponse, error) {
return &milvuspb.SelectGrantResponse{}, nil
}
func (coord *RootCoordMock) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return &internalpb.ListPolicyResponse{}, nil
}
func (coord *RootCoordMock) AlterCollection(ctx context.Context, request *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
if coord.checkHealthFunc != nil {
return coord.checkHealthFunc(ctx, req)
}
return &milvuspb.CheckHealthResponse{IsHealthy: true}, nil
}
type DescribeCollectionFunc func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error)
type ShowPartitionsFunc func(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error)
type ShowSegmentsFunc func(ctx context.Context, request *milvuspb.ShowSegmentsRequest) (*milvuspb.ShowSegmentsResponse, error)
type DescribeSegmentsFunc func(ctx context.Context, request *rootcoordpb.DescribeSegmentsRequest) (*rootcoordpb.DescribeSegmentsResponse, error)
type ImportFunc func(ctx context.Context, req *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error)
type DropCollectionFunc func(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error)
type GetGetCredentialFunc func(ctx context.Context, req *rootcoordpb.GetCredentialRequest) (*rootcoordpb.GetCredentialResponse, error)
type mockRootCoord struct {
types.RootCoord
DescribeCollectionFunc
ShowPartitionsFunc
ShowSegmentsFunc
DescribeSegmentsFunc
ImportFunc
DropCollectionFunc
GetGetCredentialFunc
}
func (m *mockRootCoord) GetCredential(ctx context.Context, request *rootcoordpb.GetCredentialRequest) (*rootcoordpb.GetCredentialResponse, error) {
if m.GetGetCredentialFunc != nil {
return m.GetGetCredentialFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
if m.DescribeCollectionFunc != nil {
return m.DescribeCollectionFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) DescribeCollectionInternal(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return m.DescribeCollection(ctx, request)
}
func (m *mockRootCoord) ShowPartitions(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
if m.ShowPartitionsFunc != nil {
return m.ShowPartitionsFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) ShowPartitionsInternal(ctx context.Context, request *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
return m.ShowPartitions(ctx, request)
}
func (m *mockRootCoord) ShowSegments(ctx context.Context, request *milvuspb.ShowSegmentsRequest) (*milvuspb.ShowSegmentsResponse, error) {
if m.ShowSegmentsFunc != nil {
return m.ShowSegmentsFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) Import(ctx context.Context, request *milvuspb.ImportRequest) (*milvuspb.ImportResponse, error) {
if m.ImportFunc != nil {
return m.ImportFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) DropCollection(ctx context.Context, request *milvuspb.DropCollectionRequest) (*commonpb.Status, error) {
if m.DropCollectionFunc != nil {
return m.DropCollectionFunc(ctx, request)
}
return nil, errors.New("mock")
}
func (m *mockRootCoord) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return &internalpb.ListPolicyResponse{}, nil
}
func (m *mockRootCoord) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
return &milvuspb.CheckHealthResponse{
IsHealthy: true,
}, nil
}
func newMockRootCoord() *mockRootCoord {
return &mockRootCoord{}
}