feat: Add option to specify the assignment rules for scaling policy

castai/resource_workload_scaling_policy.go

@@ -47,6 +47,14 @@ const (
 	FieldApplyThresholdStrategyPercentageType      = "PERCENTAGE"
 	FieldApplyThresholdStrategyDefaultAdaptiveType = "DEFAULT_ADAPTIVE"
 	FieldApplyThresholdStrategyCustomAdaptiveType  = "CUSTOM_ADAPTIVE"
+	FieldAssignmentRules                           = "assignment_rules"
+)
+
+const (
+	K8sLabelInOperator           = "In"
+	K8sLabelNotInOperator        = "NotIn"
+	K8sLabelExistsOperator       = "Exists"
+	K8sLabelDoesNotExistOperator = "DoesNotExist"
 )
 
 var (
@@ -72,6 +80,63 @@ func resourceWorkloadScalingPolicy() *schema.Resource {
 				Description:      "CAST AI cluster id",
 				ValidateDiagFunc: validation.ToDiagFunc(validation.IsUUID),
 			},
+			FieldAssignmentRules: {
+				Type:        schema.TypeList,
+				Optional:    true,
+				Description: "Allows defining conditions for automatically assigning workloads to this scaling policy.",
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"rules": {
+							Type:     schema.TypeList,
+							Required: true,
+							Elem: &schema.Resource{
+								Schema: map[string]*schema.Schema{
+									"namespace": {
+										Type:        schema.TypeList,
+										Optional:    true,
+										MaxItems:    1,
+										Description: "Allows assigning a scaling policy based on the workload's namespace.",
+										Elem: &schema.Resource{
+											Schema: map[string]*schema.Schema{
+												"names": {
+													Type:        schema.TypeList,
+													Optional:    true,
+													Description: "Defines matching by namespace names.",
+													Elem:        &schema.Schema{Type: schema.TypeString},
+												},
+												// TODO(WOOP-714): enable label expressions
+												//"labels_expressions": k8sLabelExpressionsSchema(),
+											},
+										},
+									},
+									"workload": {
+										Type:        schema.TypeList,
+										Optional:    true,
+										MaxItems:    1,
+										Description: "Allows assigning a scaling policy based on the workload's metadata.",
+										Elem: &schema.Resource{
+											Schema: map[string]*schema.Schema{
+												"gvk": {
+													Type:     schema.TypeList,
+													Optional: true,
+													Description: `Group, version, and kind for Kubernetes resources. Format: kind[.version][.group].
+It can be either:
+ - only kind, e.g. "Deployment"
+ - group and kind: e.g."Deployment.apps"
+ - group, version and kind: e.g."Deployment.v1.apps"`,
+
+													Elem: &schema.Schema{Type: schema.TypeString},
+												},
+												"labels_expressions": k8sLabelExpressionsSchema(),
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
 			"name": {
 				Type:             schema.TypeString,
 				Required:         true,
@@ -200,6 +265,37 @@ func resourceWorkloadScalingPolicy() *schema.Resource {
 	}
 }
 
+func k8sLabelExpressionsSchema() *schema.Schema {
+	return &schema.Schema{
+		Type:        schema.TypeList,
+		Optional:    true,
+		Description: "Defines matching by label selector requirements.",
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				"key": {
+					Type:        schema.TypeString,
+					Required:    true,
+					Description: "The label key to match.",
+				},
+				"operator": {
+					Type:        schema.TypeString,
+					Required:    true,
+					Description: "The operator to use for matching the label.",
+					ValidateDiagFunc: validation.ToDiagFunc(validation.StringInSlice([]string{
+						K8sLabelInOperator, K8sLabelNotInOperator, K8sLabelExistsOperator, K8sLabelDoesNotExistOperator,
+					}, false)),
+				},
+				"values": {
+					Type:        schema.TypeList,
+					Optional:    true,
+					Description: "A list of values to match against the label key. Allowed for `In` and `NotIn` operators.",
+					Elem:        &schema.Schema{Type: schema.TypeString},
+				},
+			},
+		},
+	}
+}
+
 func workloadScalingPolicyResourceSchema(resource, function string, overhead, minRecommended float64) *schema.Resource {
 	return &schema.Resource{
 		Schema: map[string]*schema.Schema{
@@ -414,6 +510,12 @@ func resourceWorkloadScalingPolicyCreate(ctx context.Context, d *schema.Resource
 
 	req.RecommendationPolicies.AntiAffinity = toAntiAffinity(toSection(d, "anti_affinity"))
 
+	ar, err := toAssignmentRules(toSection(d, FieldAssignmentRules))
+	if err != nil {
+		return diag.FromErr(err)
+	}
+	req.AssignmentRules = ar
+
 	create, err := client.WorkloadOptimizationAPICreateWorkloadScalingPolicyWithResponse(ctx, clusterID, req)
 	if err != nil {
 		return diag.FromErr(err)
@@ -489,6 +591,10 @@ func resourceWorkloadScalingPolicyRead(ctx context.Context, d *schema.ResourceDa
 		return diag.FromErr(fmt.Errorf("setting anti-affinity: %w", err))
 	}
 
+	if err := d.Set(FieldAssignmentRules, toAssignmentRulesMap(getResourceFrom(d, FieldAssignmentRules), sp.AssignmentRules)); err != nil {
+		return diag.FromErr(fmt.Errorf("setting assignment rules: %w", err))
+	}
+
 	return nil
 }
 
@@ -511,6 +617,7 @@ func resourceWorkloadScalingPolicyUpdate(ctx context.Context, d *schema.Resource
 		"memory_event",
 		"anti_affinity",
 		FieldConfidence,
+		FieldAssignmentRules,
 	) {
 		tflog.Info(ctx, "scaling policy up to date")
 		return nil
@@ -526,9 +633,15 @@ func resourceWorkloadScalingPolicyUpdate(ctx context.Context, d *schema.Resource
 	if err != nil {
 		return diag.FromErr(err)
 	}
+	ar, err := toAssignmentRules(toSection(d, FieldAssignmentRules))
+	if err != nil {
+		return diag.FromErr(err)
+	}
+
 	req := sdk.WorkloadOptimizationAPIUpdateWorkloadScalingPolicyJSONRequestBody{
-		Name:      d.Get("name").(string),
-		ApplyType: sdk.WorkloadoptimizationV1ApplyType(d.Get("apply_type").(string)),
+		Name:            d.Get("name").(string),
+		ApplyType:       sdk.WorkloadoptimizationV1ApplyType(d.Get("apply_type").(string)),
+		AssignmentRules: ar,
 		RecommendationPolicies: sdk.WorkloadoptimizationV1RecommendationPolicies{
 			ManagementOption: sdk.WorkloadoptimizationV1ManagementOption(d.Get("management_option").(string)),
 			Cpu:              cpu,
@@ -1058,3 +1171,214 @@ func getWorkloadScalingPolicyByName(ctx context.Context, client sdk.ClientWithRe
 	}
 	return nil, fmt.Errorf("policy with name %q not found", name)
 }
+
+func toAssignmentRules(in map[string]any) (*[]sdk.WorkloadoptimizationV1ScalingPolicyAssignmentRule, error) {
+	if len(in) == 0 {
+		return nil, nil
+	}
+
+	rules, ok := in["rules"].([]any)
+	if !ok || len(rules) == 0 {
+		return &[]sdk.WorkloadoptimizationV1ScalingPolicyAssignmentRule{}, nil
+	}
+
+	result := make([]sdk.WorkloadoptimizationV1ScalingPolicyAssignmentRule, len(rules))
+	for i, rule := range rules {
+		ruleMap := rule.(map[string]any)
+
+		nsRules, err := toNamespaceAssignmentRule(ruleMap)
+		if err != nil {
+			return nil, fmt.Errorf("field assignment_rules[%d].namespace: %w", i, err)
+		}
+
+		wRules, err := toWorkloadAssignmentRule(ruleMap)
+		if err != nil {
+			return nil, fmt.Errorf("field assignment_rules[%d].workload: %w", i, err)
+		}
+		if nsRules == nil && wRules == nil {
+			continue
+		}
+		result[i] = sdk.WorkloadoptimizationV1ScalingPolicyAssignmentRule{
+			Workload:  wRules,
+			Namespace: nsRules,
+		}
+	}
+
+	return &result, nil
+}
+
+func toNamespaceAssignmentRule(ruleMap map[string]any) (*sdk.WorkloadoptimizationV1KubernetesNamespaceMatcher, error) {
+	namespaceMap := getFirstElem(ruleMap, "namespace")
+	if namespaceMap == nil {
+		return nil, nil
+	}
+
+	namespaceMatcher := &sdk.WorkloadoptimizationV1KubernetesNamespaceMatcher{}
+
+	if names := readOptionalValue[[]any](namespaceMap, "names"); names != nil {
+		namespaceMatcher.Names = lo.ToPtr(toStringList(*names))
+	}
+
+	labels, err := toKubernetesLabelExpressionMatcher(namespaceMap)
+	if err != nil {
+		return nil, err
+	}
+	namespaceMatcher.LabelsExpressions = labels
+
+	return namespaceMatcher, nil
+}
+
+func toKubernetesLabelExpressionMatcher(namespaceMap map[string]any) (*[]sdk.WorkloadoptimizationV1KubernetesLabelExpressionMatcher, error) {
+	exprs := readOptionalValue[[]any](namespaceMap, "labels_expressions")
+	if exprs == nil {
+		return nil, nil
+	}
+	expressions := make([]sdk.WorkloadoptimizationV1KubernetesLabelExpressionMatcher, len(*exprs))
+	for j, expr := range *exprs {
+		exprMap := expr.(map[string]any)
+		key, err := mustGetValue[string](exprMap, "key")
+		if err != nil {
+			return nil, err
+		}
+
+		operator, err := mustGetValue[string](exprMap, "operator")
+		if err != nil {
+			return nil, err
+		}
+
+		expressions[j] = sdk.WorkloadoptimizationV1KubernetesLabelExpressionMatcher{
+			Key:      *key,
+			Operator: toLabelSelectorOperator(*operator),
+		}
+
+		values := readOptionalValue[[]any](exprMap, "values")
+		if values != nil {
+			expressions[j].Values = toStringList(*values)
+		}
+	}
+	return &expressions, nil
+}
+
+func toWorkloadAssignmentRule(ruleMap map[string]any) (*sdk.WorkloadoptimizationV1KubernetesWorkloadMatcher, error) {
+	workloadMap := getFirstElem(ruleMap, "workload")
+	if workloadMap == nil {
+		return nil, nil
+	}
+
+	workloadMatcher := &sdk.WorkloadoptimizationV1KubernetesWorkloadMatcher{}
+
+	if gvk := readOptionalValue[[]any](workloadMap, "gvk"); gvk != nil {
+		workloadMatcher.Gvk = lo.ToPtr(toStringList(*gvk))
+	}
+
+	labels, err := toKubernetesLabelExpressionMatcher(workloadMap)
+	if err != nil {
+		return nil, err
+	}
+	workloadMatcher.LabelsExpressions = labels
+
+	return workloadMatcher, nil
+}
+
+func getFirstElem(in map[string]any, key string) map[string]any {
+	list, ok := in[key].([]any)
+	if !ok || len(list) == 0 {
+		return nil
+	}
+	elem, ok := list[0].(map[string]any)
+	if !ok {
+		return nil
+	}
+	return elem
+}
+
+func toLabelSelectorOperator(in string) sdk.WorkloadoptimizationV1KubernetesLabelSelectorOperator {
+	switch in {
+	case K8sLabelInOperator:
+		return sdk.KUBERNETESLABELSELECTOROPIN
+	case K8sLabelNotInOperator:
+		return sdk.KUBERNETESLABELSELECTOROPNOTIN
+	case K8sLabelExistsOperator:
+		return sdk.KUBERNETESLABELSELECTOROPEXISTS
+	case K8sLabelDoesNotExistOperator:
+		return sdk.KUBERNETESLABELSELECTOROPDOESNOTEXIST
+	}
+	return sdk.KUBERNETESLABELSELECTOROPUNSPECIFIED
+}
+
+func toAssignmentRulesMap(previous map[string]any, rules *[]sdk.WorkloadoptimizationV1ScalingPolicyAssignmentRule) []any {
+	if rules == nil {
+		return nil
+	}
+
+	if len(*rules) == 0 && len(previous) == 0 {
+		return nil
+	}
+
+	result := make([]map[string]any, len(*rules))
+	for i, rule := range *rules {
+		ruleMap := make(map[string]any)
+
+		if rule.Namespace != nil {
+			namespaceMap := make(map[string]any)
+
+			if rule.Namespace.Names != nil {
+				namespaceMap["names"] = *rule.Namespace.Names
+			}
+
+			if rule.Namespace.LabelsExpressions != nil && len(*rule.Namespace.LabelsExpressions) > 0 {
+				namespaceMap["labels_expressions"] = toK8sLabelsExpressionsMap(*rule.Namespace.LabelsExpressions)
+			}
+
+			ruleMap["namespace"] = []map[string]any{namespaceMap}
+		}
+
+		if rule.Workload != nil {
+			workloadMap := make(map[string]any)
+
+			if rule.Workload.Gvk != nil && len(*rule.Workload.Gvk) > 0 {
+				workloadMap["gvk"] = *rule.Workload.Gvk
+			}
+
+			if rule.Workload.LabelsExpressions != nil && len(*rule.Workload.LabelsExpressions) > 0 {
+				workloadMap["labels_expressions"] = toK8sLabelsExpressionsMap(*rule.Workload.LabelsExpressions)
+			}
+
+			ruleMap["workload"] = []map[string]any{workloadMap}
+		}
+
+		result[i] = ruleMap
+	}
+
+	return []any{
+		map[string]any{
+			"rules": result,
+		},
+	}
+}
+
+func toK8sLabelsExpressionsMap(in []sdk.WorkloadoptimizationV1KubernetesLabelExpressionMatcher) []map[string]any {
+	expressions := make([]map[string]any, len(in))
+	for j, expr := range in {
+		expressions[j] = map[string]any{
+			"key":      expr.Key,
+			"operator": labelSelectorOperatorMap(expr.Operator),
+			"values":   expr.Values,
+		}
+	}
+	return expressions
+}
+
+func labelSelectorOperatorMap(in sdk.WorkloadoptimizationV1KubernetesLabelSelectorOperator) string {
+	switch in {
+	case sdk.KUBERNETESLABELSELECTOROPIN:
+		return K8sLabelInOperator
+	case sdk.KUBERNETESLABELSELECTOROPNOTIN:
+		return K8sLabelNotInOperator
+	case sdk.KUBERNETESLABELSELECTOROPEXISTS:
+		return K8sLabelExistsOperator
+	case sdk.KUBERNETESLABELSELECTOROPDOESNOTEXIST:
+		return K8sLabelDoesNotExistOperator
+	}
+	return "unspecified"
+}