Unify upgradeResourceState between provider and provider2 (#2005)

This PR extracts the non provider2 specific part of
`planResourceChangeImpl.upgradeState` into it's own function:
`upgradeResourceStateGRPC`. This part is a pure refactor.

The PR then re-implements `upgradeResourceState` in terms of
`upgradeResourceStateGRPC`.

The driving motivation for this change is to reduce complexity and allow
both upgraders to handle large JSON numbers.

Replaces #1735

pkg/tfshim/sdk-v2/provider.go

@@ -124,7 +124,7 @@ func (p v2Provider) Apply(
 	if !ok {
 		return nil, fmt.Errorf("unknown resource %v", t)
 	}
-	state, err := upgradeResourceState(ctx, p.tf, r, stateFromShim(s))
+	state, err := upgradeResourceState(ctx, t, p.tf, r, stateFromShim(s))
 	if err != nil {
 		return nil, fmt.Errorf("failed to upgrade resource state: %w", err)
 	}
@@ -148,7 +148,7 @@ func (p v2Provider) Refresh(
 		return nil, fmt.Errorf("unknown resource %v", t)
 	}
 
-	state, err := upgradeResourceState(ctx, p.tf, r, stateFromShim(s))
+	state, err := upgradeResourceState(ctx, t, p.tf, r, stateFromShim(s))
 	if err != nil {
 		return nil, fmt.Errorf("failed to upgrade resource state: %w", err)
 	}

pkg/tfshim/sdk-v2/provider2.go

@@ -4,11 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"strconv"
 
 	"github.com/golang/glog"
 	"github.com/hashicorp/go-cty/cty"
-	ctyjson "github.com/hashicorp/go-cty/cty/json"
 	"github.com/hashicorp/go-cty/cty/msgpack"
 	"github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
@@ -330,61 +328,11 @@ func (p *planResourceChangeImpl) upgradeState(
 	res := p.tf.ResourcesMap[t]
 	state := p.unpackInstanceState(t, s)
 
-	// TODO[pulumi/pulumi-terraform-bridge#1667]: This is not quite right but we need
-	// the old TF state to get it right.
-	jsonBytes, err := ctyjson.Marshal(state.stateValue, state.stateValue.Type())
+	newState, newMeta, err := upgradeResourceStateGRPC(ctx, t, res, state.stateValue, state.meta, p.server.gserver)
 	if err != nil {
 		return nil, err
 	}
 
-	version := int64(0)
-	if versionValue, ok := state.meta["schema_version"]; ok {
-		versionString, ok := versionValue.(string)
-		if !ok {
-			return nil, fmt.Errorf("unexpected type %T for schema_version", versionValue)
-		}
-		v, err := strconv.ParseInt(versionString, 0, 32)
-		if err != nil {
-			return nil, err
-		}
-		version = v
-	}
-
-	//nolint:lll
-	// https://github.com/opentofu/opentofu/blob/2ef3047ec6bb266e8d91c55519967212c1a0975d/internal/tofu/upgrade_resource_state.go#L52
-	if version > int64(res.SchemaVersion) {
-		return nil, fmt.Errorf(
-			"State version %d is greater than schema version %d for resource %s. "+
-				"Please upgrade the provider to work with this resource.",
-			version, res.SchemaVersion, t,
-		)
-	}
-
-	// Note upgrade is always called, even if the versions match
-	//nolint:lll
-	// https://github.com/opentofu/opentofu/blob/2ef3047ec6bb266e8d91c55519967212c1a0975d/internal/tofu/upgrade_resource_state.go#L72
-
-	resp, err := p.server.gserver.UpgradeResourceState(ctx, &tfprotov5.UpgradeResourceStateRequest{
-		TypeName: t,
-		RawState: &tfprotov5.RawState{JSON: jsonBytes},
-		Version:  version,
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	newState, err := msgpack.Unmarshal(resp.UpgradedState.MsgPack, res.CoreConfigSchema().ImpliedType())
-	if err != nil {
-		return nil, err
-	}
-
-	newMeta := make(map[string]interface{}, len(state.meta))
-	// copy old meta into new meta
-	for k, v := range state.meta {
-		newMeta[k] = v
-	}
-	newMeta["schema_version"] = strconv.Itoa(res.SchemaVersion)
-
 	return &v2InstanceState2{
 		resourceType: t,
 		stateValue:   newState,

pkg/tfshim/sdk-v2/provider2_test.go

@@ -11,7 +11,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestUpgradeResourceState(t *testing.T) {
+func TestProvider2UpgradeResourceState(t *testing.T) {
 	const tfToken = "test_token"
 	for _, tc := range []struct {
 		name     string

pkg/tfshim/sdk-v2/provider_diff.go

@@ -66,7 +66,7 @@ func (p v2Provider) Diff(
 		}
 	} else {
 		// Upgrades are needed only if we have non-empty prior state.
-		state, err = upgradeResourceState(ctx, p.tf, r, state)
+		state, err = upgradeResourceState(ctx, t, p.tf, r, state)
 		if err != nil {
 			return nil, fmt.Errorf("failed to upgrade resource state: %w", err)
 		}

pkg/tfshim/sdk-v2/provider_test.go

@@ -0,0 +1,106 @@
+package sdkv2
+
+import (
+	"context"
+	"testing"
+
+	"github.com/hashicorp/go-cty/cty"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProvider1UpgradeResourceState(t *testing.T) {
+	t.Parallel()
+
+	type tc struct {
+		name   string
+		schema *schema.Resource
+		input  func() *terraform.InstanceState
+		expect func(t *testing.T, actual *terraform.InstanceState, tc tc)
+	}
+
+	tests := []tc{
+		{
+			name: "roundtrip int64",
+			schema: &schema.Resource{
+				UseJSONNumber: true,
+				Schema: map[string]*schema.Schema{
+					"x": {Type: schema.TypeInt, Optional: true},
+				},
+			},
+			input: func() *terraform.InstanceState {
+				n, err := cty.ParseNumberVal("641577219598130723")
+				require.NoError(t, err)
+				v := cty.ObjectVal(map[string]cty.Value{"x": n})
+				s := terraform.NewInstanceStateShimmedFromValue(v, 0)
+				s.Meta["schema_version"] = "0"
+				s.ID = "id"
+				s.RawState = v
+				s.Attributes["id"] = s.ID
+				return s
+			},
+			expect: func(t *testing.T, actual *terraform.InstanceState, tc tc) {
+				assert.Equal(t, tc.input().Attributes, actual.Attributes)
+			},
+		},
+		{
+			name: "type change",
+			schema: &schema.Resource{
+				Schema: map[string]*schema.Schema{
+					"x1": {Type: schema.TypeInt, Optional: true},
+				},
+				SchemaVersion: 1,
+				StateUpgraders: []schema.StateUpgrader{{
+					Version: 0,
+					Type: cty.Object(map[string]cty.Type{
+						"id": cty.String,
+						"x0": cty.String,
+					}),
+					Upgrade: func(_ context.Context, rawState map[string]any, _ interface{}) (map[string]any, error) {
+						return map[string]any{
+							"id": rawState["id"],
+							"x1": len(rawState["x0"].(string)),
+						}, nil
+					},
+				}},
+			},
+			input: func() *terraform.InstanceState {
+				s := terraform.NewInstanceStateShimmedFromValue(cty.ObjectVal(map[string]cty.Value{
+					"x0": cty.StringVal("123"),
+				}), 0)
+				s.Meta["schema_version"] = "0"
+				s.ID = "id"
+				return s
+			},
+			expect: func(t *testing.T, actual *terraform.InstanceState, tc tc) {
+				t.Logf("Actual = %#v", actual)
+				assert.Equal(t, map[string]string{
+					"id": "id",
+					"x1": "3",
+				}, actual.Attributes)
+			},
+		},
+	}
+
+	const tfToken = "test_token"
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := context.Background()
+
+			require.NoError(t, tt.schema.InternalValidate(tt.schema.Schema, true))
+
+			p := &schema.Provider{ResourcesMap: map[string]*schema.Resource{tfToken: tt.schema}}
+
+			actual, err := upgradeResourceState(ctx, tfToken, p, tt.schema, tt.input())
+			require.NoError(t, err)
+
+			tt.expect(t, actual, tt)
+		})
+	}
+}