From d48fd2f91188956a35d4cabf5531a8ef8ba42438 Mon Sep 17 00:00:00 2001
From: Patrice Ferlet <metal3d@gmail.com>
Date: Sun, 21 Apr 2024 16:34:21 +0200
Subject: [PATCH] make a better override + add more values (serviceAccount,
 nodeSelector...)

---
 generator/deployment.go | 55 +++++++++++++++++++++++++++++++++++------
 generator/values.go     |  2 ++
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/generator/deployment.go b/generator/deployment.go
index 26e7c7e..b674cef 100644
--- a/generator/deployment.go
+++ b/generator/deployment.go
@@ -81,6 +81,11 @@ func NewDeployment(service types.ServiceConfig, chart *HelmChart) *Deployment {
 					ObjectMeta: metav1.ObjectMeta{
 						Labels: GetMatchLabels(service.Name, appName),
 					},
+					Spec: corev1.PodSpec{
+						NodeSelector: map[string]string{
+							"katenary.v3/node-selector": "replace",
+						},
+					},
 				},
 			},
 		},
@@ -161,6 +166,9 @@ func (d *Deployment) AddContainer(service types.ServiceConfig) {
 		Name: `{{ .Values.pullSecrets | toYaml | indent __indent__ }}`,
 	}}
 
+	// add ServiceAccount to the deployment
+	d.Spec.Template.Spec.ServiceAccountName = `{{ .Values.` + service.Name + `.serviceAccount | quote }}`
+
 	d.AddHealthCheck(service, &container)
 
 	d.Spec.Template.Spec.Containers = append(d.Spec.Template.Spec.Containers, container)
@@ -539,24 +547,55 @@ func (d *Deployment) Yaml() ([]byte, error) {
 	}
 
 	// for impagePullSecrets, replace the name with the value from values.yaml
-	inpullsecrets := false
 	for i, line := range content {
 		if strings.Contains(line, "imagePullSecrets:") {
-			inpullsecrets = true
-		}
-		if inpullsecrets && strings.Contains(line, "- name: ") && inpullsecrets {
-			line = strings.Replace(line, "- name: ", "", 1)
-			line = strings.ReplaceAll(line, "'", "")
+			spaces = strings.Repeat(" ", utils.CountStartingSpaces(line))
+			line = spaces + "{{- if .Values.pullSecrets }}"
+			line += "\n" + spaces + "imagePullSecrets:\n"
+			line += spaces + "{{- .Values.pullSecrets | toYaml | nindent __indent__ }}"
+			line += "\n" + spaces + "{{- end }}"
 			content[i] = line
-			inpullsecrets = false
 		}
 	}
 
 	// Find the replicas line and replace it with the value from values.yaml
 	for i, line := range content {
+		// manage nodeSelector
+		if strings.Contains(line, "nodeSelector:") {
+			spaces = strings.Repeat(" ", utils.CountStartingSpaces(line))
+			pre := spaces + `{{- if .Values.` + serviceName + `.nodeSelector }}`
+			post := spaces + "{{- end }}"
+			ns := spaces + "nodeSelector:\n"
+			ns += spaces + `  {{- .Values.` + serviceName + `.nodeSelector | toYaml | nindent __indent__ }}`
+			//line = strings.Replace(line, "katenary.v3/node-selector: replace", ns, 1)
+			line = pre + "\n" + ns + "\n" + post
+		}
+		// manage replicas
 		if strings.Contains(line, "replicas:") {
 			line = regexp.MustCompile("replicas: .*$").ReplaceAllString(line, "replicas: {{ .Values."+serviceName+".replicas }}")
-			content[i] = line
+		}
+
+		// manage serviceAccount, add condition to use the serviceAccount from values.yaml
+		if strings.Contains(line, "serviceAccountName:") {
+			spaces = strings.Repeat(" ", utils.CountStartingSpaces(line))
+			pre := spaces + `{{- if ne .Values.` + serviceName + `.serviceAccount "" }}`
+			post := spaces + "{{- end }}"
+			line = strings.ReplaceAll(line, "'", "")
+			line = pre + "\n" + line + "\n" + post
+		}
+
+		content[i] = line
+	}
+
+	// find the katenary.v3/node-selector line, and remove it
+	for i, line := range content {
+		if strings.Contains(line, "katenary.v3/node-selector") {
+			content = append(content[:i], content[i+1:]...)
+			continue
+		}
+		if strings.Contains(line, "- name: '{{ .Values.pullSecrets ") {
+			content = append(content[:i], content[i+1:]...)
+			continue
 		}
 	}
 
diff --git a/generator/values.go b/generator/values.go
index 055a6b3..725db51 100644
--- a/generator/values.go
+++ b/generator/values.go
@@ -57,6 +57,8 @@ type Value struct {
 	Environment     map[string]any               `yaml:"environment,omitempty"`
 	Replicas        *uint32                      `yaml:"replicas,omitempty"`
 	CronJob         *CronJobValue                `yaml:"cronjob,omitempty"`
+	NodeSelector    map[string]string            `yaml:"nodeSelector"`
+	ServiceAccount  string                       `yaml:"serviceAccount"`
 }
 
 // CronJobValue is a cronjob configuration that will be saved in values.yaml.