Fix bug: describe collection return wrong shard number (#7841)

Signed-off-by: Nzhenshan.cao <zhenshan.cao@zilliz.com>

internal/common/common.go

@@ -24,4 +24,5 @@ const (
 	TimeStampField     = 1
 	RowIDFieldName     = "RowID"
 	TimeStampFieldName = "Timestamp"
+	DefaultShardsNum   = int32(2)
 )

internal/proxy/rootcoord_mock_test.go

@@ -225,7 +225,7 @@ func (coord *RootCoordMock) CreateCollection(ctx context.Context, req *milvuspb.
 
 	var shardsNum int32
 	if req.ShardsNum <= 0 {
-		shardsNum = 2
+		shardsNum = common.DefaultShardsNum
 	} else {
 		shardsNum = req.ShardsNum
 	}
@@ -364,6 +364,9 @@ func (coord *RootCoordMock) DescribeCollection(ctx context.Context, req *milvusp
 	}
 
 	meta := coord.collID2Meta[collID]
+	if meta.shardsNum == 0 {
+		meta.shardsNum = int32(len(meta.virtualChannelNames))
+	}
 
 	return &milvuspb.DescribeCollectionResponse{
 		Status: &commonpb.Status{
@@ -372,6 +375,7 @@ func (coord *RootCoordMock) DescribeCollection(ctx context.Context, req *milvusp
 		},
 		Schema:               meta.schema,
 		CollectionID:         collID,
+		ShardsNum:            meta.shardsNum,
 		VirtualChannelNames:  meta.virtualChannelNames,
 		PhysicalChannelNames: meta.physicalChannelNames,
 		CreatedTimestamp:     meta.createdUtcTimestamp,

internal/proxy/task.go

@@ -2701,7 +2701,7 @@ func (dct *describeCollectionTask) Execute(ctx context.Context) error {
 		dct.result.PhysicalChannelNames = result.PhysicalChannelNames
 		dct.result.CreatedTimestamp = result.CreatedTimestamp
 		dct.result.CreatedUtcTimestamp = result.CreatedUtcTimestamp
-
+		dct.result.ShardsNum = result.ShardsNum
 		for _, field := range result.Schema.Fields {
 			if field.FieldID >= common.StartOfUserFieldID {
 				dct.result.Schema.Fields = append(dct.result.Schema.Fields, &schemapb.FieldSchema{

internal/proxy/task_test.go

@@ -13,18 +13,16 @@ import (
 	"testing"
 	"time"
 
-	"github.com/milvus-io/milvus/internal/proto/querypb"
-
-	"github.com/milvus-io/milvus/internal/msgstream"
-
-	"github.com/milvus-io/milvus/internal/util/distance"
-
 	"github.com/golang/protobuf/proto"
+	"github.com/milvus-io/milvus/internal/common"
 	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/msgstream"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
+	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/proto/schemapb"
+	"github.com/milvus-io/milvus/internal/util/distance"
 	"github.com/milvus-io/milvus/internal/util/funcutil"
 	"github.com/milvus-io/milvus/internal/util/uniquegenerator"
 	"github.com/stretchr/testify/assert"
@@ -1212,6 +1210,129 @@ func TestDescribeCollectionTask(t *testing.T) {
 	assert.Equal(t, commonpb.ErrorCode_UnexpectedError, task.result.Status.ErrorCode)
 }
 
+func TestDescribeCollectionTask_ShardsNum1(t *testing.T) {
+	Params.Init()
+	rc := NewRootCoordMock()
+	rc.Start()
+	defer rc.Stop()
+	ctx := context.Background()
+	InitMetaCache(rc)
+	prefix := "TestDescribeCollectionTask"
+	dbName := ""
+	collectionName := prefix + funcutil.GenRandomStr()
+
+	shardsNum := int32(2)
+	int64Field := "int64"
+	floatVecField := "fvec"
+	dim := 128
+
+	schema := constructCollectionSchema(int64Field, floatVecField, dim, collectionName)
+	marshaledSchema, err := proto.Marshal(schema)
+	assert.NoError(t, err)
+
+	createColReq := &milvuspb.CreateCollectionRequest{
+		Base: &commonpb.MsgBase{
+			MsgType:   commonpb.MsgType_DropCollection,
+			MsgID:     100,
+			Timestamp: 100,
+		},
+		DbName:         dbName,
+		CollectionName: collectionName,
+		Schema:         marshaledSchema,
+		ShardsNum:      shardsNum,
+	}
+
+	rc.CreateCollection(ctx, createColReq)
+	globalMetaCache.GetCollectionID(ctx, collectionName)
+
+	//CreateCollection
+	task := &describeCollectionTask{
+		Condition: NewTaskCondition(ctx),
+		DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:   commonpb.MsgType_DescribeCollection,
+				MsgID:     100,
+				Timestamp: 100,
+			},
+			DbName:         dbName,
+			CollectionName: collectionName,
+		},
+		ctx:       ctx,
+		rootCoord: rc,
+		result:    nil,
+	}
+	err = task.PreExecute(ctx)
+	assert.Nil(t, err)
+
+	err = task.Execute(ctx)
+	assert.Nil(t, err)
+	assert.Equal(t, commonpb.ErrorCode_Success, task.result.Status.ErrorCode)
+	assert.Equal(t, shardsNum, task.result.ShardsNum)
+
+}
+
+func TestDescribeCollectionTask_ShardsNum2(t *testing.T) {
+	Params.Init()
+	rc := NewRootCoordMock()
+	rc.Start()
+	defer rc.Stop()
+	ctx := context.Background()
+	InitMetaCache(rc)
+	prefix := "TestDescribeCollectionTask"
+	dbName := ""
+	collectionName := prefix + funcutil.GenRandomStr()
+
+	int64Field := "int64"
+	floatVecField := "fvec"
+	dim := 128
+
+	schema := constructCollectionSchema(int64Field, floatVecField, dim, collectionName)
+	marshaledSchema, err := proto.Marshal(schema)
+	assert.NoError(t, err)
+
+	createColReq := &milvuspb.CreateCollectionRequest{
+		Base: &commonpb.MsgBase{
+			MsgType:   commonpb.MsgType_DropCollection,
+			MsgID:     100,
+			Timestamp: 100,
+		},
+		DbName:         dbName,
+		CollectionName: collectionName,
+		Schema:         marshaledSchema,
+	}
+
+	rc.CreateCollection(ctx, createColReq)
+	globalMetaCache.GetCollectionID(ctx, collectionName)
+
+	//CreateCollection
+	task := &describeCollectionTask{
+		Condition: NewTaskCondition(ctx),
+		DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:   commonpb.MsgType_DescribeCollection,
+				MsgID:     100,
+				Timestamp: 100,
+			},
+			DbName:         dbName,
+			CollectionName: collectionName,
+		},
+		ctx:       ctx,
+		rootCoord: rc,
+		result:    nil,
+	}
+	task.PreExecute(ctx)
+
+	// missing collectionID in globalMetaCache
+	err = task.Execute(ctx)
+	assert.Nil(t, err)
+
+	err = task.Execute(ctx)
+	assert.Nil(t, err)
+	assert.Equal(t, commonpb.ErrorCode_Success, task.result.Status.ErrorCode)
+	assert.Equal(t, common.DefaultShardsNum, task.result.ShardsNum)
+	rc.Stop()
+}
+
 func TestCreatePartitionTask(t *testing.T) {
 	Params.Init()
 	rc := NewRootCoordMock()

internal/querycoord/mock_3rd_component_test.go

@@ -20,6 +20,7 @@ import (
 	"strconv"
 	"sync"
 
+	"github.com/milvus-io/milvus/internal/common"
 	"github.com/milvus-io/milvus/internal/kv"
 	minioKV "github.com/milvus-io/milvus/internal/kv/minio"
 	"github.com/milvus-io/milvus/internal/log"
@@ -40,7 +41,6 @@ const (
 	defaultCollectionID = UniqueID(2021)
 	defaultPartitionID  = UniqueID(2021)
 	defaultSegmentID    = UniqueID(2021)
-	defaultShardsNum    = 2
 )
 
 func genCollectionSchema(collectionID UniqueID, isBinary bool) *schemapb.CollectionSchema {
@@ -330,7 +330,7 @@ func (data *dataCoordMock) GetRecoveryInfo(ctx context.Context, req *datapb.GetR
 		channelInfos := make([]*datapb.VchannelInfo, 0)
 		data.collections = append(data.collections, collectionID)
 		collectionName := funcutil.RandomString(8)
-		for i := 0; i < defaultShardsNum; i++ {
+		for i := int32(0); i < common.DefaultShardsNum; i++ {
 			vChannel := fmt.Sprintf("%s_%d_%d_v", collectionName, collectionID, i)
 			channelInfo := &datapb.VchannelInfo{
 				CollectionID: collectionID,

internal/rootcoord/root_coord.go

@@ -65,8 +65,6 @@ const (
 
 	// MetricRequestsSuccess used to count the num of successful requests
 	MetricRequestsSuccess = "success"
-
-	DefaultShardsNum = int32(2)
 )
 
 func metricProxy(v int64) string {

internal/rootcoord/root_coord_test.go

@@ -22,13 +22,13 @@ import (
 	"testing"
 	"time"
 
-	"github.com/milvus-io/milvus/internal/util/metricsinfo"
-
 	"github.com/golang/protobuf/proto"
+	"github.com/milvus-io/milvus/internal/common"
 	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
 	"github.com/milvus-io/milvus/internal/msgstream"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/util/metricsinfo"
 
 	"github.com/milvus-io/milvus/internal/proto/etcdpb"
 	"github.com/milvus-io/milvus/internal/proto/indexpb"
@@ -2043,9 +2043,9 @@ func TestRootCoord2(t *testing.T) {
 		assert.Equal(t, commonpb.ErrorCode_Success, rsp.Status.ErrorCode)
 		assert.Equal(t, collName, rsp.Schema.Name)
 		assert.Equal(t, collMeta.ID, rsp.CollectionID)
-		assert.Equal(t, DefaultShardsNum, int32(len(rsp.VirtualChannelNames)))
-		assert.Equal(t, DefaultShardsNum, int32(len(rsp.PhysicalChannelNames)))
-		assert.Equal(t, DefaultShardsNum, rsp.ShardsNum)
+		assert.Equal(t, common.DefaultShardsNum, int32(len(rsp.VirtualChannelNames)))
+		assert.Equal(t, common.DefaultShardsNum, int32(len(rsp.PhysicalChannelNames)))
+		assert.Equal(t, common.DefaultShardsNum, rsp.ShardsNum)
 	})
 	err = core.Stop()
 	assert.Nil(t, err)

internal/rootcoord/task.go

@@ -16,8 +16,7 @@ import (
 	"fmt"
 
 	"github.com/golang/protobuf/proto"
-	"go.uber.org/zap"
-
+	"github.com/milvus-io/milvus/internal/common"
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/etcdpb"
@@ -27,6 +26,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/schemapb"
 	"github.com/milvus-io/milvus/internal/util/tsoutil"
 	"github.com/milvus-io/milvus/internal/util/typeutil"
+	"go.uber.org/zap"
 )
 
 type reqTask interface {
@@ -92,7 +92,7 @@ func (t *CreateCollectionReqTask) Execute(ctx context.Context) error {
 		return fmt.Errorf("collection name = %s, schema.Name=%s", t.Req.CollectionName, schema.Name)
 	}
 	if t.Req.ShardsNum <= 0 {
-		t.Req.ShardsNum = DefaultShardsNum
+		t.Req.ShardsNum = common.DefaultShardsNum
 	}
 	log.Debug("CreateCollectionReqTask Execute", zap.Any("CollectionName", t.Req.CollectionName),
 		zap.Any("ShardsNum", t.Req.ShardsNum))