From ffe8c9697c530b4f8f0fda97e5f068316e1fe7b2 Mon Sep 17 00:00:00 2001
From: Lionel Nicolas <lionel.nicolas@nividic.org>
Date: Fri, 16 Oct 2020 20:46:04 -0400
Subject: [PATCH] add ability to set container labels on nodes using --label/-l

This supports node filters as is already done for '--volume' and '--port'.
---
 cmd/cluster/clusterCreate.go | 44 ++++++++++++++++++++++++++++++++++++
 cmd/util/labels.go           | 39 ++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 cmd/util/labels.go

diff --git a/cmd/cluster/clusterCreate.go b/cmd/cluster/clusterCreate.go
index 3cf609cf..c953cfe5 100644
--- a/cmd/cluster/clusterCreate.go
+++ b/cmd/cluster/clusterCreate.go
@@ -128,6 +128,7 @@ func NewCmdClusterCreate() *cobra.Command {
 	cmd.Flags().String("token", "", "Specify a cluster token. By default, we generate one.")
 	cmd.Flags().StringArrayP("volume", "v", nil, "Mount volumes into the nodes (Format: `[SOURCE:]DEST[@NODEFILTER[;NODEFILTER...]]`\n - Example: `k3d cluster create --agents 2 -v \"/my/path@agent[0,1]\" -v \"/tmp/test:/tmp/other@server[0]\"`")
 	cmd.Flags().StringArrayP("port", "p", nil, "Map ports from the node containers to the host (Format: `[HOST:][HOSTPORT:]CONTAINERPORT[/PROTOCOL][@NODEFILTER]`)\n - Example: `k3d cluster create --agents 2 -p \"8080:80@agent[0]\" -p \"8081@agent[1]\"`")
+	cmd.Flags().StringArrayP("label", "l", nil, "Add label to node container (Format: `KEY[=VALUE][@NODEFILTER[;NODEFILTER...]]`\n - Example: `k3d cluster create --agents 2 -l \"my.label@agent[0,1]\" -v \"other.label=somevalue@server[0]\"`")
 	cmd.Flags().BoolVar(&createClusterOpts.WaitForServer, "wait", true, "Wait for the server(s) to be ready before returning. Use '--timeout DURATION' to not wait forever.")
 	cmd.Flags().DurationVar(&createClusterOpts.Timeout, "timeout", 0*time.Second, "Rollback changes if cluster couldn't be created in specified duration.")
 	cmd.Flags().BoolVar(&updateDefaultKubeconfig, "update-default-kubeconfig", true, "Directly update the default kubeconfig with the new cluster's context")
@@ -317,6 +318,32 @@ func parseCreateClusterCmd(cmd *cobra.Command, args []string, createClusterOpts
 
 	log.Tracef("PortFilterMap: %+v", portFilterMap)
 
+	// --label
+	labelFlags, err := cmd.Flags().GetStringArray("label")
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	// labelFilterMap will add container label to applied node filters
+	labelFilterMap := make(map[string][]string, 1)
+	for _, labelFlag := range labelFlags {
+
+		// split node filter from the specified label
+		label, filters, err := cliutil.SplitFiltersFromFlag(labelFlag)
+		if err != nil {
+			log.Fatalln(err)
+		}
+
+		// create new entry or append filter to existing entry
+		if _, exists := labelFilterMap[label]; exists {
+			labelFilterMap[label] = append(labelFilterMap[label], filters...)
+		} else {
+			labelFilterMap[label] = filters
+		}
+	}
+
+	log.Tracef("LabelFilterMap: %+v", labelFilterMap)
+
 	/********************
 	 *									*
 	 * generate cluster *
@@ -411,6 +438,23 @@ func parseCreateClusterCmd(cmd *cobra.Command, args []string, createClusterOpts
 		}
 	}
 
+	// append labels
+	for label, filters := range labelFilterMap {
+		nodes, err := cliutil.FilterNodes(cluster.Nodes, filters)
+		if err != nil {
+			log.Fatalln(err)
+		}
+		for _, node := range nodes {
+			// ensure node.Labels map is initialized (see also ClusterCreate.nodeSetup)
+			if node.Labels == nil {
+				node.Labels = make(map[string]string)
+			}
+
+			labelKey, labelValue := cliutil.SplitLabelKeyValue(label)
+			node.Labels[labelKey] = labelValue
+		}
+	}
+
 	/**********************
 	 * Utility Containers *
 	 **********************/
diff --git a/cmd/util/labels.go b/cmd/util/labels.go
new file mode 100644
index 00000000..2a60e0ab
--- /dev/null
+++ b/cmd/util/labels.go
@@ -0,0 +1,39 @@
+/*
+Copyright © 2020 The k3d Author(s)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+package util
+
+import (
+	"strings"
+)
+
+// SplitLabelKeyValue separates the label key from the label value (if any)
+func SplitLabelKeyValue(label string) (string, string) {
+	// split only on first '=' sign (like `docker run` do)
+	labelSlice := strings.SplitN(label, "=", 2)
+
+	if len(labelSlice) > 1 {
+		return labelSlice[0], labelSlice[1]
+	}
+
+	// defaults to label key with empty value (like `docker run` do)
+	return label, ""
+}