feat: Add node affinity to Node Template

castai/resource_node_template.go

@@ -56,6 +56,9 @@ const (
 	FieldNodeTemplateStorageOptimized                         = "storage_optimized"
 	FieldNodeTemplateUseSpotFallbacks                         = "use_spot_fallbacks"
 	FieldNodeTemplateCustomPriority                           = "custom_priority"
+	FieldNodeTemplateNodeAffinity                             = "node_affinity"
+	FieldNodeTemplateAzName                                   = "az_name"
+	FieldNodeTemplateInstanceTypes                            = "instance_types"
 )
 
 const (
@@ -345,6 +348,32 @@ func resourceNodeTemplate() *schema.Resource {
 								},
 							},
 						},
+						FieldNodeTemplateNodeAffinity: {
+							Type:     schema.TypeList,
+							Optional: true,
+							Elem: &schema.Resource{
+								Schema: map[string]*schema.Schema{
+									FieldNodeTemplateInstanceTypes: {
+										Type:     schema.TypeList,
+										Required: true,
+										Elem: &schema.Schema{
+											Type: schema.TypeString,
+										},
+										Description: "Instance types in this node group.",
+									},
+									FieldNodeTemplateAzName: {
+										Required:    true,
+										Type:        schema.TypeString,
+										Description: "Availability zone name.",
+									},
+									FieldNodeTemplateName: {
+										Required:    true,
+										Type:        schema.TypeString,
+										Description: "Name of node group.",
+									},
+								},
+							},
+						},
 					},
 				},
 			},
@@ -489,6 +518,9 @@ func flattenConstraints(c *sdk.NodetemplatesV1TemplateConstraints) ([]map[string
 	if c.CustomPriority != nil && len(*c.CustomPriority) > 0 {
 		out[FieldNodeTemplateCustomPriority] = flattenCustomPriority(*c.CustomPriority)
 	}
+	if c.NodeAffinity != nil && len(*c.NodeAffinity) > 0 {
+		out[FieldNodeTemplateNodeAffinity] = flattenNodeAffinity(*c.NodeAffinity)
+	}
 	if c.InstanceFamilies != nil {
 		out[FieldNodeTemplateInstanceFamilies] = flattenInstanceFamilies(c.InstanceFamilies)
 	}
@@ -596,6 +628,19 @@ func flattenCustomPriority(priorities []sdk.NodetemplatesV1TemplateConstraintsCu
 	})
 }
 
+func flattenNodeAffinity(affinities []sdk.NodetemplatesV1TemplateConstraintsNodeAffinity) any {
+	return lo.Map(affinities, func(item sdk.NodetemplatesV1TemplateConstraintsNodeAffinity, index int) map[string]any {
+		result := map[string]any{}
+		if item.InstanceTypes != nil {
+			result[FieldNodeTemplateInstanceTypes] = *item.InstanceTypes
+		}
+
+		result[FieldNodeTemplateName] = lo.FromPtr(item.Name)
+		result[FieldNodeTemplateAzName] = lo.FromPtr(item.AzName)
+		return result
+	})
+}
+
 func resourceNodeTemplateDelete(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*ProviderConfig).api
 	clusterID := d.Get(FieldClusterID).(string)
@@ -1019,6 +1064,21 @@ func toTemplateConstraints(obj map[string]any) *sdk.NodetemplatesV1TemplateConst
 			}))
 		}
 	}
+	if v, ok := obj[FieldNodeTemplateNodeAffinity].([]any); ok && len(v) > 0 {
+		if ok {
+			out.NodeAffinity = lo.ToPtr(lo.FilterMap(v, func(item any, _ int) (sdk.NodetemplatesV1TemplateConstraintsNodeAffinity, bool) {
+				val, ok := item.(map[string]any)
+				if !ok {
+					return sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{}, false
+				}
+				res := toTemplateConstraintsNodeAffinity(val)
+				if res == nil {
+					return sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{}, false
+				}
+				return *res, true
+			}))
+		}
+	}
 
 	return out
 }
