fix issues with organization not being set after applies and add support for dont_delete policies

Author: SwiftEngineer

internal/provider/data_retention_policy.go

@@ -61,11 +61,31 @@ func modelFromTFEDataRetentionPolicyDontDelete(model modelTFEDataRetentionPolicy
 	}
 }
 
+func modelFromTFELegacyDataRetentionPolicy(ctx context.Context, model modelTFEDataRetentionPolicy, legacy *tfe.DataRetentionPolicy) (modelTFEDataRetentionPolicy, diag.Diagnostics) {
+	deleteOlderThan := modelTFEDeleteOlderThan{
+		Days: types.NumberValue(big.NewFloat(float64(legacy.DeleteOlderThanNDays))),
+	}
+	deleteOlderThanObject, diags := types.ObjectValueFrom(ctx, deleteOlderThan.AttributeTypes(), deleteOlderThan)
+
+	return modelTFEDataRetentionPolicy{
+		ID:              types.StringValue(legacy.ID),
+		Organization:    model.Organization,
+		WorkspaceId:     model.WorkspaceId,
+		DeleteOlderThan: deleteOlderThanObject,
+		DontDelete:      types.ObjectNull(map[string]attr.Type{}),
+	}, diags
+}
+
 func modelFromTFEDataRetentionPolicyChoice(ctx context.Context, model modelTFEDataRetentionPolicy, choice *tfe.DataRetentionPolicyChoice) (modelTFEDataRetentionPolicy, diag.Diagnostics) {
 	if choice.DataRetentionPolicyDeleteOlder != nil {
 		return modelFromTFEDataRetentionPolicyDeleteOlder(ctx, model, choice.DataRetentionPolicyDeleteOlder)
 	}
 
 	var emptyDiag []diag.Diagnostic
-	return modelFromTFEDataRetentionPolicyDontDelete(model, choice.DataRetentionPolicyDontDelete), emptyDiag
+	if choice.DataRetentionPolicyDontDelete != nil {
+		return modelFromTFEDataRetentionPolicyDontDelete(model, choice.DataRetentionPolicyDontDelete), emptyDiag
+	}
+
+	legacyPolicy := choice.ConvertToLegacyStruct()
+	return modelFromTFELegacyDataRetentionPolicy(ctx, model, legacyPolicy)
 }

internal/provider/resource_tfe_data_retention_policy.go

@@ -3,8 +3,6 @@ package provider
 import (
 	"context"
 	"fmt"
-	"strings"
-
 	"github.com/hashicorp/go-tfe"
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
@@ -16,14 +14,16 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework-validators/objectvalidator"
 	"github.com/hashicorp/terraform-plugin-framework/types/basetypes"
 	"github.com/hashicorp/terraform-plugin-framework/types"
+	"github.com/hashicorp/terraform-plugin-framework/diag"
+	"strings"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema/numberplanmodifier"
+	"github.com/hashicorp/terraform-plugin-framework-validators/numbervalidator"
 )
 
 // Ensure provider defined types fully satisfy framework interfaces.
 var _ resource.Resource = &resourceTFEDataRetentionPolicy{}
 var _ resource.ResourceWithConfigure = &resourceTFEDataRetentionPolicy{}
 var _ resource.ResourceWithImportState = &resourceTFEDataRetentionPolicy{}
