Implement Resource Private State Management (#433)

* Add privatestate pkg and Data type (#399)

Reference: #399

* Fixing map merging and tests. Adding doc (#399)

Reference: #399

* Amend diagnostic messages and add logging (#399)

Reference: #399

* Setting up private data for use in ReadResource RPC (#399)

Reference: #399

* Refactoring and fixing tests (#399)

Reference: #399

* Call privatestate.NewData directly from fromproto5/6 (#399)

Reference: #399

* Adding tests for proto5/6server private (#399)

Reference: #399

* Additional tests for privatestate.NewData (#399)

Reference: #399

* Additional tests for privatestate (#399)

Reference: #399

* Calling Data.Bytes directly from toproto5/6/readresource (#399)

Reference: #399

* Adding test coverage to verify SetKey behaviour (#399)

Reference: #399

* Adding test coverage for fwserver readresource behaviour (#399)

Reference: #399

* Apply suggestions from code review

Co-authored-by: Brian Flad <[email protected]>

* Adding test coverage for proto5/6server readresource behaviour (#399)

Reference: #399

* Update resource/read.go

Co-authored-by: Brian Flad <[email protected]>

* Configuring private state data for use with create, update and delete (applyresourcechange) operations (#399)

Reference: #399

* Refactoring Data to use a pointer to a ProviderData struct which holds an unexported data field containing map[string][]byte to prevent direct manipulation of the data and inadvertent addition of values that are not invalid (i.e., invalid JSON and/or UTF-8) (#399)

Reference: #399

* Expanding test coverage around privatestate.Data (#399)

Reference: #399

* Ensuring that all instances of privatestate.Provider data that passed through to provider code have been initialised (#399)

Reference: #399

* Fixing tests following rebase (#399)

Reference: #399

* Implementing private state management for import resource state (#399)

Reference: #399

* Implementing private state management for ModifyAttribute and ModifyResource (#399)

Reference: #399

* Apply suggestions from code review

Co-authored-by: Brian Flad <[email protected]>

* Updating following code review (#399)

Reference: #399

* Apply suggestions from code review

Co-authored-by: Brian Flad <[email protected]>

* Updating following code review (#399)

Reference: #399

* Refactoring to use EmptyData to avoid Data.Provider from being nil (#399)

Reference: #399

* Updating changelog entry (#399)

Reference: #399

Co-authored-by: Brian Flad <[email protected]>

Author: bendbennett

.changelog/433.txt

@@ -0,0 +1,3 @@
+```release-note:note
+A new internal package has been introduced which enables provider developers to read/write framework-managed private state data.
+```

internal/fromproto5/applyresourcechange.go

@@ -3,11 +3,13 @@ package fromproto5
 import (
 	"context"
 
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
 )
 
 // ApplyResourceChangeRequest returns the *fwserver.ApplyResourceChangeRequest
@@ -34,7 +36,6 @@ func ApplyResourceChangeRequest(ctx context.Context, proto5 *tfprotov5.ApplyReso
 	}
 
 	fw := &fwserver.ApplyResourceChangeRequest{
-		PlannedPrivate: proto5.PlannedPrivate,
 		ResourceSchema: resourceSchema,
 		ResourceType:   resourceType,
 	}
@@ -63,5 +64,11 @@ func ApplyResourceChangeRequest(ctx context.Context, proto5 *tfprotov5.ApplyReso
 
 	fw.ProviderMeta = providerMeta
 
+	privateData, privateDataDiags := privatestate.NewData(ctx, proto5.PlannedPrivate)
+
+	diags.Append(privateDataDiags...)
+
+	fw.PlannedPrivate = privateData
+
 	return fw, diags
 }

internal/fromproto5/applyresourcechange_test.go

@@ -5,15 +5,17 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+	"github.com/hashicorp/terraform-plugin-go/tftypes"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fromproto5"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
 	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
 	"github.com/hashicorp/terraform-plugin-framework/types"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
-	"github.com/hashicorp/terraform-plugin-go/tftypes"
 )
 
 func TestApplyResourceChangeRequest(t *testing.T) {
@@ -44,6 +46,14 @@ func TestApplyResourceChangeRequest(t *testing.T) {
 		},
 	}
 
+	testProviderKeyValue := privatestate.MustMarshalToJson(map[string][]byte{
+		"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+	})
+
+	testProviderData := privatestate.MustProviderData(context.Background(), testProviderKeyValue)
+
+	testEmptyProviderData := privatestate.EmptyProviderData(context.Background())
+
 	testCases := map[string]struct {
 		input               *tfprotov5.ApplyResourceChangeRequest
 		resourceSchema      fwschema.Schema
@@ -125,14 +135,50 @@ func TestApplyResourceChangeRequest(t *testing.T) {
 				ResourceSchema: testFwSchema,
 			},
 		},
-		"plannedprivate": {
+		"plannedprivate-malformed-json": {
 			input: &tfprotov5.ApplyResourceChangeRequest{
-				PlannedPrivate: []byte("{}"),
+				PlannedPrivate: []byte(`{`),
 			},
 			resourceSchema: testFwSchema,
 			expected: &fwserver.ApplyResourceChangeRequest{
+				ResourceSchema: testFwSchema,
+			}, expectedDiagnostics: diag.Diagnostics{
+				diag.NewErrorDiagnostic(
+					"Error Decoding Private State",
+					"An error was encountered when decoding private state: unexpected end of JSON input.\n\n"+
+						"This is always a problem with Terraform or terraform-plugin-framework. Please report this to the provider developer.",
+				),
+			},
+		},
+		"plannedprivate-empty-json": {
+			input: &tfprotov5.ApplyResourceChangeRequest{
 				PlannedPrivate: []byte("{}"),
+			},
+			resourceSchema: testFwSchema,
+			expected: &fwserver.ApplyResourceChangeRequest{
 				ResourceSchema: testFwSchema,
+				PlannedPrivate: &privatestate.Data{
+					Framework: map[string][]byte{},
+					Provider:  testEmptyProviderData,
+				},
+			},
+		},
+		"plannedprivate": {
+			input: &tfprotov5.ApplyResourceChangeRequest{
+				PlannedPrivate: privatestate.MustMarshalToJson(map[string][]byte{
+					".frameworkKey":  []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+				}),
+			},
+			resourceSchema: testFwSchema,
+			expected: &fwserver.ApplyResourceChangeRequest{
+				ResourceSchema: testFwSchema,
+				PlannedPrivate: &privatestate.Data{
+					Framework: map[string][]byte{
+						".frameworkKey": []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					},
+					Provider: testProviderData,
+				},
 			},
 		},
 		"priorstate-missing-schema": {
@@ -209,7 +255,7 @@ func TestApplyResourceChangeRequest(t *testing.T) {
 
 			got, diags := fromproto5.ApplyResourceChangeRequest(context.Background(), testCase.input, testCase.resourceType, testCase.resourceSchema, testCase.providerMetaSchema)
 
-			if diff := cmp.Diff(got, testCase.expected); diff != "" {
+			if diff := cmp.Diff(got, testCase.expected, cmp.AllowUnexported(privatestate.ProviderData{})); diff != "" {
 				t.Errorf("unexpected difference: %s", diff)
 			}
 

internal/fromproto5/planresourcechange.go

@@ -3,11 +3,13 @@ package fromproto5
 import (
 	"context"
 
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
 )
 
 // PlanResourceChangeRequest returns the *fwserver.PlanResourceChangeRequest
@@ -34,7 +36,6 @@ func PlanResourceChangeRequest(ctx context.Context, proto5 *tfprotov5.PlanResour
 	}
 
 	fw := &fwserver.PlanResourceChangeRequest{
-		PriorPrivate:   proto5.PriorPrivate,
 		ResourceSchema: resourceSchema,
 		ResourceType:   resourceType,
 	}
@@ -63,5 +64,11 @@ func PlanResourceChangeRequest(ctx context.Context, proto5 *tfprotov5.PlanResour
 
 	fw.ProviderMeta = providerMeta
 
+	privateData, privateDataDiags := privatestate.NewData(ctx, proto5.PriorPrivate)
+
+	diags.Append(privateDataDiags...)
+
+	fw.PriorPrivate = privateData
+
 	return fw, diags
 }

internal/fromproto5/planresourcechange_test.go

@@ -5,15 +5,17 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+	"github.com/hashicorp/terraform-plugin-go/tftypes"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fromproto5"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
 	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
 	"github.com/hashicorp/terraform-plugin-framework/types"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
-	"github.com/hashicorp/terraform-plugin-go/tftypes"
 )
 
 func TestPlanResourceChangeRequest(t *testing.T) {
@@ -44,6 +46,12 @@ func TestPlanResourceChangeRequest(t *testing.T) {
 		},
 	}
 
+	testProviderKeyValue := privatestate.MustMarshalToJson(map[string][]byte{
+		"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+	})
+
+	testProviderData := privatestate.MustProviderData(context.Background(), testProviderKeyValue)
+
 	testCases := map[string]struct {
 		input               *tfprotov5.PlanResourceChangeRequest
 		resourceSchema      fwschema.Schema
@@ -99,11 +107,19 @@ func TestPlanResourceChangeRequest(t *testing.T) {
 		},
 		"priorprivate": {
 			input: &tfprotov5.PlanResourceChangeRequest{
-				PriorPrivate: []byte("{}"),
+				PriorPrivate: privatestate.MustMarshalToJson(map[string][]byte{
+					".frameworkKey":  []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+				}),
 			},
 			resourceSchema: testFwSchema,
 			expected: &fwserver.PlanResourceChangeRequest{
-				PriorPrivate:   []byte("{}"),
+				PriorPrivate: &privatestate.Data{
+					Framework: map[string][]byte{
+						".frameworkKey": []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					},
+					Provider: testProviderData,
+				},
 				ResourceSchema: testFwSchema,
 			},
 		},
@@ -209,7 +225,7 @@ func TestPlanResourceChangeRequest(t *testing.T) {
 
 			got, diags := fromproto5.PlanResourceChangeRequest(context.Background(), testCase.input, testCase.resourceType, testCase.resourceSchema, testCase.providerMetaSchema)
 
-			if diff := cmp.Diff(got, testCase.expected); diff != "" {
+			if diff := cmp.Diff(got, testCase.expected, cmp.AllowUnexported(privatestate.ProviderData{})); diff != "" {
 				t.Errorf("unexpected difference: %s", diff)
 			}
 

internal/fromproto5/readresource.go

@@ -3,11 +3,13 @@ package fromproto5
 import (
 	"context"
 
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
 )
 
 // ReadResourceRequest returns the *fwserver.ReadResourceRequest
@@ -20,7 +22,6 @@ func ReadResourceRequest(ctx context.Context, proto5 *tfprotov5.ReadResourceRequ
 	var diags diag.Diagnostics
 
 	fw := &fwserver.ReadResourceRequest{
-		Private:      proto5.Private,
 		ResourceType: resourceType,
 	}
 
@@ -36,5 +37,11 @@ func ReadResourceRequest(ctx context.Context, proto5 *tfprotov5.ReadResourceRequ
 
 	fw.ProviderMeta = providerMeta
 
+	privateData, privateDataDiags := privatestate.NewData(ctx, proto5.Private)
+
+	diags.Append(privateDataDiags...)
+
+	fw.Private = privateData
+
 	return fw, diags
 }

internal/fromproto5/readresource_test.go

@@ -5,15 +5,17 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+	"github.com/hashicorp/terraform-plugin-go/tftypes"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fromproto5"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
 	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
 	"github.com/hashicorp/terraform-plugin-framework/types"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
-	"github.com/hashicorp/terraform-plugin-go/tftypes"
 )
 
 func TestReadResourceRequest(t *testing.T) {
@@ -44,6 +46,14 @@ func TestReadResourceRequest(t *testing.T) {
 		},
 	}
 
+	testProviderKeyValue := privatestate.MustMarshalToJson(map[string][]byte{
+		"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+	})
+
+	testProviderData := privatestate.MustProviderData(context.Background(), testProviderKeyValue)
+
+	testEmptyProviderData := privatestate.EmptyProviderData(context.Background())
+
 	testCases := map[string]struct {
 		input               *tfprotov5.ReadResourceRequest
 		resourceSchema      fwschema.Schema
@@ -87,13 +97,47 @@ func TestReadResourceRequest(t *testing.T) {
 				},
 			},
 		},
-		"private": {
+		"private-malformed-json": {
+			input: &tfprotov5.ReadResourceRequest{
+				Private: []byte(`{`),
+			},
+			resourceSchema: testFwSchema,
+			expected:       &fwserver.ReadResourceRequest{},
+			expectedDiagnostics: diag.Diagnostics{
+				diag.NewErrorDiagnostic(
+					"Error Decoding Private State",
+					"An error was encountered when decoding private state: unexpected end of JSON input.\n\n"+
+						"This is always a problem with Terraform or terraform-plugin-framework. Please report this to the provider developer.",
+				),
+			},
+		},
+		"private-empty-json": {
 			input: &tfprotov5.ReadResourceRequest{
 				Private: []byte("{}"),
 			},
 			resourceSchema: testFwSchema,
 			expected: &fwserver.ReadResourceRequest{
-				Private: []byte("{}"),
+				Private: &privatestate.Data{
+					Framework: map[string][]byte{},
+					Provider:  testEmptyProviderData,
+				},
+			},
+		},
+		"private": {
+			input: &tfprotov5.ReadResourceRequest{
+				Private: privatestate.MustMarshalToJson(map[string][]byte{
+					".frameworkKey":  []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					"providerKeyOne": []byte(`{"pKeyOne": {"k0": "zero", "k1": 1}}`),
+				}),
+			},
+			resourceSchema: testFwSchema,
+			expected: &fwserver.ReadResourceRequest{
+				Private: &privatestate.Data{
+					Framework: map[string][]byte{
+						".frameworkKey": []byte(`{"fKeyOne": {"k0": "zero", "k1": 1}}`),
+					},
+					Provider: testProviderData,
+				},
 			},
 		},
 		"providermeta-missing-data": {
@@ -136,7 +180,7 @@ func TestReadResourceRequest(t *testing.T) {
 
 			got, diags := fromproto5.ReadResourceRequest(context.Background(), testCase.input, testCase.resourceType, testCase.resourceSchema, testCase.providerMetaSchema)
 
-			if diff := cmp.Diff(got, testCase.expected); diff != "" {
+			if diff := cmp.Diff(got, testCase.expected, cmp.AllowUnexported(privatestate.ProviderData{})); diff != "" {
 				t.Errorf("unexpected difference: %s", diff)
 			}
 

internal/fromproto6/applyresourcechange.go

@@ -3,11 +3,13 @@ package fromproto6
 import (
 	"context"
 
+	"github.com/hashicorp/terraform-plugin-go/tfprotov6"
+
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwschema"
 	"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
+	"github.com/hashicorp/terraform-plugin-framework/internal/privatestate"
 	"github.com/hashicorp/terraform-plugin-framework/provider"
-	"github.com/hashicorp/terraform-plugin-go/tfprotov6"
 )
 
 // ApplyResourceChangeRequest returns the *fwserver.ApplyResourceChangeRequest
@@ -34,7 +36,6 @@ func ApplyResourceChangeRequest(ctx context.Context, proto6 *tfprotov6.ApplyReso
 	}
 
 	fw := &fwserver.ApplyResourceChangeRequest{
-		PlannedPrivate: proto6.PlannedPrivate,
 		ResourceSchema: resourceSchema,
 		ResourceType:   resourceType,
 	}
@@ -63,5 +64,11 @@ func ApplyResourceChangeRequest(ctx context.Context, proto6 *tfprotov6.ApplyReso
 
 	fw.ProviderMeta = providerMeta
 
+	privateData, privateDataDiags := privatestate.NewData(ctx, proto6.PlannedPrivate)
+
+	diags.Append(privateDataDiags...)
+
+	fw.PlannedPrivate = privateData
+
 	return fw, diags
 }