From 4cbe6751b0db8b35adb0e3d482d005b05cc91d67 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Fri, 18 Aug 2023 14:48:13 +0800
Subject: [PATCH] Fix embedEtcdState command missing (#190)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 .gitignore                            |  4 ++++
 states/app_state.go                   |  1 +
 states/backup_mock_connect.go         | 11 ++++-------
 states/etcd/show/collection_loaded.go |  8 +++++++-
 states/etcd/show/index.go             |  5 ++++-
 5 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index f28fec4..98350b3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,6 +10,7 @@ birdwatcher-build
 
 # logs
 audit.log
+audit_*.log
 bw_debug.log
 
 # vim
@@ -17,3 +18,6 @@ bw_debug.log
 
 # backup file
 *.bak.gz
+*.tar.gz
+backups/
+
diff --git a/states/app_state.go b/states/app_state.go
index 163590e..7ac6c5f 100644
--- a/states/app_state.go
+++ b/states/app_state.go
@@ -48,6 +48,7 @@ func (app *ApplicationState) Process(cmd string) (framework.State, error) {
 	for key, state := range app.states {
 		next := state.NextState()
 		if next != nil {
+			app.config.Log("[debug] set next", key, next.Label())
 			state.SetNext(nil)
 			app.states[key] = next
 		}
diff --git a/states/backup_mock_connect.go b/states/backup_mock_connect.go
index ed2e165..d762ca2 100644
--- a/states/backup_mock_connect.go
+++ b/states/backup_mock_connect.go
@@ -57,11 +57,11 @@ func (s *embedEtcdMockState) Close() {
 // SetupCommands setups the command.
 // also called after each command run to reset flag values.
 func (s *embedEtcdMockState) SetupCommands() {
-	cmd := &cobra.Command{}
+	cmd := s.GetCmd()
 
 	rootPath := path.Join(s.instanceName, metaPath)
 
-	cmd.AddCommand(
+	s.MergeCobraCommands(cmd,
 		// show [subcommand] options...
 		etcd.ShowCommand(s.client, rootPath),
 
@@ -83,12 +83,9 @@ func (s *embedEtcdMockState) SetupCommands() {
 
 		getListMetricsNodeCmd(s),
 	)
-	cmd.AddCommand(etcd.RawCommands(s.client)...)
+	s.MergeCobraCommands(cmd, etcd.RawCommands(s.client)...)
 
-	s.MergeFunctionCommands(cmd, s)
-
-	s.CmdState.RootCmd = cmd
-	s.SetupFn = s.SetupCommands
+	s.UpdateState(cmd, s, s.SetupCommands)
 }
 
 func (s *embedEtcdMockState) SetInstance(instanceName string) {
diff --git a/states/etcd/show/collection_loaded.go b/states/etcd/show/collection_loaded.go
index fab614d..64715df 100644
--- a/states/etcd/show/collection_loaded.go
+++ b/states/etcd/show/collection_loaded.go
@@ -10,6 +10,7 @@ import (
 	"github.com/milvus-io/birdwatcher/models"
 	"github.com/milvus-io/birdwatcher/states/etcd/common"
 	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
+	"github.com/samber/lo"
 )
 
 const (
@@ -18,7 +19,7 @@ const (
 
 type CollectionLoadedParam struct {
 	framework.ParamBase `use:"show collection-loaded" desc:"display information of loaded collection from querycoord" alias:"collection-load"`
-	//CollectionID int64 `name:""`
+	CollectionID        int64 `name:"collection" default:"0" desc:"collection id to check"`
 }
 
 // CollectionLoadedCommand return show collection-loaded command.
@@ -31,6 +32,11 @@ func (c *ComponentShow) CollectionLoadedCommand(ctx context.Context, p *Collecti
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to list collection load info")
 	}
+	if p.CollectionID > 0 {
+		infos = lo.Filter(infos, func(info *models.CollectionLoaded, _ int) bool {
+			return info.CollectionID == p.CollectionID
+		})
+	}
 
 	return framework.NewListResult[CollectionsLoaded](infos), nil
 }
diff --git a/states/etcd/show/index.go b/states/etcd/show/index.go
index 00a3bcb..5519400 100644
--- a/states/etcd/show/index.go
+++ b/states/etcd/show/index.go
@@ -14,6 +14,7 @@ import (
 
 type IndexParam struct {
 	framework.ParamBase `use:"show index" desc:"" alias:"indexes"`
+	CollectionID        int64 `name:"collection" default:"0" desc:"collection id to list index on"`
 }
 
 // IndexCommand returns show index command.
@@ -39,7 +40,9 @@ func (c *ComponentShow) IndexCommand(ctx context.Context, p *IndexParam) {
 	}
 
 	for _, index := range fieldIndexes {
-		printIndexV2(index)
+		if p.CollectionID == 0 || p.CollectionID == index.IndexInfo.GetCollectionID() {
+			printIndexV2(index)
+		}
 	}
 }