Revert 'Revert "add --node-label flag for node create command (#584, @…

…developer-guy, @EJose, @Dentrax)"' This reverts commit 7087264.
k3d-io · Jul 6, 2021 · 774af31 · 774af31
1 parent a2305bd
commit 774af31
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 21 deletions.
diff --git a/cmd/node/nodeCreate.go b/cmd/node/nodeCreate.go
@@ -23,6 +23,7 @@ package node
 
 import (
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -73,6 +74,8 @@ func NewCmdNodeCreate() *cobra.Command {
 	cmd.Flags().BoolVar(&createNodeOpts.Wait, "wait", false, "Wait for the node(s) to be ready before returning.")
 	cmd.Flags().DurationVar(&createNodeOpts.Timeout, "timeout", 0*time.Second, "Maximum waiting time for '--wait' before canceling/returning.")
 
+	cmd.Flags().StringSliceP("k3s-node-label", "", []string{}, "Specify k3s node labels in format \"foo=bar\"")
+
 	// done
 	return cmd
 }
@@ -124,6 +127,21 @@ func parseCreateNodeCmd(cmd *cobra.Command, args []string) ([]*k3d.Node, *k3d.Cl
 		log.Errorf("Provided memory limit value is invalid")
 	}
 
+	k3sNodeLabelsFlag, err := cmd.Flags().GetStringSlice("k3s-node-label")
+	if err != nil {
+		log.Errorln("No node-label specified")
+		log.Fatalln(err)
+	}
+
+	k3sNodeLabels := make(map[string]string, len(k3sNodeLabelsFlag))
+	for _, label := range k3sNodeLabelsFlag {
+		labelSplitted := strings.Split(label, "=")
+		if len(labelSplitted) != 2 {
+			log.Fatalf("unknown label format format: %s, use format \"foo=bar\"", label)
+		}
+		k3sNodeLabels[labelSplitted[0]] = labelSplitted[1]
+	}
+
 	// generate list of nodes
 	nodes := []*k3d.Node{}
 	for i := 0; i < replicas; i++ {
@@ -134,8 +152,9 @@ func parseCreateNodeCmd(cmd *cobra.Command, args []string) ([]*k3d.Node, *k3d.Cl
 			Labels: map[string]string{
 				k3d.LabelRole: roleStr,
 			},
-			Restart: true,
-			Memory:  memory,
+			K3sNodeLabels: k3sNodeLabels,
+			Restart:       true,
+			Memory:        memory,
 		}
 		nodes = append(nodes, node)
 	}

diff --git a/pkg/client/node.go b/pkg/client/node.go
@@ -492,6 +492,11 @@ func patchAgentSpec(node *k3d.Node) error {
 	if node.Cmd == nil {
 		node.Cmd = []string{"agent"}
 	}
+
+	for k, v := range node.K3sNodeLabels {
+		node.Args = append(node.Args, "--node-label", fmt.Sprintf("%s=%s", k, v))
+	}
+
 	return nil
 }
 

diff --git a/pkg/types/types.go b/pkg/types/types.go
@@ -330,25 +330,26 @@ type NodeIP struct {
 
 // Node describes a k3d node
 type Node struct {
-	Name       string            `yaml:"name" json:"name,omitempty"`
-	Role       Role              `yaml:"role" json:"role,omitempty"`
-	Image      string            `yaml:"image" json:"image,omitempty"`
-	Volumes    []string          `yaml:"volumes" json:"volumes,omitempty"`
-	Env        []string          `yaml:"env" json:"env,omitempty"`
-	Cmd        []string          // filled automatically based on role
-	Args       []string          `yaml:"extraArgs" json:"extraArgs,omitempty"`
-	Ports      nat.PortMap       `yaml:"portMappings" json:"portMappings,omitempty"`
-	Restart    bool              `yaml:"restart" json:"restart,omitempty"`
-	Created    string            `yaml:"created" json:"created,omitempty"`
-	Labels     map[string]string // filled automatically
-	Networks   []string          // filled automatically
-	ExtraHosts []string          // filled automatically
-	ServerOpts ServerOpts        `yaml:"serverOpts" json:"serverOpts,omitempty"`
-	AgentOpts  AgentOpts         `yaml:"agentOpts" json:"agentOpts,omitempty"`
-	GPURequest string            // filled automatically
-	Memory     string            // filled automatically
-	State      NodeState         // filled automatically
-	IP         NodeIP            // filled automatically
+	Name          string            `yaml:"name" json:"name,omitempty"`
+	Role          Role              `yaml:"role" json:"role,omitempty"`
+	Image         string            `yaml:"image" json:"image,omitempty"`
+	Volumes       []string          `yaml:"volumes" json:"volumes,omitempty"`
+	Env           []string          `yaml:"env" json:"env,omitempty"`
+	Cmd           []string          // filled automatically based on role
+	Args          []string          `yaml:"extraArgs" json:"extraArgs,omitempty"`
+	Ports         nat.PortMap       `yaml:"portMappings" json:"portMappings,omitempty"`
+	Restart       bool              `yaml:"restart" json:"restart,omitempty"`
+	Created       string            `yaml:"created" json:"created,omitempty"`
+	Labels        map[string]string // filled automatically
+	K3sNodeLabels map[string]string `yaml:"k3sNodeLabels" json:"k3sNodeLabels,omitempty"`
+	Networks      []string          // filled automatically
+	ExtraHosts    []string          // filled automatically
+	ServerOpts    ServerOpts        `yaml:"serverOpts" json:"serverOpts,omitempty"`
+	AgentOpts     AgentOpts         `yaml:"agentOpts" json:"agentOpts,omitempty"`
+	GPURequest    string            // filled automatically
+	Memory        string            // filled automatically
+	State         NodeState         // filled automatically
+	IP            NodeIP            // filled automatically
 }
 
 // ServerOpts describes some additional server role specific opts