feat: allow specify the context to be added (#962)

Signed-off-by: chengzw <[email protected]>
sunny0826 · Jun 1, 2024 · 13fbf30 · 13fbf30
1 parent 39fc351
commit 13fbf30
Show file tree

Hide file tree

Showing 9 changed files with 90 additions and 32 deletions.
diff --git a/cmd/add.go b/cmd/add.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/spf13/cobra"
+	"golang.org/x/exp/slices"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
 )
@@ -37,23 +38,25 @@ func (ac *AddCommand) Init() {
 		},
 		Example: addExample(),
 	}
+	ac.command.PersistentFlags().BoolP("cover", "c", false, "overwrite local kubeconfig files")
 	ac.command.Flags().StringP("file", "f", "", "path to merge kubeconfig files")
+	ac.command.Flags().StringSlice("context", []string{}, "specify the context to be added")
 	ac.command.Flags().String("context-prefix", "", "add a prefix before context name")
 	ac.command.Flags().String("context-name", "", "override context name when add kubeconfig context, when context-name is set, context-prefix and context-template parameters will be ignored")
-	ac.command.PersistentFlags().BoolP("cover", "c", false, "overwrite local kubeconfig files")
-	ac.command.Flags().Bool("select-context", false, "select the context to be added")
 	ac.command.Flags().StringSlice("context-template", []string{"context"}, "define the attributes used for composing the context name, available values: filename, user, cluster, context, namespace")
