Add show partition command (#152)

Signed-off-by: Congqi Xia <[email protected]>

models/partition.go

@@ -0,0 +1,52 @@
+package models
+
+import (
+	etcdpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/etcdpb"
+)
+
+type Partition struct {
+	ID           int64
+	CollectionID int64
+	Name         string
+	State        PartitionState
+
+	key string
+}
+
+func NewPartition(info *etcdpbv2.PartitionInfo, key string) *Partition {
+	p := &Partition{
+		ID:           info.GetPartitionID(),
+		CollectionID: info.GetCollectionId(),
+		Name:         info.GetPartitionName(),
+		State:        PartitionState(info.GetState()),
+		key:          key,
+	}
+	return p
+}
+
+type PartitionState int32
+
+const (
+	PartitionStatePartitionCreated  PartitionState = 0
+	PartitionStatePartitionCreating PartitionState = 1
+	PartitionStatePartitionDropping PartitionState = 2
+	PartitionStatePartitionDropped  PartitionState = 3
+)
+
+var PartitionStatename = map[int32]string{
+	0: "PartitionCreated",
+	1: "PartitionCreating",
+	2: "PartitionDropping",
+	3: "PartitionDropped",
+}
+
+var PartitionStatevalue = map[string]int32{
+	"PartitionCreated":  0,
+	"PartitionCreating": 1,
+	"PartitionDropping": 2,
+	"PartitionDropped":  3,
+}
+
+func (x PartitionState) String() string {
+	return EnumName(PartitionStatename, int32(x))
+}

states/etcd/commands.go

@@ -38,6 +38,8 @@ func ShowCommand(cli clientv3.KV, basePath string) *cobra.Command {
 		show.IndexCommand(cli, basePath),
 		// show segment-index
 		show.SegmentIndexCommand(cli, basePath),
+		// show partition
+		show.PartitionCommand(cli, basePath),
 
 		// show replica
 		show.ReplicaCommand(cli, basePath),

states/etcd/common/partition.go

@@ -0,0 +1,31 @@
+package common
+
+import (
+	"context"
+	"fmt"
+	"path"
+
+	"github.com/milvus-io/birdwatcher/models"
+	etcdpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/etcdpb"
+	"github.com/samber/lo"
+	clientv3 "go.etcd.io/etcd/client/v3"
+)
+
+const (
+	PartitionPrefix = `root-coord/partitions/`
+)
+
+// ListCollectionPartitions returns partition list of collection.
+func ListCollectionPartitions(ctx context.Context, cli clientv3.KV, basePath string, collectionID int64) ([]*models.Partition, error) {
+	prefix := path.Join(basePath, PartitionPrefix, fmt.Sprintf("%d", collectionID))
+
+	infos, keys, err := ListProtoObjects[etcdpbv2.PartitionInfo](ctx, cli, prefix)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return lo.Map(infos, func(info etcdpbv2.PartitionInfo, idx int) *models.Partition {
+		return models.NewPartition(&info, keys[idx])
+	}), nil
+}

states/etcd/show/partition.go

@@ -0,0 +1,47 @@
+package show
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	"github.com/spf13/cobra"
+	clientv3 "go.etcd.io/etcd/client/v3"
+)
+
+func PartitionCommand(cli clientv3.KV, basePath string) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "partition",
+		Short: "list partitions of provided collection",
+		Run: func(cmd *cobra.Command, args []string) {
+			collectionID, err := cmd.Flags().GetInt64("collection")
+			if err != nil {
+				fmt.Println(err.Error())
+				return
+			}
+
+			if collectionID == 0 {
+				fmt.Println("please provided collection id")
+				return
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			partitions, err := common.ListCollectionPartitions(ctx, cli, basePath, collectionID)
+			if err != nil {
+				fmt.Println("failed to list partition info", err.Error())
+			}
+
+			if len(partitions) == 0 {
+				fmt.Printf("no partition found for collection %d\n", collectionID)
+			}
+
+			for _, partition := range partitions {
+				fmt.Printf("Parition ID: %d\tName: %s\tState: %s\n", partition.ID, partition.Name, partition.State.String())
+			}
+		},
+	}
+
+	cmd.Flags().Int64("collection", 0, "collection id to list")
+	return cmd
+}