@@ -1083,3 +1143,22 @@ func toTemplateConstraintsCustomPriority(o map[string]any) *sdk.NodetemplatesV1T
 
 	return &out
 }
+
+func toTemplateConstraintsNodeAffinity(o map[string]any) *sdk.NodetemplatesV1TemplateConstraintsNodeAffinity {
+	if o == nil {
+		return nil
+	}
+
+	out := sdk.NodetemplatesV1TemplateConstraintsNodeAffinity{}
+	if v, ok := o[FieldNodeTemplateName].(string); ok {
+		out.Name = toPtr(v)
+	}
+	if v, ok := o[FieldNodeTemplateAzName].(string); ok {
+		out.AzName = toPtr(v)
+	}
+	if v, ok := o[FieldNodeTemplateInstanceTypes].([]any); ok {
+		out.InstanceTypes = toPtr(toStringList(v))
+	}
+
+	return &out
+}

castai/resource_node_template_test.go

@@ -87,6 +87,13 @@ func TestNodeTemplateResourceReadContext(t *testing.T) {
 				  		"spot": true,
 				  		"onDemand": true
 				  	}
+				  ],
+				  "nodeAffinity": [
+				    {	
+						"name": "foo",
+						"azName": "eu-central-1a",
+						"instanceTypes": ["m5.24xlarge"]
+					}
 				  ]
 				},
 				"version": "3",
@@ -171,6 +178,11 @@ constraints.0.max_cpu = 10000
 constraints.0.max_memory = 0
 constraints.0.min_cpu = 10
 constraints.0.min_memory = 0
+constraints.0.node_affinity.# = 1
+constraints.0.node_affinity.0.az_name = eu-central-1a
+constraints.0.node_affinity.0.instance_types.# = 1
+constraints.0.node_affinity.0.instance_types.0 = m5.24xlarge
+constraints.0.node_affinity.0.name = foo
 constraints.0.on_demand = true
 constraints.0.os.# = 1
 constraints.0.os.0 = linux
@@ -434,6 +446,11 @@ func TestAccResourceNodeTemplate_basic(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.instance_families.1", "d"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.spot", "true"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.0.on_demand", "true"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.name", "foo"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.az_name", "eu-central-1a"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.instance_types.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.instance_types.0", "m5.24xlarge"),
 				),
 			},
 			{
@@ -491,6 +508,15 @@ func TestAccResourceNodeTemplate_basic(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.instance_families.1", "d"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.spot", "true"),
 					resource.TestCheckResourceAttr(resourceName, "constraints.0.custom_priority.1.on_demand", "true"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.#", "2"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.name", "foo"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.az_name", "eu-central-1a"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.instance_types.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.0.instance_types.0", "m5.24xlarge"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.1.name", "bar"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.1.az_name", "eu-central-1b"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.1.instance_types.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "constraints.0.node_affinity.1.instance_types.0", "m5.24xlarge"),
 				),
 			},
 		},
@@ -563,6 +589,12 @@ func testAccNodeTemplateConfig(rName, clusterName string) string {
 					spot = true
 					on_demand = true
 				}
+
+				node_affinity {
+					name = "foo"
+					az_name = "eu-central-1a"
+					instance_types = ["m5.24xlarge"]
+				}
 			}
 		}
 	`, rName))
@@ -614,6 +646,17 @@ func testNodeTemplateUpdated(rName, clusterName string) string {
 					spot = true
 					on_demand = true
 				}
+
+				node_affinity {
+					name = "foo"
+					az_name = "eu-central-1a"
+					instance_types = ["m5.24xlarge"]
+				}
+				node_affinity {
+					name = "bar"
+					az_name = "eu-central-1b"
+					instance_types = ["m5.24xlarge"]
+				}
 			}
 		}
 	`, rName))