Create tags with templated strings (#48)

mtougeron · May 8, 2022 · 9499551 · 9499551
1 parent 25bd54a
commit 9499551
Show file tree

Hide file tree

Showing 3 changed files with 146 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -34,11 +34,40 @@ The `k8s-aws-ebs-tagger` watches for new PersistentVolumeClaims and when new AWS
 
 #### ignored tags
 
-The following tags are ignored
+The following tags are ignored by default
  - `kubernetes.io/*`
  - `KubernetesCluster`
  - `Name`
 
+#### Tag Templates
+
+Tag values can be Go templates using values from the PVC's `Name`, `Namespace`, `Annotations`, and `Labels`.
+
+Some examples could be:
+
+```yaml
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: touge-test
+  namespace: touge
+  labels:
+    TeamID: "Frontend"
+  annotations:
+    CostCenter: "1234"
+    aws-ebs-tagger/tags: |
+      {"Owner": "{{ .Labels.TeamID }}-{{ .Annotations.CostCenter }}"}
+---
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: app-1
+  namespace: my-app
+  annotations:
+    aws-ebs-tagger/tags: |
+      {"OwnerID": "{{ .Namespace }}/{{ .Name }}"}
+```
+
 ### Installation
 
 #### AWS IAM Role

diff --git a/kubernetes.go b/kubernetes.go
@@ -19,9 +19,11 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
+	"html/template"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -50,6 +52,13 @@ const (
 	regexpAWSVolumeID = `^aws:\/\/\w{2}-\w{4,9}-\d\w\/(vol-\w+)$`
 )
 
+type TagTemplate struct {
+	Name        string
+	Namespace   string
+	Labels      map[string]string
+	Annotations map[string]string
+}
+
 func BuildClient(kubeconfig string, kubeContext string) (*kubernetes.Clientset, error) {
 	config, err := rest.InClusterConfig()
 	if err != nil {
@@ -171,7 +180,7 @@ func buildTags(pvc *corev1.PersistentVolumeClaim) map[string]string {
 	if _, ok := annotations[annotationPrefix+"/ignore"]; ok {
 		log.Debugln(annotationPrefix + "/ignore annotation is set")
 		promIgnoredTotal.Inc()
-		return tags
+		return renderTagTemplates(pvc, tags)
 	}
 
 	// Set the default tags
@@ -191,7 +200,7 @@ func buildTags(pvc *corev1.PersistentVolumeClaim) map[string]string {
 	tagString, ok := annotations[annotationPrefix+"/tags"]
 	if !ok {
 		log.Debugln("Does not have " + annotationPrefix + "/tags annotation")
-		return tags
+		return renderTagTemplates(pvc, tags)
 	}
 	if tagFormat == "csv" {
 		customTags = parseCsv(tagString)
@@ -215,6 +224,31 @@ func buildTags(pvc *corev1.PersistentVolumeClaim) map[string]string {
 		tags[k] = v
 	}
 
+	return renderTagTemplates(pvc, tags)
+}
+
+func renderTagTemplates(pvc *corev1.PersistentVolumeClaim, tags map[string]string) map[string]string {
+
+	tplData := TagTemplate{
+		Name:        pvc.GetName(),
+		Namespace:   pvc.GetNamespace(),
+		Labels:      pvc.GetLabels(),
+		Annotations: pvc.GetAnnotations(),
+	}
+
+	for k, v := range tags {
+		tmpl, err := template.New("tag").Parse(v)
+		if err != nil {
+			continue
+		}
+		buf := new(bytes.Buffer)
+		err = tmpl.Execute(buf, tplData)
+		if err != nil {
+			continue
+		}
+		tags[k] = buf.String()
+	}
+
 	return tags
 }
 

diff --git a/kubernetes_test.go b/kubernetes_test.go
@@ -315,6 +315,7 @@ func Test_buildTags(t *testing.T) {
 				t.Errorf("buildTags() = %v, want %v", got, tt.want)
 			}
 			tagFormat = "json"
+			defaultTags = map[string]string{}
 		})
 	}
 }
@@ -323,6 +324,7 @@ func Test_annotationPrefix(t *testing.T) {
 
 	pvc := &corev1.PersistentVolumeClaim{}
 	pvc.SetName("my-pvc")
+	defaultAnnotationPrefix := annotationPrefix
 
 	tests := []struct {
 		name             string
@@ -368,6 +370,8 @@ func Test_annotationPrefix(t *testing.T) {
 			if got := buildTags(pvc); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("buildTags() = %v, want %v", got, tt.want)
 			}
+			annotationPrefix = defaultAnnotationPrefix
+			defaultTags = map[string]string{}
 		})
 	}
 }
