enhance: Support set collection consistency-level command (#323)

Add command to update collection default consistency level

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

states/etcd/common/collection.go

@@ -3,13 +3,13 @@ package common
 import (
 	"bytes"
 	"context"
-	"errors"
 	"fmt"
 	"path"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/cockroachdb/errors"
 	"github.com/golang/protobuf/proto"
 	"github.com/samber/lo"
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -232,37 +232,45 @@ func FillFieldSchemaIfEmptyV2(cli clientv3.KV, basePath string, collection *etcd
 	return nil
 }
 
-func UpdateCollection(ctx context.Context, cli clientv3.KV, basePath string, collectionID int64, fn func(coll *etcdpbv2.CollectionInfo), dryRun bool) error {
-	prefix := path.Join(basePath, CollectionMetaPrefix, strconv.FormatInt(collectionID, 10))
-	resp, err := cli.Get(ctx, prefix)
+func UpdateCollection(ctx context.Context, cli clientv3.KV, key string, fn func(coll *etcdpbv2.CollectionInfo), run bool) error {
+	resp, err := cli.Get(ctx, key)
 	if err != nil {
 		return err
 	}
 
+	if len(resp.Kvs) != 1 {
+		return errors.Newf("collection with key[%s] not found", key)
+	}
+
 	info := &etcdpbv2.CollectionInfo{}
 	err = proto.Unmarshal(resp.Kvs[0].Value, info)
 	if err != nil {
 		return err
 	}
 
-	fn(info)
-
-	bs, err := proto.Marshal(info)
+	clone := proto.Clone(info).(*etcdpbv2.CollectionInfo)
+	fn(clone)
+	bs, err := proto.Marshal(clone)
 	if err != nil {
 		return err
 	}
 
-	if dryRun {
-		fmt.Println("dry run")
-		fmt.Println("before alter")
-		fmt.Printf("schema:%s", info.Schema.String())
-		fmt.Println("after alter")
-		fmt.Printf("schema:%s", info.Schema.String())
+	fmt.Println("======dry run======")
+	fmt.Println("before alter")
+	fmt.Printf("schema:%s\n", info.String())
+	fmt.Println()
+	fmt.Println("after alter")
+	fmt.Printf("schema:%s\n", clone.String())
+	if !run {
 		return nil
 	}
 
-	_, err = cli.Put(ctx, prefix, string(bs))
-	return err
+	_, err = cli.Put(ctx, key, string(bs))
+	if err != nil {
+		return err
+	}
+	fmt.Println("Update collection done!")
+	return nil
 }
 
 func UpdateField(ctx context.Context, cli clientv3.KV, basePath string, collectionID, fieldID int64, fn func(field *schemapbv2.FieldSchema), dryRun bool) error {

states/etcd/restore/collection.go

@@ -48,11 +48,7 @@ func CollectionCommand(cli clientv3.KV, basePath string) *cobra.Command {
 				collection.State = etcdpbv2.CollectionState_CollectionCreated
 			}
 
-			if collection.DBID > 0 {
-				// err = common.UpdateCollectionWithDB(ctx, cli, basePath, collectionID, collection.DBID, updateCollState)
-			} else {
-				err = common.UpdateCollection(ctx, cli, basePath, collectionID, updateCollState, false)
-			}
+			err = common.UpdateCollection(ctx, cli, collection.Key(), updateCollState, true)
 
 			if err != nil {
 				fmt.Println(err.Error())

states/etcd/set/collection_consistency_level.go

@@ -0,0 +1,44 @@
+package set
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cockroachdb/errors"
+
+	"github.com/milvus-io/birdwatcher/framework"
+	"github.com/milvus-io/birdwatcher/models"
+	commonpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/commonpb"
+	etcdpbv2 "github.com/milvus-io/birdwatcher/proto/v2.2/etcdpb"
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
+)
+
+type CollectionConsistencyLevelParam struct {
+	framework.ParamBase `use:"set collection consistency-level" desc:"set collection default consistency level"`
+	CollectionID        int64  `name:"collection" default:"0" desc:"collection id to update"`
+	ConsistencyLevel    string `name:"consistency-level" default:"" desc:"Consistency Level to set"`
+	Run                 bool   `name:"run" default:"false"`
+}
+
+func (c *ComponentSet) CollectionConsistencyLevelCommand(ctx context.Context, p *CollectionConsistencyLevelParam) error {
+	levelVal, ok := commonpbv2.ConsistencyLevel_value[p.ConsistencyLevel]
+	if !ok {
+		return errors.Newf(`consistency level string "%s" is not valid`, p.ConsistencyLevel)
+	}
+
+	consistencyLevel := commonpbv2.ConsistencyLevel(levelVal)
+	collection, err := common.GetCollectionByIDVersion(ctx, c.client, c.basePath, etcdversion.GetVersion(), p.CollectionID)
+	if err != nil {
+		return err
+	}
+
+	if collection.ConsistencyLevel == models.ConsistencyLevel(consistencyLevel) {
+		fmt.Printf("collection consistency level is already %s\n", p.ConsistencyLevel)
+		return nil
+	}
+
+	return common.UpdateCollection(ctx, c.client, collection.Key(), func(coll *etcdpbv2.CollectionInfo) {
+		coll.ConsistencyLevel = consistencyLevel
+	}, p.Run)
+}

states/etcd/set/component.go

@@ -0,0 +1,21 @@
+package set
+
+import (
+	clientv3 "go.etcd.io/etcd/client/v3"
+
+	"github.com/milvus-io/birdwatcher/configs"
+)
+
+type ComponentSet struct {
+	client   clientv3.KV
+	config   *configs.Config
+	basePath string
+}
+
+func NewComponent(cli clientv3.KV, config *configs.Config, basePath string) *ComponentSet {
+	return &ComponentSet{
+		client:   cli,
+		config:   config,
+		basePath: basePath,
+	}
+}

states/instance.go

@@ -14,6 +14,7 @@ import (
 	"github.com/milvus-io/birdwatcher/states/etcd/audit"
 	"github.com/milvus-io/birdwatcher/states/etcd/remove"
 	"github.com/milvus-io/birdwatcher/states/etcd/repair"
+	"github.com/milvus-io/birdwatcher/states/etcd/set"
 	"github.com/milvus-io/birdwatcher/states/etcd/show"
 )
 
@@ -23,6 +24,7 @@ type InstanceState struct {
 	*show.ComponentShow
 	*remove.ComponentRemove
 	*repair.ComponentRepair
+	*set.ComponentSet
 	instanceName string
 	client       clientv3.KV
 	auditFile    *os.File
@@ -141,6 +143,7 @@ func getInstanceState(cli clientv3.KV, instanceName, metaPath string, etcdState
 		ComponentShow:   show.NewComponent(cli, config, basePath),
 		ComponentRemove: remove.NewComponent(cli, config, basePath),
 		ComponentRepair: repair.NewComponent(cli, config, basePath),
+		ComponentSet:    set.NewComponent(cli, config, basePath),
 		instanceName:    instanceName,
 		client:          kv,
 		auditFile:       file,