-var _ resource.ResourceWithModifyPlan = &resourceTFEDataRetentionPolicy{}
 
 func NewDataRetentionPolicyResource() resource.Resource {
 	return &resourceTFEDataRetentionPolicy{}
@@ -38,10 +38,6 @@ func (r *resourceTFEDataRetentionPolicy) Metadata(ctx context.Context, req resou
 	resp.TypeName = req.ProviderTypeName + "_data_retention_policy"
 }
 
-func (r *resourceTFEDataRetentionPolicy) ModifyPlan(ctx context.Context, req resource.ModifyPlanRequest, resp *resource.ModifyPlanResponse) {
-	modifyPlanForDefaultOrganizationChange(ctx, r.config.Organization, req.State, req.Config, req.Plan, resp)
-}
-
 func (r *resourceTFEDataRetentionPolicy) Schema(ctx context.Context, req resource.SchemaRequest, resp *resource.SchemaResponse) {
 	resp.Schema = schema.Schema{
 		Description: "Manages the data retention policies for a specific workspace or an the entire organization.",
@@ -60,6 +56,7 @@ func (r *resourceTFEDataRetentionPolicy) Schema(ctx context.Context, req resourc
 				Optional:    true,
 				Computed:    true,
 				PlanModifiers: []planmodifier.String{
+					stringplanmodifier.UseStateForUnknown(),
 					stringplanmodifier.RequiresReplace(),
 				},
 			},
@@ -77,17 +74,17 @@ func (r *resourceTFEDataRetentionPolicy) Schema(ctx context.Context, req resourc
 				Attributes: map[string]schema.Attribute{
 					"days": schema.NumberAttribute{
 						Description: "Number of days",
-						Required:    true,
+						Optional:    true,
 						PlanModifiers: []planmodifier.Number{
 							numberplanmodifier.RequiresReplace(),
 						},
+						Validators: []validator.Number{
+							numbervalidator.ExactlyOneOf(
+								path.MatchRelative().AtParent().AtParent().AtName("dont_delete"),
+							),
+						},
 					},
 				},
-				Validators: []validator.Object{
-					objectvalidator.ExactlyOneOf(
-						path.MatchRelative().AtParent().AtName("dont_delete"),
-					),
-				},
 			},
 			"dont_delete": schema.SingleNestedBlock{
 				Attributes: map[string]schema.Attribute{},
@@ -128,11 +125,7 @@ func (r *resourceTFEDataRetentionPolicy) Create(ctx context.Context, req resourc
 		return
 	}
 
-	var organization string
-	if plan.WorkspaceId.IsNull() {
-		resp.Diagnostics.Append(r.config.dataOrDefaultOrganization(ctx, req.Plan, &organization)...)
-		plan.Organization = types.StringValue(organization)
-	}
+	r.ensureOrganizationIsSet(ctx, &plan, req.Plan, &resp.Diagnostics)
 
 	if resp.Diagnostics.HasError() {
 		return
@@ -150,6 +143,19 @@ func (r *resourceTFEDataRetentionPolicy) Create(ctx context.Context, req resourc
 
 }
 
+func (r *resourceTFEDataRetentionPolicy) ensureOrganizationIsSet(ctx context.Context, model *modelTFEDataRetentionPolicy, data AttrGettable, diags *diag.Diagnostics) {
+	if !model.Organization.IsUnknown() || model.Organization.ValueString() != "" {
+		// skip this method if the organization has already been set
+		return
+	}
+
+	if model.WorkspaceId.IsNull() {
+		var organization string
+		diags.Append(r.config.dataOrDefaultOrganization(ctx, data, &organization)...)
+		model.Organization = types.StringValue(organization)
+	}
+}
+
 func (r *resourceTFEDataRetentionPolicy) createDeleteOlderThanRetentionPolicy(ctx context.Context, plan modelTFEDataRetentionPolicy, resp *resource.CreateResponse) {
 	deleteOlderThan := &modelTFEDeleteOlderThan{}
 
@@ -183,8 +189,12 @@ func (r *resourceTFEDataRetentionPolicy) createDeleteOlderThanRetentionPolicy(ct
 		return
 	}
 
+	// set organization if it is still not known after creating the data retention policy
+	r.ensureOrganizationSetAfterApply(&result, &resp.Diagnostics)
+
 	// Save data into Terraform state
-	resp.Diagnostics.Append(resp.State.Set(ctx, &result)...)
+	diags = resp.State.Set(ctx, &result)
+	resp.Diagnostics.Append(diags...)
 }
 
 func (r *resourceTFEDataRetentionPolicy) createDontDeleteRetentionPolicy(ctx context.Context, plan modelTFEDataRetentionPolicy, resp *resource.CreateResponse) {
@@ -213,8 +223,23 @@ func (r *resourceTFEDataRetentionPolicy) createDontDeleteRetentionPolicy(ctx con
 
 	result := modelFromTFEDataRetentionPolicyDontDelete(plan, dataRetentionPolicy)
 
+	// set organization if it is still not known after creating the data retention policy
+	r.ensureOrganizationSetAfterApply(&result, &resp.Diagnostics)
+
 	// Save data into Terraform state
-	resp.Diagnostics.Append(resp.State.Set(ctx, &result)...)
+	diags = resp.State.Set(ctx, &result)
+	resp.Diagnostics.Append(diags...)
+}
+
+func (r *resourceTFEDataRetentionPolicy) ensureOrganizationSetAfterApply(policy *modelTFEDataRetentionPolicy, diags *diag.Diagnostics) {
+	if policy.Organization.IsUnknown() {
+		workspace, err := r.config.Client.Workspaces.ReadByID(ctx, policy.WorkspaceId.ValueString())
+		if err != nil {
+			diags.AddError("Unable to create data retention policy", err.Error())
+			return
+		}
+		policy.Organization = types.StringValue(workspace.Organization.Name)
+	}
 }
 
 func (r *resourceTFEDataRetentionPolicy) Read(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse) {
@@ -289,18 +314,50 @@ func (r *resourceTFEDataRetentionPolicy) Delete(ctx context.Context, req resourc
 }
 
 func (r *resourceTFEDataRetentionPolicy) ImportState(ctx context.Context, req resource.ImportStateRequest, resp *resource.ImportStateResponse) {
-	s := strings.SplitN(req.ID, "/", 2)
+	s := strings.Split(req.ID, "/")
+	if len(s) >= 3 {
+		resp.Diagnostics.AddError("Error importing workspace settings", fmt.Sprintf(
+			"invalid workspace input format: %s (expected <ORGANIZATION>/<WORKSPACE NAME> or <ORGANIZATION>)",
+			req.ID,
+		))
+	} else if len(s) == 2 {
+		workspaceID, err := fetchWorkspaceExternalID(s[0]+"/"+s[1], r.config.Client)
+		if err != nil {
+			resp.Diagnostics.AddError("Error importing data retention policy", fmt.Sprintf(
+				"error retrieving workspace with name %s from organization %s: %s", s[1], s[0], err.Error(),
+			))
+		}
 
-	if len(s) != 2 && len(s) != 1 {
-		resp.Diagnostics.AddError(
-			"Error importing variable",
-			fmt.Sprintf("Invalid variable import format: %s (expected <ORGANIZATION>/<WORKSPACE ID> or <ORGANIZATION>)", req.ID),
-		)
-		return
+		policy, err := r.config.Client.Workspaces.ReadDataRetentionPolicyChoice(ctx, workspaceID)
+		if err != nil {
+			resp.Diagnostics.AddError("Error importing data retention policy", fmt.Sprintf(
+				"error retrieving data policy for workspace %s from organization %s: %s", s[1], s[0], err.Error(),
+			))
+		}
+
+		req.ID = r.getPolicyID(policy)
+		resp.Diagnostics.Append(resp.State.SetAttribute(ctx, path.Root("workspace_id"), workspaceID)...)
+		resp.Diagnostics.Append(resp.State.SetAttribute(ctx, path.Root("organization"), s[0])...)
+	} else if len(s) == 1 {
+		policy, err := r.config.Client.Organizations.ReadDataRetentionPolicyChoice(ctx, s[0])
+		if err != nil {
+			resp.Diagnostics.AddError("Error importing data retention policy", fmt.Sprintf(
+				"error retrieving data policy for organization %s: %s", s[0], err.Error(),
+			))
+		}
+		req.ID = r.getPolicyID(policy)
+		resp.Diagnostics.Append(resp.State.SetAttribute(ctx, path.Root("organization"), s[0])...)
+	}
+}
+
+func (r *resourceTFEDataRetentionPolicy) getPolicyID(policy *tfe.DataRetentionPolicyChoice) string {
+	if policy.DataRetentionPolicyDeleteOlder != nil {
+		return policy.DataRetentionPolicyDeleteOlder.ID
+	}
+
+	if policy.DataRetentionPolicyDontDelete != nil {
+		return policy.DataRetentionPolicyDontDelete.ID
 	}
-	org := s[0]
-	wsId := s[1]
 
-	resp.Diagnostics.Append(resp.State.SetAttribute(ctx, path.Root("organization"), org)...)
-	resp.Diagnostics.Append(resp.State.SetAttribute(ctx, path.Root("workspace_id"), wsId)...)
+	return policy.ConvertToLegacyStruct().ID
 }