feat: add node template custom taint effect support

castai · Sep 14, 2023 · fc9bc10 · fc9bc10
1 parent db2c6f0
commit fc9bc10
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 20 deletions.
diff --git a/castai/resource_node_template.go b/castai/resource_node_template.go
@@ -11,7 +11,6 @@ import (
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
 	"github.com/samber/lo"
 	"log"
-	"regexp"
 	"strings"
 	"time"
 )
@@ -55,6 +54,12 @@ const (
 	FieldNodeTemplateUseSpotFallbacks                       = "use_spot_fallbacks"
 )
 
+const (
+	TaintEffectNoSchedule       = "NoSchedule"
+	TaintEffectNoExecute        = "NoExecute"
+	TaintEffectPreferNoSchedule = "PreferNoSchedule"
+)
+
 const (
 	ArchAMD64 = "amd64"
 	ArchARM64 = "arm64"
@@ -346,11 +351,11 @@ func resourceNodeTemplate() *schema.Resource {
 						FieldEffect: {
 							Optional: true,
 							Type:     schema.TypeString,
-							Default:  "NoSchedule",
+							Default:  TaintEffectNoSchedule,
 							ValidateDiagFunc: validation.ToDiagFunc(
-								validation.StringMatch(regexp.MustCompile("^NoSchedule$"), "effect must be NoSchedule"),
+								validation.StringInSlice([]string{TaintEffectNoSchedule, TaintEffectNoExecute, TaintEffectPreferNoSchedule}, false),
 							),
-							Description: "Effect of a taint to be added to nodes created from this template. The effect must always be NoSchedule.",
+							Description: "Effect of a taint to be added to nodes created from this template.",
 						},
 					},
 				},
@@ -630,7 +635,7 @@ func updateNodeTemplate(ctx context.Context, d *schema.ResourceData, meta any, s
 			ts = append(ts, val.(map[string]any))
 		}
 
-		req.CustomTaints = toCustomTaintsWithoutEffect(ts)
+		req.CustomTaints = toCustomTaintsWithOptionalEffect(ts)
 	}
 
 	if !(*req.ShouldTaint) && req.CustomTaints != nil && len(*req.CustomTaints) > 0 {
@@ -709,7 +714,7 @@ func resourceNodeTemplateCreate(ctx context.Context, d *schema.ResourceData, met
 			ts = append(ts, val.(map[string]any))
 		}
 
-		req.CustomTaints = toCustomTaintsWithoutEffect(ts)
+		req.CustomTaints = toCustomTaintsWithOptionalEffect(ts)
 	}
 
 	if !(*req.ShouldTaint) && req.CustomTaints != nil && len(*req.CustomTaints) > 0 {
@@ -848,22 +853,25 @@ func toCustomLabel(obj map[string]any) *sdk.NodetemplatesV1Label {
 	return out
 }
 
-func toCustomTaintsWithoutEffect(objs []map[string]any) *[]sdk.NodetemplatesV1TaintWithoutEffect {
+func toCustomTaintsWithOptionalEffect(objs []map[string]any) *[]sdk.NodetemplatesV1TaintWithOptionalEffect {
 	if len(objs) == 0 {
 		return nil
 	}
 
-	out := &[]sdk.NodetemplatesV1TaintWithoutEffect{}
+	out := &[]sdk.NodetemplatesV1TaintWithOptionalEffect{}
 
 	for _, taint := range objs {
-		t := sdk.NodetemplatesV1TaintWithoutEffect{}
+		t := sdk.NodetemplatesV1TaintWithOptionalEffect{}
 
 		if v, ok := taint[FieldKey]; ok && v != "" {
-			t.Key = toPtr(v.(string))
+			t.Key = v.(string)
 		}
 		if v, ok := taint[FieldValue]; ok && v != "" {
 			t.Value = toPtr(v.(string))
 		}
+		if v, ok := taint[FieldEffect]; ok && v != "" {
+			t.Effect = toPtr(sdk.NodetemplatesV1TaintEffect(v.(string)))
+		}
 
 		*out = append(*out, t)
 	}

diff --git a/castai/resource_node_template_test.go b/castai/resource_node_template_test.go
@@ -373,11 +373,16 @@ func TestAccResourceNodeTemplate_basic(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "custom_labels.%", "2"),
 					resource.TestCheckResourceAttr(resourceName, "custom_labels."+rName+"-label-key-1", rName+"-label-value-1"),
 					resource.TestCheckResourceAttr(resourceName, "custom_labels."+rName+"-label-key-2", rName+"-label-value-2"),
-					resource.TestCheckResourceAttr(resourceName, "custom_taints.#", "2"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.#", "3"),
 					resource.TestCheckResourceAttr(resourceName, "custom_taints.0.key", rName+"-taint-key-1"),
 					resource.TestCheckResourceAttr(resourceName, "custom_taints.0.value", rName+"-taint-value-1"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.0.effect", "NoSchedule"),
 					resource.TestCheckResourceAttr(resourceName, "custom_taints.1.key", rName+"-taint-key-2"),
 					resource.TestCheckResourceAttr(resourceName, "custom_taints.1.value", rName+"-taint-value-2"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.1.value", "NoExecute"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.2.key", rName+"-taint-key-3"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.2.value", rName+"-taint-value-3"),
+					resource.TestCheckResourceAttr(resourceName, "custom_taints.2.value", "NoSchedule"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.instance_families.0.exclude.0", "m5"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.gpu.0.manufacturers.0", "NVIDIA"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.gpu.0.include_names.#", "0"),
@@ -470,7 +475,12 @@ func testAccNodeTemplateConfig(rName, clusterName string) string {
 			custom_taints {
 				key = "%[1]s-taint-key-2"
 				value = "%[1]s-taint-value-2"
-				effect = "NoSchedule"
+				effect = "NoExecute"
+			}
+
+			custom_taints {
+				key = "%[1]s-taint-key-3"
+				value = "%[1]s-taint-value-3"
 			}
 
 			constraints {

diff --git a/castai/sdk/api.gen.go b/castai/sdk/api.gen.go
diff --git a/docs/resources/node_template.md b/docs/resources/node_template.md
diff --git a/examples/gke/gke_cluster_autoscaler_policies/castai.tf b/examples/gke/gke_cluster_autoscaler_policies/castai.tf
@@ -88,10 +88,12 @@ module "castai-gke-cluster" {
         {
           key = "custom-taint-key-1"
           value = "custom-taint-value-1"
+          effect = "NoSchedule"
         },
         {
           key = "custom-taint-key-2"
           value = "custom-taint-value-2"
+          effect = "NoSchedule"
         }
       ]