enhance: Refine inspect-pk command to support remote mode

states/backup_mock_connect.go

@@ -68,10 +68,6 @@ func (s *embedEtcdMockState) SetupCommands() {
 		// remove [subcommand] options...
 		// used for testing
 		etcd.RemoveCommand(s.client, s.instanceName, rootPath),
-		// download-pk
-		getDownloadPKCmd(s.client, rootPath),
-		// inspect-pk
-		getInspectPKCmd(s.client, rootPath),
 
 		// for testing
 		etcd.RepairCommand(s.client, rootPath),
@@ -262,7 +258,7 @@ func readFixLengthHeader[T proto.Message](rd *bufio.Reader, header T) error {
 	lb := make([]byte, 8)
 	lenRead, err := rd.Read(lb)
 	if err == io.EOF || lenRead < 8 {
-		return fmt.Errorf("File does not contains valid header")
+		return errors.New("File does not contains valid header")
 	}
 
 	nextBytes := binary.LittleEndian.Uint64(lb)

states/current_version.go

@@ -23,20 +23,20 @@ func CurrentVersionCommand() *cobra.Command {
 	return cmd
 }
 
-type setCurrentVersionParam struct {
+type SetCurrentVersionParam struct {
 	framework.ParamBase `use:"set current-version" desc:"set current version for etcd meta parsing"`
 	newVersion          string
 }
 
-func (p *setCurrentVersionParam) ParseArgs(args []string) error {
+func (p *SetCurrentVersionParam) ParseArgs(args []string) error {
 	if len(args) != 1 {
 		return errors.New("invalid parameter number")
 	}
 	p.newVersion = args[0]
 	return nil
 }
 
-func (s *InstanceState) SetCurrentVersionCommand(ctx context.Context, param setCurrentVersionParam) {
+func (s *InstanceState) SetCurrentVersionCommand(ctx context.Context, param *SetCurrentVersionParam) error {
 	switch param.newVersion {
 	case models.LTEVersion2_1:
 		fallthrough
@@ -49,39 +49,5 @@ func (s *InstanceState) SetCurrentVersionCommand(ctx context.Context, param setC
 	default:
 		fmt.Println("Invalid version string:", param.newVersion)
 	}
-}
-
-// SetCurrentVersionCommand returns command for set current-version.
-func SetCurrentVersionCommand() *cobra.Command {
-	cmd := &cobra.Command{
-		Use: "set",
-	}
-
-	subCmd := &cobra.Command{
-		Use: "current-version",
-		Run: func(_ *cobra.Command, args []string) {
-			if len(args) != 1 {
-				fmt.Println("invalid parameter numbers")
-				return
-			}
-
-			newVersion := args[0]
-			switch newVersion {
-			case models.LTEVersion2_1:
-				fallthrough
-			case "LTEVersion2_1":
-				etcdversion.SetVersion(models.LTEVersion2_1)
-			case models.GTEVersion2_2:
-				fallthrough
-			case "GTEVersion2_2":
-				etcdversion.SetVersion(models.GTEVersion2_2)
-			default:
-				fmt.Println("Invalid version string:", newVersion)
-			}
-		},
-	}
-
-	cmd.AddCommand(subCmd)
-
-	return cmd
+	return nil
 }

states/download_pk.go

@@ -7,91 +7,112 @@ import (
 	"io"
 	"os"
 	"path"
+	"strings"
+	"time"
 
+	"github.com/cockroachdb/errors"
 	"github.com/gosuri/uilive"
 	"github.com/manifoldco/promptui"
 	"github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
-	"github.com/spf13/cobra"
-	clientv3 "go.etcd.io/etcd/client/v3"
 
-	"github.com/milvus-io/birdwatcher/proto/v2.0/datapb"
+	"github.com/milvus-io/birdwatcher/framework"
+	"github.com/milvus-io/birdwatcher/models"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
 )
 
-func getDownloadPKCmd(cli clientv3.KV, basePath string) *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "download-pk",
-		Short: "download pk column of a collection",
-		RunE: func(cmd *cobra.Command, args []string) error {
-			collectionID, err := cmd.Flags().GetInt64("id")
-			if err != nil {
-				return err
-			}
+type DownloadPKParam struct {
+	framework.ParamBase `use:"download-pk" desc:"download segment pk with provided collection/segment id"`
+	MinioAddress        string `name:"minioAddr" default:"" desc:"the minio address to override, leave empty to use milvus.yaml value"`
+	CollectionID        int64  `name:"collection" default:"0" desc:"collection id to download"`
+	SegmentID           int64  `name:"segment" default:"0" desc:"segment id to download"`
+}
 
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-			coll, err := common.GetCollectionByIDVersion(ctx, cli, basePath, etcdversion.GetVersion(), collectionID)
-			if err != nil {
-				fmt.Println("Collection not found for id", collectionID)
-				return nil
-			}
+func (s *InstanceState) DownloadPKCommand(ctx context.Context, p *DownloadPKParam) error {
+	collection, err := common.GetCollectionByIDVersion(ctx, s.client, s.basePath, etcdversion.GetVersion(), p.CollectionID)
+	if err != nil {
+		return err
+	}
+	pkField, ok := collection.GetPKField()
+	if !ok {
+		return errors.New("pk field not found")
+	}
 
-			var pkID int64 = -1
+	segments, err := common.ListSegmentsVersion(ctx, s.client, s.basePath, etcdversion.GetVersion(), func(s *models.Segment) bool {
+		return s.CollectionID == p.CollectionID && (p.SegmentID == 0 || p.SegmentID == s.ID)
+	})
+	if err != nil {
+		return err
+	}
 
-			for _, field := range coll.Schema.Fields {
-				if field.IsPrimaryKey {
-					pkID = field.FieldID
-					break
-				}
-			}
+	minioClient, bucketName, rootPath, err := s.GetMinioClientFromCfg(ctx, p.MinioAddress)
+	if err != nil {
+		return err
+	}
 
-			if pkID < 0 {
-				fmt.Println("collection pk not found")
-				return nil
-			}
+	return s.downloadPKs(ctx, minioClient, bucketName, rootPath, p.CollectionID, pkField.FieldID, segments, s.writeLogfile)
+}
 
-			segments, err := common.ListSegments(cli, basePath, func(segment *datapb.SegmentInfo) bool {
-				return segment.CollectionID == collectionID
-			})
-			if err != nil {
-				return err
-			}
+func (s *InstanceState) writeLogfile(ctx context.Context, obj *minio.Object) error {
+	return nil
+}
 
-			p := promptui.Prompt{
-				Label: "BucketName",
-			}
-			bucketName, err := p.Run()
-			if err != nil {
-				return err
-			}
+func (s *InstanceState) downloadPKs(ctx context.Context, cli *minio.Client, bucketName, rootPath string, collID int64, pkID int64, segments []*models.Segment, handler func(context.Context, *minio.Object) error) error {
+	folder := fmt.Sprintf("dlpks_%s", time.Now().Format("20060102150406"))
+	err := os.Mkdir(folder, 0o777)
+	if err != nil {
+		fmt.Println("Failed to create folder,", err.Error())
+	}
 
-			minioClient, err := getMinioClient()
-			if err != nil {
-				fmt.Println("cannot get minio client", err.Error())
-				return nil
-			}
-			exists, err := minioClient.BucketExists(context.Background(), bucketName)
-			if err != nil {
-				return err
-			}
-			if !exists {
-				fmt.Printf("bucket %s not exists\n", bucketName)
-				return nil
-			}
+	pd := uilive.New()
+	pf := "Downloading pk files ... %d%%(%d/%d)\n"
+	pd.Start()
+	fmt.Fprintf(pd, pf, 0, 0, len(segments))
+	defer pd.Stop()
 
-			for _, segment := range segments {
-				common.FillFieldsIfV2(cli, basePath, segment)
+	count := 0
+	for i, segment := range segments {
+		targetFolder := fmt.Sprintf("%s/%d", folder, segment.ID)
+		os.Mkdir(targetFolder, 0o777)
+		for _, fieldBinlog := range segment.GetBinlogs() {
+			if fieldBinlog.FieldID != pkID {
+				continue
 			}
-			downloadPks(minioClient, bucketName, collectionID, pkID, segments)
 
-			return nil
-		},
+			for _, binlog := range fieldBinlog.Binlogs {
+				logPath := strings.Replace(binlog.LogPath, "ROOT_PATH", rootPath, -1)
+				obj, err := cli.GetObject(ctx, bucketName, logPath, minio.GetObjectOptions{})
+				if err != nil {
+					fmt.Println("failed to download file", bucketName, logPath)
+					return err
+				}
+
+				name := path.Base(logPath)
+
+				f, err := os.Create(path.Join(targetFolder, name))
+				if err != nil {
+					fmt.Println("failed to open file")
+					return err
+				}
+				w := bufio.NewWriter(f)
+				r := bufio.NewReader(obj)
+				_, err = io.Copy(w, r)
+				if err != nil {
+					fmt.Println(err.Error())
+				}
+				w.Flush()
+				f.Close()
+				count++
+			}
+		}
+		progress := (i + 1) * 100 / len(segments)
+		fmt.Fprintf(pd, pf, progress, i+1, len(segments))
 	}
 
-	cmd.Flags().Int64("id", 0, "collection id to download")
-	return cmd
+	fmt.Println()
+	fmt.Printf("pk file download completed for collection :%d, %d file(s) downloaded\n", collID, count)
+	return nil
 }
 
 func getMinioClient() (*minio.Client, error) {
@@ -166,56 +187,3 @@ func getMinioClient() (*minio.Client, error) {
 
 	return minioClient, nil
 }
-
-func downloadPks(cli *minio.Client, bucketName string, collID, pkID int64, segments []*datapb.SegmentInfo) {
-	err := os.Mkdir(fmt.Sprintf("%d", collID), 0o777)
-	if err != nil {
-		fmt.Println("Failed to create folder,", err.Error())
-	}
-
-	pd := uilive.New()
-	pf := "Downloading pk files ... %d%%(%d/%d)\n"
-	pd.Start()
-	fmt.Fprintf(pd, pf, 0, 0, len(segments))
-	defer pd.Stop()
-
-	count := 0
-	for i, segment := range segments {
-		for _, fieldBinlog := range segment.Binlogs {
-			if fieldBinlog.FieldID != pkID {
-				continue
-			}
-
-			folder := fmt.Sprintf("%d/%d", collID, segment.ID)
-			err := os.MkdirAll(folder, 0o777)
-			if err != nil {
-				fmt.Println("Failed to create sub-folder", err.Error())
-				return
-			}
-
-			for _, binlog := range fieldBinlog.Binlogs {
-				obj, err := cli.GetObject(context.Background(), bucketName, binlog.GetLogPath(), minio.GetObjectOptions{})
-				if err != nil {
-					fmt.Println("failed to download file", bucketName, binlog.GetLogPath())
-					return
-				}
-
-				name := path.Base(binlog.GetLogPath())
-
-				f, err := os.Create(path.Join(folder, name))
-				if err != nil {
-					fmt.Println("failed to open file")
-					return
-				}
-				w := bufio.NewWriter(f)
-				r := bufio.NewReader(obj)
-				io.Copy(w, r)
-				count++
-			}
-		}
-		progress := (i + 1) * 100 / len(segments)
-		fmt.Fprintf(pd, pf, progress, i+1, len(segments))
-	}
-	fmt.Println()
-	fmt.Printf("pk file download completed for collection :%d, %d file(s) downloaded\n", collID, count)
-}

states/download_segment.go

@@ -22,7 +22,7 @@ import (
 type DownloadSegmentParam struct {
 	framework.ParamBase `use:"download-segment" desc:"download segment file with provided segment id"`
 	MinioAddress        string `name:"minioAddr" default:"" desc:"the minio address to override, leave empty to use milvus.yaml value"`
-	SegmentID           int64  `name:"segment" default:"0" desc:"segment id to downloads"`
+	SegmentID           int64  `name:"segment" default:"0" desc:"segment id to download"`
 }
 
 func (s *InstanceState) DownloadSegmentCommand(ctx context.Context, p *DownloadSegmentParam) error {