+	ac.command.Flags().Bool("select-context", false, "select the context to be added in interactive mode")
 	_ = ac.command.MarkFlagRequired("file")
 	ac.AddCommands(&DocsCommand{})
 }
 
 func (ac *AddCommand) runAdd(cmd *cobra.Command, args []string) error {
-	file, _ := ac.command.Flags().GetString("file")
 	cover, _ := ac.command.Flags().GetBool("cover")
+	file, _ := ac.command.Flags().GetString("file")
+	context, _ := ac.command.Flags().GetStringSlice("context")
 	contextPrefix, _ := ac.command.Flags().GetString("context-prefix")
 	contextName, _ := ac.command.Flags().GetString("context-name")
-	selectContext, _ := ac.command.Flags().GetBool("select-context")
 	contextTemplate, _ := ac.command.Flags().GetStringSlice("context-template")
+	selectContext, _ := ac.command.Flags().GetBool("select-context")
 
 	var newConfig *clientcmdapi.Config
 
@@ -88,15 +91,15 @@ func (ac *AddCommand) runAdd(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	err = AddToLocal(newConfig, file, contextPrefix, cover, selectContext, contextTemplate)
+	err = AddToLocal(newConfig, file, contextPrefix, cover, selectContext, contextTemplate, context)
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
 // AddToLocal add kubeConfig to local
-func AddToLocal(newConfig *clientcmdapi.Config, path, contextPrefix string, cover bool, selectContext bool, contextTemplate []string) error {
+func AddToLocal(newConfig *clientcmdapi.Config, path, contextPrefix string, cover bool, selectContext bool, contextTemplate []string, context []string) error {
 	oldConfig, err := clientcmd.LoadFromFile(cfgFile)
 	if err != nil {
 		return err
@@ -106,7 +109,7 @@ func AddToLocal(newConfig *clientcmdapi.Config, path, contextPrefix string, cove
 		fileName: getFileName(path),
 	}
 	// merge context loop
-	outConfig, err := kco.handleContexts(oldConfig, contextPrefix, selectContext, contextTemplate)
+	outConfig, err := kco.handleContexts(oldConfig, contextPrefix, selectContext, contextTemplate, context)
 	if err != nil {
 		return err
 	}
@@ -134,7 +137,7 @@ func AddToLocal(newConfig *clientcmdapi.Config, path, contextPrefix string, cove
 	return nil
 }
 
-func (kc *KubeConfigOption) handleContexts(oldConfig *clientcmdapi.Config, contextPrefix string, selectContext bool, contextTemplate []string) (*clientcmdapi.Config, error) {
+func (kc *KubeConfigOption) handleContexts(oldConfig *clientcmdapi.Config, contextPrefix string, selectContext bool, contextTemplate []string, context []string) (*clientcmdapi.Config, error) {
 	newConfig := clientcmdapi.NewConfig()
 	var newName string
 	generatedName := make(map[string]int)
@@ -152,7 +155,11 @@ func (kc *KubeConfigOption) handleContexts(oldConfig *clientcmdapi.Config, conte
 			newName = fmt.Sprintf("%s-%d", newName, generatedName[newName])
 		}
 
-		if selectContext {
+		if len(context) > 0 {
+			if !slices.Contains(context, newName) {
+				continue
+			}
+		} else if selectContext {
 			importContext := BoolUI(fmt.Sprintf("Do you want to add context「%s」? (If you select `False`, this context will not be merged)", newName))
 			if importContext == "False" {
 				continue
@@ -269,6 +276,8 @@ kubecm add -f test.yaml --context-template user,cluster --context-prefix demo
 kubecm add -f test.yaml --context-name test
 # Merge test.yaml with $HOME/.kube/config and select the context to be added in interactive mode
 kubecm add -f test.yaml --select-context
+# Merge test.yaml with $HOME/.kube/config and specify the context to be added
+kubecm add -f test.yaml --context context1,context2
 # Add kubeconfig from stdin
 cat /etc/kubernetes/admin.conf | kubecm add -f -
 `

diff --git a/cmd/add_test.go b/cmd/add_test.go
@@ -177,6 +177,35 @@ var (
 			"demo":    {AuthInfo: "single-user", Cluster: "single-cluster", Namespace: "single-ns"},
 		},
 	}
+
+	multiTestConfig = clientcmdapi.Config{
+		AuthInfos: map[string]*clientcmdapi.AuthInfo{
+			"blue-user":  {Token: "blue-token"},
+			"green-user": {Token: "green-token"}},
+		Clusters: map[string]*clientcmdapi.Cluster{
+			"cat-cluster": {Server: "http://cat.org:8080"},
+			"dog-cluster": {Server: "http://dog.org:8080"}},
+		Contexts: map[string]*clientcmdapi.Context{
+			"small": {AuthInfo: "blue-user", Cluster: "cat-cluster", Namespace: "cat-ns"},
+			"large": {AuthInfo: "green-user", Cluster: "dog-cluster", Namespace: "dog-ns"},
+		},
+	}
+
+	selectContextTestConfig = clientcmdapi.Config{
+		AuthInfos: map[string]*clientcmdapi.AuthInfo{
+			"blue-user":  {Token: "blue-token"},
+			"black-user": {Token: "black-token"},
+			"red-user":   {Token: "red-token"}},
+		Clusters: map[string]*clientcmdapi.Cluster{
+			"cat-cluster": {Server: "http://cat.org:8080"},
+			"pig-cluster": {Server: "http://pig.org:8080"},
+			"cow-cluster": {Server: "http://cow.org:8080"}},
+		Contexts: map[string]*clientcmdapi.Context{
+			"small":   {AuthInfo: "blue-user", Cluster: "cat-cluster", Namespace: "cat-ns"},
+			"root":    {AuthInfo: "black-user", Cluster: "pig-cluster", Namespace: "saw-ns"},
+			"federal": {AuthInfo: "red-user", Cluster: "cow-cluster", Namespace: "hammer-ns"},
+		},
+	}
 )
 
 func Test_checkContextName(t *testing.T) {
@@ -244,6 +273,7 @@ func TestKubeConfig_handleContexts(t *testing.T) {
 	}
 	type args struct {
 		oldConfig       *clientcmdapi.Config
+		context         []string
 		contextPrefix   string
 		contextTemplate []string
 	}
@@ -255,20 +285,21 @@ func TestKubeConfig_handleContexts(t *testing.T) {
 		wantErr bool
 	}{
 		// TODO: Add test cases.
-		{"not have new context name", fields{config: newConfig, fileName: "test"}, args{&oldTestConfig, "", []string{"context"}}, &mergedConfig, false},
-		{"single context name", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, "", []string{"context"}}, &mergeSingleTestConfig, false},
-		{"single context name - new", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, "rename", []string{"context"}}, &renameSingleTestConfig, false},
-		{"set context template", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, "", []string{"filename", "user", "cluster"}}, &contextTemplateTestConfig, false},
-		{"set context template and context prefix", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, "demo", []string{"user", "cluster"}}, &contextTemplateAndPrefixTestConfig, false},
-		{"set context name", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, "demo", []string{}}, &contextNameTestConfig, false},
+		{"not have new context name", fields{config: newConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "", []string{"context"}}, &mergedConfig, false},
+		{"single context name", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "", []string{"context"}}, &mergeSingleTestConfig, false},
+		{"single context name - new", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "rename", []string{"context"}}, &renameSingleTestConfig, false},
+		{"set context template", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "", []string{"filename", "user", "cluster"}}, &contextTemplateTestConfig, false},
+		{"set context template and context prefix", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "demo", []string{"user", "cluster"}}, &contextTemplateAndPrefixTestConfig, false},
+		{"set context name", fields{config: singleConfig, fileName: "test"}, args{&oldTestConfig, []string{}, "demo", []string{}}, &contextNameTestConfig, false},
+		{"select context", fields{config: &multiTestConfig, fileName: "test"}, args{&oldTestConfig, []string{"small"}, "", []string{"context"}}, &selectContextTestConfig, false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			kc := &KubeConfigOption{
 				config:   tt.fields.config,
 				fileName: tt.fields.fileName,
 			}
-			got, err := kc.handleContexts(tt.args.oldConfig, tt.args.contextPrefix, false, tt.args.contextTemplate)
+			got, err := kc.handleContexts(tt.args.oldConfig, tt.args.contextPrefix, false, tt.args.contextTemplate, tt.args.context)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("handleContexts() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -320,7 +351,7 @@ func TestAddToLocal(t *testing.T) {
 	}
 
 	// Test AddToLocal function
-	err = AddToLocal(newConfig, tempFile.Name(), "", true, false, []string{"context"})
+	err = AddToLocal(newConfig, tempFile.Name(), "", true, false, []string{"context"}, []string{})
 	if err != nil {
 		t.Fatalf("Failed to add to local: %v", err)
 	}

diff --git a/cmd/cloud_add.go b/cmd/cloud_add.go
@@ -36,6 +36,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 	clusterID, _ := ca.command.Flags().GetString("cluster_id")
 	regionID, _ := ca.command.Flags().GetString("region_id")
 	cover, _ := ca.command.Flags().GetBool("cover")
+	context, _ := ca.command.Flags().GetStringSlice("context")
 	selectContext, _ := ca.command.Flags().GetBool("select-context")
 	contextTemplate, _ := ca.command.Flags().GetStringSlice("context-template")
 	var num int
@@ -76,7 +77,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -89,7 +90,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, fmt.Sprintf("alicloud-%s", clusterID), "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, fmt.Sprintf("alicloud-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -129,7 +130,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -142,7 +143,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, fmt.Sprintf("tencent-%s", clusterID), "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, fmt.Sprintf("tencent-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -171,7 +172,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, clusters[clusterNum].Name, "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -184,7 +185,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			err = AddToLocal(newConfig, fmt.Sprintf("rancher-%s", clusterID), "", cover, selectContext, contextTemplate)
+			err = AddToLocal(newConfig, fmt.Sprintf("rancher-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 			if err != nil {
 				return err
 			}
@@ -221,7 +222,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 		if err != nil {
 			return err
 		}
-		err = AddToLocal(newConfig, fmt.Sprintf("aws-%s", clusterID), "", cover, selectContext, contextTemplate)
+		err = AddToLocal(newConfig, fmt.Sprintf("aws-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 		if err != nil {
 			return err
 		}
@@ -281,7 +282,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 			if err != nil {
 				return err
 			}
-			return AddToLocal(newConfig, fmt.Sprintf("azure-%s", clusterID), "", cover, selectContext, contextTemplate)
+			return AddToLocal(newConfig, fmt.Sprintf("azure-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 		}
 
 		subscriptionList, err := azure.ListSubscriptions()
@@ -334,7 +335,7 @@ func (ca *CloudAddCommand) runCloudAdd(cmd *cobra.Command, args []string) error
 		if err != nil {
 			return err
 		}
-		return AddToLocal(newConfig, fmt.Sprintf("azure-%s", clusterID), "", cover, selectContext, contextTemplate)
+		return AddToLocal(newConfig, fmt.Sprintf("azure-%s", clusterID), "", cover, selectContext, contextTemplate, context)
 
 	}
 	return nil

diff --git a/cmd/merge.go b/cmd/merge.go
@@ -30,16 +30,18 @@ func (mc *MergeCommand) Init() {
 	}
 	mc.command.Flags().StringP("folder", "f", "", "KubeConfig folder")
 	mc.command.Flags().BoolP("assumeyes", "y", false, "skip interactive file overwrite confirmation")
+	mc.command.Flags().StringSlice("context", []string{}, "specify the context to be merged")
 	mc.command.Flags().String("context-prefix", "", "add a prefix before context name")
-	mc.command.Flags().Bool("select-context", false, "select the context to be merged")
 	mc.command.Flags().StringSlice("context-template", []string{"context"}, "define the attributes used for composing the context name, available values: filename, user, cluster, context, namespace")
+	mc.command.Flags().Bool("select-context", false, "select the context to be merged in interactive mode")
 	//_ = mc.command.MarkFlagRequired("folder")
 	mc.AddCommands(&DocsCommand{})
 }
 
 func (mc MergeCommand) runMerge(command *cobra.Command, args []string) error {
 	files := args
 	folder, _ := mc.command.Flags().GetString("folder")
+	context, _ := mc.command.Flags().GetStringSlice("context")
 	contextPrefix, _ := mc.command.Flags().GetString("context-prefix")
 	selectContext, _ := mc.command.Flags().GetBool("select-context")
 	contextTemplate, _ := mc.command.Flags().GetStringSlice("context-template")
@@ -72,7 +74,7 @@ func (mc MergeCommand) runMerge(command *cobra.Command, args []string) error {
 			config:   loadConfig,
 			fileName: getFileName(yaml),
 		}
-		outConfigs, err = kco.handleContexts(outConfigs, contextPrefix, selectContext, contextTemplate)
+		outConfigs, err = kco.handleContexts(outConfigs, contextPrefix, selectContext, contextTemplate, context)
 		if err != nil {
 			return err
 		}
@@ -138,5 +140,7 @@ kubecm merge test.yaml --context-template user,cluster
 kubecm merge test.yaml --context-template user,cluster --context-prefix demo
 # Merge test.yaml with $HOME/.kube/config and select the context to be added in interactive mode
 kubecm merge test.yaml --select-context
+# Merge test.yaml with $HOME/.kube/config and specify the context to be added
+kubecm merge test.yaml --context context1,context2
 `
 }
diff --git a/docs/en-us/cli/kubecm_add.md b/docs/en-us/cli/kubecm_add.md
@@ -26,6 +26,8 @@ kubecm add -f test.yaml --context-template user,cluster --context-prefix demo
 kubecm add -f test.yaml --context-name test
 # Merge test.yaml with $HOME/.kube/config and select the context to be added in interactive mode
 kubecm add -f test.yaml --select-context
+# Merge test.yaml with $HOME/.kube/config and specify the context to be added
+kubecm add -f test.yaml --context context1,context2
 # Add kubeconfig from stdin
 cat /etc/kubernetes/admin.conf | kubecm add -f -
 
@@ -34,13 +36,14 @@ cat /etc/kubernetes/admin.conf | kubecm add -f -
 ### Options
 
 ```
+      --context strings            specify the context to be added
       --context-name string        override context name when add kubeconfig context, when context-name is set, context-prefix and context-template parameters will be ignored
       --context-prefix string      add a prefix before context name
       --context-template strings   define the attributes used for composing the context name, available values: filename, user, cluster, context, namespace (default [context])
   -c, --cover                      overwrite local kubeconfig files
   -f, --file string                path to merge kubeconfig files
   -h, --help                       help for add
-      --select-context             select the context to be added
+      --select-context             select the context to be added in interactive mode
 ```
 
 ### Options inherited from parent commands

diff --git a/docs/en-us/cli/kubecm_merge.md b/docs/en-us/cli/kubecm_merge.md
@@ -28,18 +28,21 @@ kubecm merge test.yaml --context-template user,cluster
 kubecm merge test.yaml --context-template user,cluster --context-prefix demo
 # Merge test.yaml with $HOME/.kube/config and select the context to be added in interactive mode
 kubecm merge test.yaml --select-context
+# Merge test.yaml with $HOME/.kube/config and specify the context to be added
+kubecm merge test.yaml --context context1,context2
 
 ```
 
 ### Options
 
 ```
   -y, --assumeyes                  skip interactive file overwrite confirmation
+      --context strings            specify the context to be merged
       --context-prefix string      add a prefix before context name
       --context-template strings   define the attributes used for composing the context name, available values: filename, user, cluster, context, namespace (default [context])
   -f, --folder string              KubeConfig folder
   -h, --help                       help for merge
-      --select-context             select the context to be merged
+      --select-context             select the context to be merged in interactive mode
 ```
 
 ### Options inherited from parent commands

diff --git a/docs/zh-cn/cli/kubecm_add.md b/docs/zh-cn/cli/kubecm_add.md
@@ -25,20 +25,23 @@ kubecm add -f test.yaml --context-template user,cluster --context-prefix demo
 kubecm add -f test.yaml --context-name test
 # Merge test.yaml with $HOME/.kube/config and select the context to be added in interactive mode
 kubecm add -f test.yaml --select-context
+# Merge test.yaml with $HOME/.kube/config and specify the context to be added
+kubecm add -f test.yaml --context context1,context2
 # Add kubeconfig from stdin
 cat /etc/kubernetes/admin.conf | kubecm add -f -
 ```
 
 ### 选项
 
 ```
+      --context strings            specify the context to be added
       --context-name string        override context name when add kubeconfig context, when context-name is set, context-prefix and context-template parameters will be ignored
       --context-prefix string      add a prefix before context name
       --context-template strings   define the attributes used for composing the context name, available values: filename, user, cluster, context, namespace (default [context])
   -c, --cover                      overwrite local kubeconfig files
   -f, --file string                path to merge kubeconfig files
   -h, --help                       help for add
-      --select-context             select the context to be added
+      --select-context             select the context to be added in interactive mode
 ```
 
 ### 全局选项