Add bash completion for necoperf-cli (#13)

* Add bash completion for necoperf-cli Signed-off-by: zeroalphat <taichi-takemura@cybozu.co.jp> --------- Signed-off-by: zeroalphat <taichi-takemura@cybozu.co.jp>
cybozu-go · Oct 25, 2023 · a9bd043 · a9bd043
1 parent b2b9d16
commit a9bd043
Show file tree

Hide file tree

Showing 2 changed files with 119 additions and 4 deletions.
diff --git a/cmd/necoperf-cli/cmd/completion.go b/cmd/necoperf-cli/cmd/completion.go
@@ -0,0 +1,112 @@
+package cmd
+
+import (
+	"context"
+	"strings"
+
+	"github.com/spf13/cobra"
+	corev1 "k8s.io/api/core/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	k8sConfig "sigs.k8s.io/controller-runtime/pkg/client/config"
+)
+
+func validArgsCompletionFunc(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	if len(args) != 0 {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	cfg, err := k8sConfig.GetConfig()
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	k8sClient, err := client.New(cfg, client.Options{})
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	pods := &corev1.PodList{}
+	err = k8sClient.List(context.Background(), pods, client.InNamespace(config.namespace))
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	var podNames []string
+	for _, pod := range pods.Items {
+		if !strings.HasPrefix(pod.Name, toComplete) {
+			continue
+		}
+		podNames = append(podNames, pod.Name)
+	}
+
+	return podNames, cobra.ShellCompDirectiveNoFileComp
+}
+
+func containerCompletionFunc(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	if len(args) == 0 {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	cfg, err := k8sConfig.GetConfig()
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+	k8sClient, err := client.New(cfg, client.Options{})
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	pod := &corev1.Pod{}
+	err = k8sClient.Get(context.Background(), client.ObjectKey{
+		Namespace: config.namespace,
+		Name:      args[0],
+	}, pod)
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	var containerNames []string
+	for _, c := range pod.Spec.Containers {
+		if !strings.HasPrefix(c.Name, toComplete) {
+			continue
+		}
+		containerNames = append(containerNames, c.Name)
+	}
+
+	return containerNames, cobra.ShellCompDirectiveNoFileComp
+}
+
+func namespaceCompletionFunc(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	cfg, err := k8sConfig.GetConfig()
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+	k8sClient, err := client.New(cfg, client.Options{})
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	ns := &corev1.NamespaceList{}
+	err = k8sClient.List(context.Background(), ns)
+	if err != nil {
+		cobra.CompError(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	var namespaces []string
+	for _, n := range ns.Items {
+		if !strings.HasPrefix(n.Name, toComplete) {
+			continue
+		}
+		namespaces = append(namespaces, n.Name)
+	}
+
+	return namespaces, cobra.ShellCompDirectiveNoFileComp
+}
diff --git a/cmd/necoperf-cli/cmd/profile.go b/cmd/necoperf-cli/cmd/profile.go
@@ -34,10 +34,11 @@ func isPodReady(pod *corev1.Pod) bool {
 
 func NewProfileCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "profile",
-		Short: "Perform CPU profiling on the target container",
-		Long:  "Perform CPU profiling on the target container",
-		Args:  cobra.ExactArgs(1),
+		Use:               "profile",
+		Short:             "Perform CPU profiling on the target container",
+		Long:              "Perform CPU profiling on the target container",
+		Args:              cobra.ExactArgs(1),
+		ValidArgsFunction: validArgsCompletionFunc,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			cmd.SilenceUsage = true
 			config.podName = args[0]
@@ -95,6 +96,8 @@ func NewProfileCommand() *cobra.Command {
 	cmd.Flags().StringVarP(&config.containerName, "container", "c", "", "Specify the container name to profile")
 	cmd.Flags().DurationVar(&config.timeout, "timeout", 30*time.Second, "Time to run cpu profiling on server")
 	cmd.Flags().StringVar(&config.outputDir, "output-dir", "/tmp", "Directory to output profiling result")
+	cmd.RegisterFlagCompletionFunc("namespace", namespaceCompletionFunc)
+	cmd.RegisterFlagCompletionFunc("container", containerCompletionFunc)
 
 	return cmd
 }