Pass PlannedPrivate from PlanResourceChange to ApplyResourceChange (#2407)

The `PlannedPrivate` field of the PRC response includes metadata that's
required for correctly re-assembling the diff.
For example, it includes `NewExtra` which includes the original value
for attributes with a `StateFunc`.
Without `PlannedPrivate`, those attributes get saved to state encoded
twice. This can lead to mangled data in state and perma-diffs.

Fixes #2408

Author: flostadler

pkg/tests/cross-tests/input_cross_test.go

@@ -601,3 +601,28 @@ func TestCreateDoesNotPanicWithStateUpgraders(t *testing.T) {
 		},
 	})
 }
+
+// TestStateFunc ensures that resources with a StateFunc set on their schema are correctly handled. This includes
+// ensuring that the PlannedPrivate blob is passed from PlanResourceChange to ApplyResourceChange. If this is passed
+// correctly, the provider will see the original value of the field, rather than the value that was produced by the StateFunc.
+func TestStateFunc(t *testing.T) {
+	input := "hello"
+
+	runCreateInputCheck(t, inputTestCase{
+		Resource: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				"test": {
+					Type:     schema.TypeString,
+					Optional: true,
+					ForceNew: true,
+					StateFunc: func(v interface{}) string {
+						return v.(string) + " world"
+					},
+				},
+			},
+		},
+		Config: map[string]any{
+			"test": input,
+		},
+	})
+}

pkg/tests/schema_pulumi_test.go

@@ -3090,3 +3090,47 @@ resource "prov_test" "mainRes" {
 		})
 	}
 }
+
+func TestStateFunc(t *testing.T) {
+	resMap := map[string]*schema.Resource{
+		"prov_test": {
+			CreateContext: func(ctx context.Context, d *schema.ResourceData, i interface{}) diag.Diagnostics {
+				d.SetId("id")
+				var diags diag.Diagnostics
+				v, ok := d.GetOk("test")
+				assert.True(t, ok, "test property not set")
+
+				err := d.Set("test", v.(string)+" world")
+				require.NoError(t, err)
+				return diags
+			},
+			Schema: map[string]*schema.Schema{
+				"test": {
+					Type:     schema.TypeString,
+					Optional: true,
+					ForceNew: true,
+					StateFunc: func(v interface{}) string {
+						return v.(string) + " world"
+					},
+				},
+			},
+		},
+	}
+	tfp := &schema.Provider{ResourcesMap: resMap}
+	bridgedProvider := pulcheck.BridgedProvider(t, "prov", tfp)
+	program := `
+name: test
+runtime: yaml
+resources:
+  mainRes:
+    type: prov:index:Test
+	properties:
+	  test: "hello"
+outputs:
+  testOut: ${mainRes.test}
+`
+	pt := pulcheck.PulCheck(t, bridgedProvider, program)
+	res := pt.Up()
+	require.Equal(t, "hello world", res.Outputs["testOut"].Value)
+	pt.Preview(optpreview.ExpectNoChanges())
+}

pkg/tfshim/sdk-v2/provider2.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"maps"
 	"sort"
 	"strings"
 	"time"
@@ -111,8 +112,9 @@ func (s *v2InstanceState2) Meta() map[string]interface{} {
 type v2InstanceDiff2 struct {
 	v2InstanceDiff
 
-	config       cty.Value
-	plannedState cty.Value
+	config         cty.Value
+	plannedState   cty.Value
+	plannedPrivate map[string]interface{}
 }
 
 func (d *v2InstanceDiff2) String() string {
@@ -129,7 +131,8 @@ func (d *v2InstanceDiff2) GoString() string {
     },
     config:         %#v,
     plannedState:   %#v,
-}`, d.v2InstanceDiff.tf, d.config, d.plannedState)
+    plannedPrivate:    %#v,
+}`, d.v2InstanceDiff.tf, d.config, d.plannedState, d.plannedPrivate)
 }
 
 var _ shim.InstanceDiff = (*v2InstanceDiff2)(nil)
@@ -258,12 +261,14 @@ func (p *planResourceChangeImpl) Diff(
 	if err != nil {
 		return nil, err
 	}
+
 	return &v2InstanceDiff2{
 		v2InstanceDiff: v2InstanceDiff{
 			tf: plan.PlannedDiff,
 		},
-		config:       cfg,
-		plannedState: plan.PlannedState,
+		config:         cfg,
+		plannedState:   plan.PlannedState,
+		plannedPrivate: plan.PlannedPrivate,
 	}, nil
 }
 
@@ -283,7 +288,17 @@ func (p *planResourceChangeImpl) Apply(
 	}
 	diff := p.unpackDiff(ty, d)
 	cfg, st, pl := diff.config, state.stateValue, diff.plannedState
-	priv := diff.v2InstanceDiff.tf.Meta
+
+	// Merge plannedPrivate and v2InstanceDiff.tf.Meta into a single map. This is necessary because
+	// timeouts are stored in the Meta and not in plannedPrivate.
+	priv := make(map[string]interface{})
+	if len(diff.plannedPrivate) > 0 {
+		maps.Copy(priv, diff.plannedPrivate)
+	}
+	if len(diff.v2InstanceDiff.tf.Meta) > 0 {
+		maps.Copy(priv, diff.v2InstanceDiff.tf.Meta)
+	}
+
 	resp, err := p.server.ApplyResourceChange(ctx, t, ty, cfg, st, pl, priv, meta)
 	if err != nil {
 		return nil, err
@@ -475,9 +490,9 @@ func (s *grpcServer) PlanResourceChange(
 	providerMeta *cty.Value,
 	ignores shim.IgnoreChanges,
 ) (*struct {
-	PlannedState cty.Value
-	PlannedMeta  map[string]interface{}
-	PlannedDiff  *terraform.InstanceDiff
+	PlannedState   cty.Value
+	PlannedPrivate map[string]interface{}
+	PlannedDiff    *terraform.InstanceDiff
 }, error) {
 	configVal, err := msgpack.Marshal(config, ty)
 	if err != nil {
@@ -548,13 +563,13 @@ func (s *grpcServer) PlanResourceChange(
 		}
 	}
 	return &struct {
-		PlannedState cty.Value
-		PlannedMeta  map[string]interface{}
-		PlannedDiff  *terraform.InstanceDiff
+		PlannedState   cty.Value
+		PlannedPrivate map[string]interface{}
+		PlannedDiff    *terraform.InstanceDiff
 	}{
-		PlannedState: plannedState,
-		PlannedMeta:  meta,
-		PlannedDiff:  resp.InstanceDiff,
+		PlannedState:   plannedState,
+		PlannedPrivate: meta,
+		PlannedDiff:    resp.InstanceDiff,
 	}, nil
 }