Add probe pk command

states/probe.go

@@ -19,9 +19,12 @@ import (
 	commonpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/commonpb"
 	indexpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/indexpb"
 	"github.com/milvus-io/birdwatcher/proto/v2.2/internalpb"
+	internalpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/internalpb"
 	"github.com/milvus-io/birdwatcher/proto/v2.2/planpb"
 	querypbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/querypb"
+	schemapbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/schemapb"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
 	"github.com/spf13/cobra"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"google.golang.org/grpc"
@@ -36,6 +39,8 @@ func GetProbeCmd(cli clientv3.KV, basePath string) *cobra.Command {
 	probeCmd.AddCommand(
 		// probe query
 		getProbeQueryCmd(cli, basePath),
+		// probe pk
+		getProbePKCmd(cli, basePath),
 	)
 
 	return probeCmd
@@ -131,6 +136,193 @@ func getProbeQueryCmd(cli clientv3.KV, basePath string) *cobra.Command {
 	return cmd
 }
 
+func getProbePKCmd(cli clientv3.KV, basePath string) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "pk",
+		Short: "probe pk in segment",
+		Run: func(cmd *cobra.Command, args []string) {
+			collID, err := cmd.Flags().GetInt64("collection")
+			if err != nil {
+				fmt.Println(err.Error())
+				return
+			}
+			pk, err := cmd.Flags().GetString("pk")
+			if err != nil {
+				fmt.Println(err.Error())
+				return
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			coll, err := common.GetCollectionByIDVersion(ctx, cli, basePath, etcdversion.GetVersion(), collID)
+			if err != nil {
+				fmt.Println(err.Error())
+				return
+			}
+
+			pkf, _ := coll.GetPKField()
+			var datatype schemapbv2.DataType
+			var val *planpb.GenericValue
+			switch pkf.DataType {
+			case models.DataTypeVarChar:
+				datatype = schemapbv2.DataType_VarChar
+				val = &planpb.GenericValue{
+					Val: &planpb.GenericValue_StringVal{
+						StringVal: pk,
+					},
+				}
+			case models.DataTypeInt64:
+				datatype = schemapbv2.DataType_Int64
+				pkv, err := strconv.ParseInt(pk, 10, 64)
+				if err != nil {
+					fmt.Println(err.Error())
+					return
+				}
+				val = &planpb.GenericValue{
+					Val: &planpb.GenericValue_Int64Val{
+						Int64Val: pkv,
+					},
+				}
+			}
+
+			var outputFields []int64
+			fieldIDName := make(map[int64]string)
+			for _, f := range coll.Schema.Fields {
+				if f.FieldID >= 100 {
+					outputFields = append(outputFields, f.FieldID)
+				}
+				fieldIDName[f.FieldID] = f.Name
+			}
+			plan := &planpb.PlanNode{
+				Node: &planpb.PlanNode_Predicates{
+					Predicates: &planpb.Expr{
+						Expr: &planpb.Expr_TermExpr{
+							TermExpr: &planpb.TermExpr{
+								ColumnInfo: &planpb.ColumnInfo{
+									FieldId:      pkf.FieldID,
+									DataType:     datatype,
+									IsAutoID:     pkf.AutoID,
+									IsPrimaryKey: pkf.IsPrimaryKey,
+								},
+								Values: []*planpb.GenericValue{
+									val,
+								},
+							},
+						},
+					},
+				},
+				OutputFieldIds: outputFields,
+			}
+
+			bs, _ := proto.Marshal(plan)
+
+			sessions, err := common.ListSessions(cli, basePath)
+			if err != nil {
+				fmt.Println("failed to list online sessions", err.Error())
+				return
+			}
+
+			qc, err := getQueryCoordClient(sessions)
+			if err != nil {
+				fmt.Println("failed to connect querycoord", err.Error())
+				return
+			}
+
+			resp, err := qc.GetShardLeaders(ctx, &querypbv2.GetShardLeadersRequest{
+				Base:         &commonpbv2.MsgBase{},
+				CollectionID: collID,
+			})
+			if err != nil {
+				fmt.Println(err.Error())
+				return
+			}
+
+			if resp.GetStatus().GetErrorCode() != commonpbv2.ErrorCode_Success {
+				fmt.Println("failed to list shard leader", resp.GetStatus().GetErrorCode())
+				return
+			}
+
+			qns, err := getQueryNodeClients(sessions)
+			if err != nil {
+				fmt.Println("failed to connect querynodes", err.Error())
+				return
+			}
+			if len(qns) == 0 {
+				fmt.Println("no querynode online")
+				return
+			}
+
+			for nodeID, qn := range qns {
+				resp, err := qn.GetDataDistribution(ctx, &querypbv2.GetDataDistributionRequest{
+					Base: &commonpbv2.MsgBase{TargetID: nodeID},
+				})
+				if err != nil {
+					fmt.Println("failed to get data distribution from node", nodeID, err.Error())
+					continue
+				}
+				for _, segInfo := range resp.GetSegments() {
+					if segInfo.GetCollection() != collID {
+						continue
+					}
+					result, err := qn.Query(ctx, &querypbv2.QueryRequest{
+						SegmentIDs:      []int64{segInfo.ID},
+						DmlChannels:     []string{segInfo.GetChannel()},
+						FromShardLeader: true, // query single segment
+						Scope:           querypbv2.DataScope_Historical,
+						Req: &internalpbv2.RetrieveRequest{
+							Base:               &commonpbv2.MsgBase{TargetID: nodeID, MsgID: time.Now().Unix()},
+							CollectionID:       segInfo.Collection,
+							PartitionIDs:       []int64{segInfo.Partition},
+							SerializedExprPlan: bs,
+							OutputFieldsId:     outputFields,
+							TravelTimestamp:    uint64((time.Now().UnixNano() / int64(time.Millisecond)) << 18),
+							Limit:              -1, // unlimited
+						},
+					})
+					if err != nil {
+						fmt.Println(err.Error())
+						continue
+					}
+					if result.GetStatus().GetErrorCode() != commonpbv2.ErrorCode_Success {
+						fmt.Printf("failed to probe pk on segment %d, reason: %s\n", segInfo.GetID(), result.GetStatus().GetReason())
+						continue
+					}
+					if GetSizeOfIDs(result.GetIds()) == 0 {
+						continue
+					}
+					fmt.Printf("PK %s found on segment %d\n", pk, segInfo.GetID())
+					for _, fd := range result.GetFieldsData() {
+						fmt.Printf("Field %s, value: %v\n", fieldIDName[fd.GetFieldId()], fd.GetField())
+					}
+				}
+			}
+		},
+	}
+
+	cmd.Flags().Int64("collection", 0, "collection id for probe")
+	cmd.Flags().String("pk", "", "pk value to probe")
+
+	return cmd
+}
+
+func GetSizeOfIDs(data *schemapbv2.IDs) int {
+	result := 0
+	if data.IdField == nil {
+		return result
+	}
+
+	switch data.GetIdField().(type) {
+	case *schemapbv2.IDs_IntId:
+		result = len(data.GetIntId().GetData())
+	case *schemapbv2.IDs_StrId:
+		result = len(data.GetStrId().GetData())
+	default:
+	}
+
+	return result
+}
+
 func getQueryCoordClient(sessions []*models.Session) (querypbv2.QueryCoordClient, error) {
 	for _, session := range sessions {
 		if strings.ToLower(session.ServerName) != "querycoord" {