Add flag to set current-context in local KubeConfig (digitalocean#440)

* Add flag to set current-context in KubeConfig
jvasilevsky · May 15, 2019 · cf9ddde · cf9ddde
1 parent 2eafd47
commit cf9ddde
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 All changes to `doctl` will be documented in this file.
 
+## [1.18.0] - UNRELEASED
+
+ - #440 Add flag to set local KubeConfig's current-context. - @eddiezane
+
 ## [1.17.0] - 2019-05-08
 
  - #438 Remove need to opt-in to database commands. - @andrewsomething

diff --git a/args.go b/args.go
@@ -52,6 +52,8 @@ const (
 	ArgAutoUpgrade = "auto-upgrade"
 	// ArgCommandWait is a wait for a resource to be created argument.
 	ArgCommandWait = "wait"
+	// ArgSetCurrentContext is a flag to set the new kubeconfig context as current.
+	ArgSetCurrentContext = "set-current-context"
 	// ArgDropletID is a droplet id argument.
 	ArgDropletID = "droplet-id"
 	// ArgDropletIDs is a list of droplet IDs.

diff --git a/commands/kubernetes.go b/commands/kubernetes.go
@@ -124,13 +124,15 @@ format is in the form "name=your-name;size=size_slug;count=5;tag=tag1;tag=tag2"
 	- tag:    tags to apply to the node pool, repeat to add multiple tags at once.`)
 	AddBoolFlag(cmdKubeClusterCreate, doctl.ArgClusterUpdateKubeconfig, "", true, "whether to add the created cluster to your kubeconfig")
 	AddBoolFlag(cmdKubeClusterCreate, doctl.ArgCommandWait, "", true, "whether to wait for the created cluster to become running")
+	AddBoolFlag(cmdKubeClusterCreate, doctl.ArgSetCurrentContext, "", true, "whether to set the current kubectl context to that of the new cluster")
 	AddStringFlag(cmdKubeClusterCreate, doctl.ArgMaintenanceWindow, "", "any=00:00", "maintenance window to be set to the cluster. Syntax is in the format: 'day=HH:MM', where time is in UTC time zone. Day can be one of: ['any', 'monday', 'tuesday', 'wednesday', 'thursday', 'friday', 'saturday', 'sunday']")
 
 	cmdKubeClusterUpdate := CmdBuilder(cmd, RunKubernetesClusterUpdate, "update <id|name>", "update a cluster's properties", Writer, aliasOpt("u"))
 	AddStringFlag(cmdKubeClusterUpdate, doctl.ArgClusterName, "", "", "new cluster name")
 	AddStringSliceFlag(cmdKubeClusterUpdate, doctl.ArgTag, "", nil, "tags to apply to the cluster, repeat to add multiple tags at once")
 	AddBoolFlag(cmdKubeClusterUpdate, doctl.ArgAutoUpgrade, "", false, "whether to enable auto-upgrade for the cluster")
 	AddBoolFlag(cmdKubeClusterUpdate, doctl.ArgClusterUpdateKubeconfig, "", true, "whether to update the cluster in your kubeconfig")
+	AddBoolFlag(cmdKubeClusterUpdate, doctl.ArgSetCurrentContext, "", true, "whether to set the current kubectl context to that of the new cluster")
 	AddStringFlag(cmdKubeClusterUpdate, doctl.ArgMaintenanceWindow, "", "any=00:00", "maintenance window to be set to the cluster. Syntax is in the format: 'day=HH:MM', where time is in UTC time zone. Day can be one of: ['any', 'monday', 'tuesday', 'wednesday', 'thursday', 'friday', 'saturday', 'sunday']")
 
 	cmdKubeClusterUpgrade := CmdBuilder(cmd, RunKubernetesClusterUpgrade, "upgrade <id|name>", "upgrade a cluster to a new version", Writer)
@@ -158,7 +160,8 @@ func kubernetesKubeconfig() *Command {
 	CmdBuilder(cmd, RunKubernetesKubeconfigShow, "show <cluster-id|cluster-name>", "show a cluster's kubeconfig to standard out", Writer, aliasOpt("p", "g"))
 	cmdExecCredential := CmdBuilder(cmd, RunKubernetesKubeconfigExecCredential, "exec-credential <cluster-id>", "INTERNAL print a cluster's exec credential", Writer, hiddenCmd())
 	AddStringFlag(cmdExecCredential, doctl.ArgVersion, "", "", "")
-	CmdBuilder(cmd, RunKubernetesKubeconfigSave, "save <cluster-id|cluster-name>", "save a cluster's credentials to your local kubeconfig", Writer, aliasOpt("s"))
+	cmdSaveConfig := CmdBuilder(cmd, RunKubernetesKubeconfigSave, "save <cluster-id|cluster-name>", "save a cluster's credentials to your local kubeconfig", Writer, aliasOpt("s"))
+	AddBoolFlag(cmdSaveConfig, doctl.ArgSetCurrentContext, "", true, "whether to set the current kubectl context to that of the new cluster")
 	CmdBuilder(cmd, RunKubernetesKubeconfigRemove, "remove <cluster-id|cluster-name>", "remove a cluster's credentials from your local kubeconfig", Writer, aliasOpt("d", "rm"))
 	return cmd
 }
@@ -283,6 +286,10 @@ func RunKubernetesClusterCreate(defaultNodeSize string, defaultNodeCount int) fu
 		if err != nil {
 			return err
 		}
+		setCurrentContext, err := c.Doit.GetBool(c.NS, doctl.ArgSetCurrentContext)
+		if err != nil {
+			return err
+		}
 
 		kube := c.Kubernetes()
 
@@ -301,7 +308,7 @@ func RunKubernetesClusterCreate(defaultNodeSize string, defaultNodeCount int) fu
 
 		if update {
 			notice("cluster created, fetching credentials")
-			tryUpdateKubeconfig(kube, cluster.ID, clusterName)
+			tryUpdateKubeconfig(kube, cluster.ID, clusterName, setCurrentContext)
 		}
 
 		return displayClusters(c, true, *cluster)
@@ -317,6 +324,10 @@ func RunKubernetesClusterUpdate(c *CmdConfig) error {
 	if err != nil {
 		return err
 	}
+	setCurrentContext, err := c.Doit.GetBool(c.NS, doctl.ArgSetCurrentContext)
+	if err != nil {
+		return err
+	}
 	clusterIDorName := c.Args[0]
 	clusterID, err := clusterIDize(c.Kubernetes(), clusterIDorName)
 	if err != nil {
@@ -336,13 +347,13 @@ func RunKubernetesClusterUpdate(c *CmdConfig) error {
 
 	if update {
 		notice("cluster updated, fetching new credentials")
-		tryUpdateKubeconfig(kube, clusterID, clusterIDorName)
+		tryUpdateKubeconfig(kube, clusterID, clusterIDorName, setCurrentContext)
 	}
 
 	return displayClusters(c, true, *cluster)
 }
 
-func tryUpdateKubeconfig(kube do.KubernetesService, clusterID, clusterName string) {
+func tryUpdateKubeconfig(kube do.KubernetesService, clusterID, clusterName string, setCurrentContext bool) {
 	var (
 		kubeconfig []byte
 		err        error
@@ -362,7 +373,7 @@ func tryUpdateKubeconfig(kube do.KubernetesService, clusterID, clusterName strin
 			break
 		}
 	}
-	if err := writeOrAddToKubeconfig(clusterID, kubeconfig); err != nil {
+	if err := writeOrAddToKubeconfig(clusterID, kubeconfig, setCurrentContext); err != nil {
 		warn("couldn't write cluster credentials: %v", err)
 	}
 }
@@ -698,7 +709,12 @@ func RunKubernetesKubeconfigSave(c *CmdConfig) error {
 		return err
 	}
 
-	return writeOrAddToKubeconfig(clusterID, kubeconfig)
+	setCurrentContext, err := c.Doit.GetBool(c.NS, doctl.ArgSetCurrentContext)
+	if err != nil {
+		return err
+	}
+
+	return writeOrAddToKubeconfig(clusterID, kubeconfig, setCurrentContext)
 }
 
 // RunKubernetesKubeconfigRemove retrieves an existing kubernetes config and removes it from your local kubeconfig.
@@ -1118,7 +1134,7 @@ func buildNodePoolUpdateRequestFromArgs(c *CmdConfig, r *godo.KubernetesNodePool
 	return nil
 }
 
-func writeOrAddToKubeconfig(clusterID string, kubeconfig []byte) error {
+func writeOrAddToKubeconfig(clusterID string, kubeconfig []byte, setCurrentContext bool) error {
 	remote, err := clientcmd.Load(kubeconfig)
 	if err != nil {
 		return err
@@ -1130,7 +1146,7 @@ func writeOrAddToKubeconfig(clusterID string, kubeconfig []byte) error {
 	}
 
 	notice("adding cluster credentials to kubeconfig file found in %q", kubectlDefaults.GlobalFile)
-	if err := mergeKubeconfig(clusterID, remote, currentConfig); err != nil {
+	if err := mergeKubeconfig(clusterID, remote, currentConfig, setCurrentContext); err != nil {
 		return fmt.Errorf("couldn't use the kubeconfig info received, %v", err)
 	}
 	return clientcmd.ModifyConfig(kubectlDefaults, *currentConfig, false)
@@ -1156,7 +1172,7 @@ func removeFromKubeconfig(kubeconfig []byte) error {
 // mergeKubeconfig merges a remote cluster's config file with a local config file,
 // assuming that the current context in the remote config file points to the
 // cluster details to add to the local config.
-func mergeKubeconfig(clusterID string, remote, local *clientcmdapi.Config) error {
+func mergeKubeconfig(clusterID string, remote, local *clientcmdapi.Config, setCurrentContext bool) error {
 	remoteCtx, ok := remote.Contexts[remote.CurrentContext]
 	if !ok {
 		// this is a bug in the backend, we received incomplete/non-sensical data
@@ -1175,6 +1191,11 @@ func mergeKubeconfig(clusterID string, remote, local *clientcmdapi.Config) error
 	local.Contexts[remote.CurrentContext] = remoteCtx
 	local.Clusters[remoteCtx.Cluster] = remoteCluster
 
+	if setCurrentContext {
+		notice("setting current-context to %s", remote.CurrentContext)
+		local.CurrentContext = remote.CurrentContext
+	}
+
 	// configure kubectl to call doctl to retrieve credentials
 	local.AuthInfos[remoteCtx.AuthInfo] = &clientcmdapi.AuthInfo{
 		Exec: &clientcmdapi.ExecConfig{