From e7125672d4777f0ddbb903e28c2f598b62107c8d Mon Sep 17 00:00:00 2001
From: "yihao.dai" <yihao.dai@zilliz.com>
Date: Tue, 13 Jun 2023 22:08:38 +0800
Subject: [PATCH] Check metric type at watchChannelTask (#24860)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 internal/querynodev2/services.go      |  7 +++++++
 internal/querynodev2/services_test.go | 13 +++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/internal/querynodev2/services.go b/internal/querynodev2/services.go
index f782f78d1..8aea1e62d 100644
--- a/internal/querynodev2/services.go
+++ b/internal/querynodev2/services.go
@@ -192,6 +192,7 @@ func (node *QueryNode) WatchDmChannels(ctx context.Context, req *querypb.WatchDm
 
 	log.Info("received watch channel request",
 		zap.Int64("version", req.GetVersion()),
+		zap.String("metricType", req.GetLoadMeta().GetMetricType()),
 	)
 
 	// check node healthy
@@ -211,6 +212,12 @@ func (node *QueryNode) WatchDmChannels(ctx context.Context, req *querypb.WatchDm
 		return status, nil
 	}
 
+	// check metric type
+	if req.GetLoadMeta().GetMetricType() == "" {
+		err := fmt.Errorf("empty metric type, collection = %d", req.GetCollectionID())
+		return merr.Status(err), nil
+	}
+
 	if !node.subscribingChannels.Insert(channel.GetChannelName()) {
 		msg := "channel subscribing..."
 		log.Warn(msg)
diff --git a/internal/querynodev2/services_test.go b/internal/querynodev2/services_test.go
index 4197dd1ba..29638f6dd 100644
--- a/internal/querynodev2/services_test.go
+++ b/internal/querynodev2/services_test.go
@@ -250,6 +250,7 @@ func (suite *ServiceSuite) TestWatchDmChannelsInt64() {
 			LoadType:     querypb.LoadType_LoadCollection,
 			CollectionID: suite.collectionID,
 			PartitionIDs: suite.partitionIDs,
+			MetricType:   defaultMetricType,
 		},
 	}
 
@@ -298,6 +299,7 @@ func (suite *ServiceSuite) TestWatchDmChannelsVarchar() {
 			LoadType:     querypb.LoadType_LoadCollection,
 			CollectionID: suite.collectionID,
 			PartitionIDs: suite.partitionIDs,
+			MetricType:   defaultMetricType,
 		},
 	}
 
@@ -342,6 +344,9 @@ func (suite *ServiceSuite) TestWatchDmChannels_Failed() {
 			},
 		},
 		Schema: segments.GenTestCollectionSchema(suite.collectionName, schemapb.DataType_Int64),
+		LoadMeta: &querypb.LoadMetaInfo{
+			MetricType: defaultMetricType,
+		},
 	}
 
 	// init msgstream failed
@@ -364,6 +369,14 @@ func (suite *ServiceSuite) TestWatchDmChannels_Failed() {
 	status, err = suite.node.WatchDmChannels(ctx, req)
 	suite.NoError(err)
 	suite.Equal(commonpb.ErrorCode_NotReadyServe, status.GetErrorCode())
+
+	// empty metric type
+	req.LoadMeta.MetricType = ""
+	req.Base.TargetID = paramtable.GetNodeID()
+	suite.node.UpdateStateCode(commonpb.StateCode_Healthy)
+	status, err = suite.node.WatchDmChannels(ctx, req)
+	suite.NoError(err)
+	suite.Equal(commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
 }
 
 func (suite *ServiceSuite) TestUnsubDmChannels_Normal() {
-- 
GitLab