feat: helm v2-alpha: add nodeSelector, toleration and affinity to the chart

Author: robinlioret

pkg/plugins/common/kustomize/v2/scaffolds/internal/templates/config/manager/config.go

@@ -96,6 +96,14 @@ spec:
       #             operator: In
       #             values:
       #               - linux
+			# TODO(user): Uncomment the following code to configure the nodeSelector.
+			# nodeSelector: 
+			#		kubernetes.io/os: linux
+			# TODO(user): Uncomment the following code to configure the tolerations.
+      #	tolerations:
+			# - key: "key1"
+			#		operator: "Exists"
+			#		effect: "NoSchedule"
       securityContext:
         # Projects are configured by default to adhere to the "restricted" Pod Security Standards.
         # This ensures that deployments meet the highest security requirements for Kubernetes.

pkg/plugins/optional/helm/v2alpha/scaffolds/internal/kustomize/chart_converter.go

@@ -105,6 +105,9 @@ func (c *ChartConverter) ExtractDeploymentConfig() map[string]interface{} {
 	}
 
 	extractPodSecurityContext(specMap, config)
+	extractPodNodeSelector(specMap, config)
+	extractPodTolerations(specMap, config)
+	extractPodAffinity(specMap, config)
 
 	container := firstManagerContainer(specMap)
 	if container == nil {
@@ -149,6 +152,48 @@ func extractPodSecurityContext(specMap map[string]interface{}, config map[string
 	config["podSecurityContext"] = podSecurityContext
 }
 
+func extractPodNodeSelector(specMap map[string]interface{}, config map[string]interface{}) {
+	raw, found, err := unstructured.NestedFieldNoCopy(specMap, "nodeSelector")
+	if !found || err != nil {
+		return
+	}
+
+	result, ok := raw.(map[string]interface{})
+	if !ok || len(result) == 0 {
+		return
+	}
+
+	config["podNodeSelector"] = result
+}
+
+func extractPodTolerations(specMap map[string]interface{}, config map[string]interface{}) {
+	raw, found, err := unstructured.NestedFieldNoCopy(specMap, "tolerations")
+	if !found || err != nil {
+		return
+	}
+
+	result, ok := raw.([]interface{})
+	if !ok || len(result) == 0 {
+		return
+	}
+
+	config["podTolerations"] = result
+}
+
+func extractPodAffinity(specMap map[string]interface{}, config map[string]interface{}) {
+	raw, found, err := unstructured.NestedFieldNoCopy(specMap, "affinity")
+	if !found || err != nil {
+		return
+	}
+
+	result, ok := raw.(map[string]interface{})
+	if !ok || len(result) == 0 {
+		return
+	}
+
+	config["podAffinity"] = result
+}
+
 func firstManagerContainer(specMap map[string]interface{}) map[string]interface{} {
 	containers, found, err := unstructured.NestedFieldNoCopy(specMap, "containers")
 	if !found || err != nil {

pkg/plugins/optional/helm/v2alpha/scaffolds/internal/kustomize/helm_templater.go

@@ -190,6 +190,24 @@ func (t *HelmTemplater) templateDeploymentFields(yamlContent string) string {
 	yamlContent = t.templateVolumeMounts(yamlContent)
 	yamlContent = t.templateVolumes(yamlContent)
 	yamlContent = t.templateControllerManagerArgs(yamlContent)
+	yamlContent = t.templateBasicWithStatement(
+		yamlContent,
+		"nodeSelector",
+		"spec.template.spec",
+		".Values.manager.nodeSelector",
+	)
+	yamlContent = t.templateBasicWithStatement(
+		yamlContent,
+		"affinity",
+		"spec.template.spec",
+		".Values.manager.affinity",
+	)
+	yamlContent = t.templateBasicWithStatement(
+		yamlContent,
+		"tolerations",
+		"spec.template.spec",
+		".Values.manager.tolerations",
+	)
 
 	return yamlContent
 }
@@ -620,6 +638,88 @@ func (t *HelmTemplater) templateImageReference(yamlContent string) string {
 	return yamlContent
 }
 
+func (t *HelmTemplater) templateBasicWithStatement(
+	yamlContent string,
+	key string,
+	parentKey string,
+	valuePath string,
+) string {
+	lines := strings.Split(yamlContent, "\n")
+	yamlKey := fmt.Sprintf("%s:", key)
+
+	var start, end int
+	var indentLen int
+	if !strings.Contains(yamlContent, yamlKey) {
+		// Find parent block start if the key is missing
+		pKeyParts := strings.Split(parentKey, ".")
+		pKeyIdx := 0
+		pKeyInit := false
+		currIndent := 0
+		for i := 0; i < len(lines); i++ {
+			_, lineIndent := leadingWhitespace(lines[i])
+			if pKeyInit && lineIndent <= currIndent {
+				return yamlContent
+			}
+			if !strings.HasPrefix(strings.TrimSpace(lines[i]), pKeyParts[pKeyIdx]) {
+				continue
+			}
+
+			// Parent key part found
+			pKeyIdx++
+			pKeyInit = true
+			if pKeyIdx >= len(pKeyParts) {
+				start = i + 1
+				end = start
+				break
+			}
+		}
+		_, indentLen = leadingWhitespace(lines[start])
+	} else {
+		// Find the existing block
+		for i := 0; i < len(lines); i++ {
+			if !strings.HasPrefix(strings.TrimSpace(lines[i]), key) {
+				continue
+			}
+			start = i
+			end = i + 1
+			trimmed := strings.TrimSpace(lines[i])
+			if len(trimmed) == len(yamlKey) {
+				_, indentLenSearch := leadingWhitespace(lines[i])
+				for j := end; j < len(lines); j++ {
+					_, indentLenLine := leadingWhitespace(lines[j])
+					if indentLenLine <= indentLenSearch {
+						end = j
+						break
+					}
+				}
+			}
+		}
+		_, indentLen = leadingWhitespace(lines[start])
+	}
+
+	indentStr := strings.Repeat(" ", indentLen)
+
+	var builder strings.Builder
+	builder.WriteString(indentStr)
+	builder.WriteString("{{- with ")
+	builder.WriteString(valuePath)
+	builder.WriteString(" }}\n")
+	builder.WriteString(indentStr)
+	builder.WriteString(yamlKey)
+	builder.WriteString(" {{ toYaml . | nindent ")
+	builder.WriteString(strconv.Itoa(indentLen + 4))
+	builder.WriteString(" }}\n")
+	builder.WriteString(indentStr)
+	builder.WriteString("{{- end }}\n")
+
+	newBlock := strings.TrimRight(builder.String(), "\n")
+
+	newLines := append([]string{}, lines[:start]...)
+	newLines = append(newLines, strings.Split(newBlock, "\n")...)
+	newLines = append(newLines, lines[end:]...)
+	return strings.Join(newLines, "\n")
+}
+
 // makeWebhookAnnotationsConditional makes only cert-manager annotations conditional, not the entire webhook
 func (t *HelmTemplater) makeWebhookAnnotationsConditional(yamlContent string) string {
 	// Find cert-manager.io/inject-ca-from annotation and make it conditional

pkg/plugins/optional/helm/v2alpha/scaffolds/internal/templates/values_basic.go

@@ -210,15 +210,7 @@ func (f *HelmValuesBasic) addDeploymentConfig(buf *bytes.Buffer) {
 		buf.WriteString("  # Environment variables\n")
 		buf.WriteString("  env:\n")
 		if envYaml, err := yaml.Marshal(env); err == nil {
-			// Indent the YAML properly
-			lines := bytes.Split(envYaml, []byte("\n"))
-			for _, line := range lines {
-				if len(line) > 0 {
-					buf.WriteString("    ")
-					buf.Write(line)
-					buf.WriteString("\n")
-				}
-			}
+			f.IndentYamlProperly(buf, envYaml)
 		} else {
 			buf.WriteString("    []\n")
 		}
@@ -233,14 +225,7 @@ func (f *HelmValuesBasic) addDeploymentConfig(buf *bytes.Buffer) {
 		buf.WriteString("  # Pod-level security settings\n")
 		buf.WriteString("  podSecurityContext:\n")
 		if secYaml, err := yaml.Marshal(podSecCtx); err == nil {
-			lines := bytes.Split(secYaml, []byte("\n"))
-			for _, line := range lines {
-				if len(line) > 0 {
-					buf.WriteString("    ")
-					buf.Write(line)
-					buf.WriteString("\n")
-				}
-			}
+			f.IndentYamlProperly(buf, secYaml)
 		}
 		buf.WriteString("\n")
 	} else {
@@ -252,14 +237,7 @@ func (f *HelmValuesBasic) addDeploymentConfig(buf *bytes.Buffer) {
 		buf.WriteString("  # Container-level security settings\n")
 		buf.WriteString("  securityContext:\n")
 		if secYaml, err := yaml.Marshal(secCtx); err == nil {
-			lines := bytes.Split(secYaml, []byte("\n"))
-			for _, line := range lines {
-				if len(line) > 0 {
-					buf.WriteString("    ")
-					buf.Write(line)
-					buf.WriteString("\n")
-				}
-			}
+			f.IndentYamlProperly(buf, secYaml)
 		}
 		buf.WriteString("\n")
 	} else {
@@ -271,19 +249,59 @@ func (f *HelmValuesBasic) addDeploymentConfig(buf *bytes.Buffer) {
 		buf.WriteString("  # Resource limits and requests\n")
 		buf.WriteString("  resources:\n")
 		if resYaml, err := yaml.Marshal(resources); err == nil {
-			lines := bytes.Split(resYaml, []byte("\n"))
-			for _, line := range lines {
-				if len(line) > 0 {
-					buf.WriteString("    ")
-					buf.Write(line)
-					buf.WriteString("\n")
-				}
-			}
+			f.IndentYamlProperly(buf, resYaml)
 		}
 		buf.WriteString("\n")
 	} else {
 		f.addDefaultResources(buf)
 	}
+
+	buf.WriteString("  # Pod's affinity\n")
+	if affinity, exists := f.DeploymentConfig["podAffinity"]; exists && affinity != nil {
+		buf.WriteString("  affinity:\n")
+		if affYaml, err := yaml.Marshal(affinity); err == nil {
+			f.IndentYamlProperly(buf, affYaml)
+		}
+		buf.WriteString("\n")
+	} else {
+		buf.WriteString("  affinity: {}\n")
+		buf.WriteString("\n")
+	}
+
+	buf.WriteString("  # Pod's node selector\n")
+	if nodeSelector, exists := f.DeploymentConfig["podNodeSelector"]; exists && nodeSelector != nil {
+		buf.WriteString("  nodeSelector:\n")
+		if nodYaml, err := yaml.Marshal(nodeSelector); err == nil {
+			f.IndentYamlProperly(buf, nodYaml)
+		}
+		buf.WriteString("\n")
+	} else {
+		buf.WriteString("  nodeSelector: {}\n")
+		buf.WriteString("\n")
+	}
+
+	buf.WriteString("  # Pod's tolerations\n")
+	if tolerations, exists := f.DeploymentConfig["podTolerations"]; exists && tolerations != nil {
+		buf.WriteString("  tolerations:\n")
+		if tolYaml, err := yaml.Marshal(tolerations); err == nil {
+			f.IndentYamlProperly(buf, tolYaml)
+		}
+		buf.WriteString("\n")
+	} else {
+		buf.WriteString("  tolerations: []\n")
+		buf.WriteString("\n")
+	}
+}
+
+func (f *HelmValuesBasic) IndentYamlProperly(buf *bytes.Buffer, envYaml []byte) {
+	lines := bytes.Split(envYaml, []byte("\n"))
+	for _, line := range lines {
+		if len(line) > 0 {
+			buf.WriteString("    ")
+			buf.Write(line)
+			buf.WriteString("\n")
+		}
+	}
 }
 
 // addDefaultDeploymentSections adds default sections when no deployment config is available

pkg/plugins/optional/helm/v2alpha/scaffolds/internal/templates/values_basic_test.go

@@ -159,6 +159,76 @@ var _ = Describe("HelmValuesBasic", func() {
 			Expect(content).To(ContainSubstring("resources:"))
 			Expect(content).To(ContainSubstring("cpu: 100m"))
 			Expect(content).To(ContainSubstring("memory: 128Mi"))
+			Expect(content).To(ContainSubstring("affinity: {}"))
+			Expect(content).To(ContainSubstring("nodeSelector: {}"))
+			Expect(content).To(ContainSubstring("tolerations: []"))
+		})
+	})
+
+	Context("with nodeSelector, affinity and tolerations configuration", func() {
+		BeforeEach(func() {
+			deploymentConfig := map[string]interface{}{
+				"podNodeSelector": map[string]string{
+					"kubernetes.io/os": "linux",
+				},
+				"podTolerations": []map[string]string{
+					{
+						"key":      "key1",
+						"operator": "Equal",
+						"effect":   "NoSchedule",
+					},
+				},
+				"podAffinity": map[string]interface{}{
+					"nodeAffinity": map[string]interface{}{
+						"requiredDuringSchedulingIgnoredDuringExecution": map[string]interface{}{
+							"nodeSelectorTerms": []interface{}{
+								map[string]interface{}{
+									"matchExpressions": []interface{}{
+										map[string]interface{}{
+											"key":      "topology.kubernetes.io/zone",
+											"operator": "In",
+											"values":   []string{"antarctica-east1", "antarctica-east2"},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			}
+
+			valuesTemplate = &HelmValuesBasic{
+				HasWebhooks:      false,
+				DeploymentConfig: deploymentConfig,
+			}
+			valuesTemplate.InjectProjectName("test-project")
+			err := valuesTemplate.SetTemplateDefaults()
+			Expect(err).NotTo(HaveOccurred())
+		})
+
+		It("should include default values", func() {
+			content := valuesTemplate.GetBody()
+			Expect(content).To(ContainSubstring(`  # Pod's node selector
+  nodeSelector:
+    kubernetes.io/os: linux`))
+
+			Expect(content).To(ContainSubstring(`  # Pod's tolerations
+  tolerations:
+    - effect: NoSchedule
+      key: key1
+      operator: Equal`))
+
+			Expect(content).To(ContainSubstring(`  # Pod's affinity
+  affinity:
+    nodeAffinity:
+        requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+                - matchExpressions:
+                    - key: topology.kubernetes.io/zone
+                      operator: In
+                      values:
+                        - antarctica-east1
+                        - antarctica-east2`))
 		})
 	})