@@ -484,3 +488,79 @@ func Test_processPersistentVolumeClaim(t *testing.T) {
 	}
 
 }
+
+func Test_templatedTags(t *testing.T) {
+
+	pvc := &corev1.PersistentVolumeClaim{}
+	pvc.SetName("my-pvc")
+	pvc.SetNamespace("my-namespace")
+
+	tests := []struct {
+		name        string
+		defaultTags map[string]string
+		annotations map[string]string
+		labels      map[string]string
+		want        map[string]string
+	}{
+		{
+			name:        "default tag with template",
+			defaultTags: map[string]string{"foo": "{{ .Name }}-{{ .Namespace }}"},
+			annotations: map[string]string{},
+			labels:      map[string]string{},
+			want:        map[string]string{"foo": "my-pvc-my-namespace"},
+		},
+		{
+			name:        "default tag overwritten with tag template",
+			defaultTags: map[string]string{"foo": "bar"},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Name }}-{{ .Namespace }}\"}"},
+			labels:      map[string]string{},
+			want:        map[string]string{"foo": "my-pvc-my-namespace"},
+		},
+		{
+			name:        "template using annotation",
+			defaultTags: map[string]string{},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Name }}-{{ .Annotations.TeamID }}\"}", "TeamID": "1234"},
+			labels:      map[string]string{},
+			want:        map[string]string{"foo": "my-pvc-1234"},
+		},
+		{
+			name:        "template using label",
+			defaultTags: map[string]string{},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Name }}-{{ .Labels.TeamID }}\"}"},
+			labels:      map[string]string{"TeamID": "1234"},
+			want:        map[string]string{"foo": "my-pvc-1234"},
+		},
+		{
+			name:        "template using label and annotation",
+			defaultTags: map[string]string{},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Name }}-{{ .Labels.TeamID }}\",\"bar\": \"{{ .Name }}-{{ .Annotations.DeptID }}\"}", "DeptID": "ABC"},
+			labels:      map[string]string{"TeamID": "1234"},
+			want:        map[string]string{"foo": "my-pvc-1234", "bar": "my-pvc-ABC"},
+		},
+		{
+			name:        "template using invalid label",
+			defaultTags: map[string]string{},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Name }}-{{ .Labels.SomeLabel }}\"}"},
+			labels:      map[string]string{"TeamID": "1234"},
+			want:        map[string]string{"foo": "my-pvc-"},
+		},
+		{
+			name:        "template using invalid field",
+			defaultTags: map[string]string{},
+			annotations: map[string]string{annotationPrefix + "/tags": "{\"foo\": \"{{ .Blah }}-{{ .Labels.TeamID }}\"}"},
+			labels:      map[string]string{"TeamID": "1234"},
+			want:        map[string]string{"foo": "{{ .Blah }}-{{ .Labels.TeamID }}"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pvc.SetAnnotations(tt.annotations)
+			pvc.SetLabels(tt.labels)
+			defaultTags = tt.defaultTags
+			if got := buildTags(pvc); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("buildTags() = %v, want %v", got, tt.want)
+			}
+			defaultTags = map[string]string{}
+		})
+	}
+}