support auto completion

cmd/completion/completion.go

@@ -0,0 +1,89 @@
+// Copyright 2022 Chaos Mesh Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package completion
+
+import (
+	"os"
+
+	"github.com/spf13/cobra"
+
+	"github.com/chaos-mesh/chaosd/pkg/utils"
+)
+
+// NewCompletionCommand returns the completion command
+func NewCompletionCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:   "completion [bash|zsh|fish|powershell]",
+		Short: "Generate completion script",
+		Long: `To load completions:
+Bash:
+
+$ source <(./bin/chaosd completion bash)
+
+# To load completions for each session, execute once:
+Linux:
+	$ ./bin/chaosd completion bash > /etc/bash_completion.d/chaosd
+MacOS:
+	$ ./bin/chaosd completion bash > /usr/local/etc/bash_completion.d/chaosd
+
+Zsh:
+
+$ compdef _chaosd ./bin/chaosd
+
+# If shell completion is not already enabled in your environment you will need
+# to enable it.  You can execute the following once:
+
+$ echo "autoload -U compinit; compinit" >> ~/.zshrc
+
+# To load completions for each session, execute once:
+$ ./bin/chaosd completion zsh > "${fpath[1]}/_chaosd"
+
+# You will need to start a new shell for this setup to take effect.
+
+Fish:
+
+$ ./bin/chaosd completion fish | source
+
+# To load completions for each session, execute once:
+$ ./bin/chaosd completion fish > ~/.config/fish/completions/chaosd.fish
+
+Powershell:
+
+PS> ./bin/chaosd completion powershell | Out-String | Invoke-Expression
+
+# To load completions for every new session, run:
+PS> ./bin/chaosd completion powershell > chaosd.ps1
+# and source this file from your powershell profile.
+	`,
+		DisableFlagsInUseLine: true,
+		ValidArgs:             []string{"bash", "zsh", "fish", "powershell"},
+		Args:                  cobra.ExactValidArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			switch args[0] {
+			case "bash":
+				if err := cmd.Root().GenBashCompletion(os.Stdout); err != nil {
+					utils.ExitWithError(utils.ExitError, err)
+				}
+			case "zsh":
+				if err := cmd.Root().GenZshCompletion(os.Stdout); err != nil {
+					utils.ExitWithError(utils.ExitError, err)
+				}
+			case "fish":
+				if err := cmd.Root().GenFishCompletion(os.Stdout, true); err != nil {
+					utils.ExitWithError(utils.ExitError, err)
+				}
+			}
+		},
+	}
+}

cmd/main.go

