From 616dfba575ba3159c541fa940144cd162315665b Mon Sep 17 00:00:00 2001
From: jayson <jayson@laboys.org>
Date: Tue, 16 Jan 2024 15:19:16 +0800
Subject: [PATCH] shell autocomplete for k8s flags

---
 cmd/root.go               | 65 ++++++++++++++++++++++++++++++++++++++-
 internal/client/client.go |  2 +-
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 2beadb4f79..ea4d15cc6c 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -4,11 +4,14 @@
 package cmd
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"runtime/debug"
+	"strings"
 
 	"github.com/derailed/k9s/internal/config/data"
+	"k8s.io/client-go/tools/clientcmd/api"
 
 	"github.com/derailed/k9s/internal/client"
 	"github.com/derailed/k9s/internal/color"
@@ -44,11 +47,19 @@ var (
 	out = colorable.NewColorableStdout()
 )
 
+type FlagError struct{ err error }
+
+func (e *FlagError) Error() string { return e.err.Error() }
+
 func init() {
 	if err := config.InitLogLoc(); err != nil {
 		fmt.Printf("Fail to init k9s logs location %s\n", err)
 	}
 
+	rootCmd.SetFlagErrorFunc(func(command *cobra.Command, err error) error {
+		return &FlagError{err: err}
+	})
+
 	rootCmd.AddCommand(versionCmd(), infoCmd())
 	initK9sFlags()
 	initK8sFlags()
@@ -57,7 +68,10 @@ func init() {
 // Execute root command.
 func Execute() {
 	if err := rootCmd.Execute(); err != nil {
-		panic(err)
+		var flagError *FlagError
+		if !errors.As(err, &flagError) {
+			panic(err)
+		}
 	}
 }
 
@@ -281,6 +295,7 @@ func initK8sFlags() {
 
 	initAsFlags()
 	initCertFlags()
+	initK8sFlagCompletion()
 }
 
 func initAsFlags() {
@@ -335,3 +350,51 @@ func initCertFlags() {
 		"Bearer token for authentication to the API server",
 	)
 }
+
+func initK8sFlagCompletion() {
+	_ = rootCmd.RegisterFlagCompletionFunc("context", k8sFlagCompletionFunc(func(cfg *api.Config) map[string]*api.Context {
+		return cfg.Contexts
+	}))
+
+	_ = rootCmd.RegisterFlagCompletionFunc("cluster", k8sFlagCompletionFunc(func(cfg *api.Config) map[string]*api.Cluster {
+		return cfg.Clusters
+	}))
+
+	_ = rootCmd.RegisterFlagCompletionFunc("user", k8sFlagCompletionFunc(func(cfg *api.Config) map[string]*api.AuthInfo {
+		return cfg.AuthInfos
+	}))
+
+	_ = rootCmd.RegisterFlagCompletionFunc("namespace", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		conn, err := client.InitConnection(client.NewConfig(k8sFlags))
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		nss, err := conn.ValidNamespaceNames()
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		return filterFlagCompletions(nss, toComplete)
+	})
+}
+
+func k8sFlagCompletionFunc[T any](picker func(cfg *api.Config) map[string]T) func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
+	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		k8sCfg, err := client.NewConfig(k8sFlags).RawConfig()
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+		return filterFlagCompletions(picker(&k8sCfg), toComplete)
+	}
+}
+
+func filterFlagCompletions[T any](m map[string]T, toComplete string) ([]string, cobra.ShellCompDirective) {
+	var completions []string
+	for name := range m {
+		if strings.HasPrefix(name, toComplete) {
+			completions = append(completions, name)
+		}
+	}
+	return completions, cobra.ShellCompDirectiveNoFileComp
+}
diff --git a/internal/client/client.go b/internal/client/client.go
index 02b1e6aa56..76dace3118 100644
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -71,7 +71,7 @@ func InitConnection(config *Config) (*APIClient, error) {
 	if err != nil {
 		log.Error().Err(err).Msgf("Fail to locate metrics-server")
 	}
-	if errors.Is(err, noMetricServerErr) || errors.Is(err, metricsUnsupportedErr) {
+	if err == nil || errors.Is(err, noMetricServerErr) || errors.Is(err, metricsUnsupportedErr) {
 		return &a, nil
 	}
 	a.connOK = false