@@ -25,6 +25,7 @@ import (
 	ctrlzap "sigs.k8s.io/controller-runtime/pkg/log/zap"
 
 	"github.com/chaos-mesh/chaosd/cmd/attack"
+	"github.com/chaos-mesh/chaosd/cmd/completion"
 	"github.com/chaos-mesh/chaosd/cmd/recover"
 	"github.com/chaos-mesh/chaosd/cmd/search"
 	"github.com/chaos-mesh/chaosd/cmd/server"
@@ -50,6 +51,7 @@ func init() {
 		recover.NewRecoverCommand(),
 		search.NewSearchCommand(),
 		version.NewVersionCommand(),
+		completion.NewCompletionCommand(),
 	)
 
 	_ = utils.SetRuntimeEnv()

cmd/recover/completion.go

@@ -0,0 +1,46 @@
+// Copyright 2021 Chaos Mesh Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package recover
+
+import (
+	"github.com/pkg/errors"
+
+	"github.com/chaos-mesh/chaosd/pkg/core"
+	"github.com/chaos-mesh/chaosd/pkg/server/chaosd"
+)
+
+type completionContext struct {
+	uids chan string
+	err  chan error
+}
+
+func newCompletionCtx() *completionContext {
+	return &completionContext{
+		uids: make(chan string),
+		err:  make(chan error),
+	}
+}
+
+func listUid(ctx *completionContext, chaos *chaosd.Server) {
+	exps, err := chaos.Search(&core.SearchCommand{Status: core.Success})
+	if err != nil {
+		ctx.err <- errors.Wrap(err, "list exp")
+		return
+	}
+
+	for _, exp := range exps {
+		ctx.uids <- exp.Uid
+	}
+	close(ctx.uids)
+}

cmd/recover/recover.go

@@ -14,11 +14,15 @@
 package recover
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/spf13/cobra"
 	"go.uber.org/fx"
 
+	"github.com/pingcap/errors"
+	"github.com/pingcap/log"
+
 	"github.com/chaos-mesh/chaosd/cmd/server"
 	"github.com/chaos-mesh/chaosd/pkg/server/chaosd"
 	"github.com/chaos-mesh/chaosd/pkg/utils"
@@ -38,9 +42,10 @@ func NewRecoverCommand() *cobra.Command {
 	)
 
 	cmd := &cobra.Command{
-		Use:   "recover UID",
-		Short: "Recover a chaos experiment",
-		Args:  cobra.MinimumNArgs(1),
+		Use:               "recover UID",
+		Short:             "Recover a chaos experiment",
+		Args:              cobra.MinimumNArgs(1),
+		ValidArgsFunction: completeUid,
 		Run: func(cmd *cobra.Command, args []string) {
 			if len(args) == 0 {
 				utils.ExitWithMsg(utils.ExitBadArgs, "UID is required")
@@ -60,3 +65,33 @@ func recoverCommandF(chaos *chaosd.Server, options *recoverCommand) {
 
 	utils.NormalExit(fmt.Sprintf("Recover %s successfully", options.uid))
 }
+
+func completeUid(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	completionCtx := newCompletionCtx()
+	completionDep := fx.Options(
+		server.Module,
+		fx.Provide(func() *completionContext {
+			return completionCtx
+		}),
+	)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		if err := utils.FxNewAppWithoutLog(completionDep, fx.Invoke(listUid)).Start(ctx); err != nil {
+			log.Error(errors.Wrap(err, "start application").Error())
+		}
+	}()
+	var uids []string
+	for {
+		select {
+		case uid := <-completionCtx.uids:
+			if len(uid) == 0 {
+				return uids, cobra.ShellCompDirectiveNoFileComp
+			}
+			uids = append(uids, uid)
+		case err := <-completionCtx.err:
+			log.Error(err.Error())
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+	}
+}

go.mod

@@ -6,12 +6,17 @@ require (
 	github.com/chaos-mesh/chaos-mesh v0.9.1-0.20210525104133-41e37dd1ac16
 	github.com/chaos-mesh/chaos-mesh/api/v1alpha1 v0.0.0
 	github.com/containerd/containerd v1.2.3
+	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
 	github.com/docker/docker v1.4.2-0.20180625184442-8e610b2b55bf
+	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/gin-gonic/gin v1.6.3
 	github.com/go-logr/zapr v0.1.0
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/uuid v1.1.2
 	github.com/hashicorp/go-multierror v1.1.0
 	github.com/joomcode/errorx v1.0.1
+	github.com/json-iterator/go v1.1.11 // indirect
 	github.com/mitchellh/go-ps v0.0.0-20170309133038-4fdf99ab2936
 	github.com/olekukonko/tablewriter v0.0.4
 	github.com/onsi/gomega v1.9.0
@@ -21,18 +26,20 @@ require (
 	github.com/pkg/errors v0.9.1
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/shirou/gopsutil v0.0.0-20180427012116-c95755e4bcd7
-	github.com/spf13/cobra v1.1.1
+	github.com/spf13/cobra v1.4.0
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.7.0
 	github.com/swaggo/files v0.0.0-20190704085106-630677cd5c14
 	github.com/swaggo/gin-swagger v1.2.0
 	github.com/swaggo/swag v1.6.7
 	go.uber.org/fx v1.13.1
-	go.uber.org/zap v1.15.0
-	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 // indirect
+	go.uber.org/zap v1.17.0
+	golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 // indirect
+	golang.org/x/oauth2 v0.0.0-20210402161424-2e8d93401602 // indirect
 	golang.org/x/tools v0.1.1 // indirect
-	google.golang.org/grpc v1.29.1
-	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/genproto v0.0.0-20210402141018-6c239bbf2bb1 // indirect
+	google.golang.org/grpc v1.38.0
 	gorm.io/driver/sqlite v1.1.4
 	gorm.io/gorm v1.20.7
 	k8s.io/api v0.18.2