Switch upgradeResourceState to use SDKv2's gRPC method

pkg/tfshim/sdk-v2/diagnostics.go

@@ -30,7 +30,7 @@ func warningsAndErrors(diags diag.Diagnostics) ([]string, []error) {
 	return warnings, errors
 }
 
-func errors(diags diag.Diagnostics) error {
+func diagToError(diags diag.Diagnostics) error {
 	var err error
 	for _, d := range diags {
 		if d.Severity == diag.Error {

pkg/tfshim/sdk-v2/provider.go

@@ -103,7 +103,7 @@ func (p v2Provider) Configure(ctx context.Context, c shim.ResourceConfig) error
 	// See ConfigureProvider e.g.
 	// https://github.com/hashicorp/terraform-plugin-sdk/blob/main/helper/schema/grpc_provider.go#L564
 	ctxHack := context.WithValue(ctx, schema.StopContextKey, p.stopContext(context.Background()))
-	return errors(p.tf.Configure(ctxHack, configFromShim(c)))
+	return diagToError(p.tf.Configure(ctxHack, configFromShim(c)))
 }
 
 func (p v2Provider) stopContext(ctx context.Context) context.Context {
@@ -124,12 +124,12 @@ 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)
 	}
 	state, diags := r.Apply(ctx, state, diffFromShim(d), p.tf.Meta())
-	return stateToShim(r, state), errors(diags)
+	return stateToShim(r, state), diagToError(diags)
 }
 
 func (p v2Provider) Refresh(
@@ -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)
 	}
@@ -158,7 +158,7 @@ func (p v2Provider) Refresh(
 	}
 
 	state, diags := r.RefreshWithoutUpgrade(context.TODO(), state, p.tf.Meta())
-	return stateToShim(r, state), errors(diags)
+	return stateToShim(r, state), diagToError(diags)
 }
 
 func (p v2Provider) ReadDataDiff(
@@ -195,7 +195,7 @@ func (p v2Provider) ReadDataApply(
 		return nil, fmt.Errorf("unknown resource %v", t)
 	}
 	state, diags := r.ReadDataApply(ctx, diffFromShim(d), p.tf.Meta())
-	return stateToShim(r, state), errors(diags)
+	return stateToShim(r, state), diagToError(diags)
 }
 
 func (p v2Provider) Meta(_ context.Context) interface{} {

pkg/tfshim/sdk-v2/provider2.go

@@ -331,7 +331,7 @@ func (p *planResourceChangeImpl) upgradeState(
 		instanceState.Attributes = map[string]string{}
 	}
 	instanceState.Meta = state.meta
-	newInstanceState, err := upgradeResourceState(ctx, p.tf, res, instanceState)
+	newInstanceState, err := upgradeResourceState(ctx, t, p.tf, res, instanceState)
 	if err != nil {
 		return nil, err
 	}
@@ -368,7 +368,7 @@ func (s *grpcServer) handle(ctx context.Context, diags []*tfprotov5.Diagnostic,
 		dd = append(dd, rd)
 		logDiag(ctx, rd)
 	}
-	derr := errors(dd)
+	derr := diagToError(dd)
 	if derr != nil && err != nil {
 		return multierror.Append(derr, err)
 	}

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/upgrade_state.go

@@ -6,82 +6,94 @@ import (
 	"strconv"
 
 	"github.com/hashicorp/go-cty/cty"
+	ctyjson "github.com/hashicorp/go-cty/cty/json"
+	"github.com/hashicorp/go-cty/cty/msgpack"
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
 
-func upgradeResourceState(ctx context.Context, p *schema.Provider, res *schema.Resource,
+func upgradeResourceState(ctx context.Context, typeName string, p *schema.Provider, res *schema.Resource,
 	instanceState *terraform.InstanceState) (*terraform.InstanceState, error) {
-
 	if instanceState == nil {
 		return nil, nil
 	}
 
-	m := instanceState.Attributes
-
 	// Ensure that we have an ID in the attributes.
-	m["id"] = instanceState.ID
+	if state := instanceState.RawState.AsValueMap(); !has(state, "id") {
+		state["id"] = cty.StringVal(instanceState.ID)
+		instanceState.RawState = cty.ObjectVal(state)
+	}
 
-	version, hasVersion := 0, false
+	version, hasVersion := int64(0), false
 	if versionValue, ok := instanceState.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
+			return nil, fmt.Errorf("failed to parse schema_version: %w", err)
 		}
-		version, hasVersion = int(v), true
+		version, hasVersion = v, true
 	}
 
-	// First, build a JSON state from the InstanceState.
-	json, version, err := schema.UpgradeFlatmapState(ctx, version, m, res, p.Meta())
+	json, err := ctyjson.Marshal(instanceState.RawState, res.CoreConfigSchema().ImpliedType())
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to marshal RawState: %w", err)
 	}
 
-	// Next, migrate the JSON state up to the current version.
-	json, err = schema.UpgradeJSONState(ctx, version, json, res, p.Meta())
+	// Now, we perform the UpgradeResourceState operation by calling into TF's UpgradeResourceState.
+
+	resp, err := schema.NewGRPCProviderServer(p).
+		UpgradeResourceState(ctx, &tfprotov5.UpgradeResourceStateRequest{
+			TypeName: typeName,
+			Version:  version,
+			RawState: &tfprotov5.RawState{JSON: json},
+		})
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("upgrade resource state GRPC: %w", err)
 	}
 
-	configBlock := res.CoreConfigSchema()
-
-	// Strip out removed fields.
-	schema.RemoveAttributes(ctx, json, configBlock.ImpliedType())
-
-	// now we need to turn the state into the default json representation, so
-	// that it can be re-decoded using the actual schema.
-	v, err := schema.JSONMapToStateValue(json, configBlock)
-	if err != nil {
-		return nil, err
+	// Handle returned diagnostics.
+	var dd diag.Diagnostics
+	for _, d := range resp.Diagnostics {
+		if d == nil {
+			continue
+		}
+		rd := recoverDiagnostic(*d)
+		dd = append(dd, rd)
+		logDiag(ctx, rd)
+	}
+	if err := diagToError(dd); err != nil {
+		return nil, fmt.Errorf("diag: %w", err)
 	}
 
-	// Now we need to make sure blocks are represented correctly, which means
-	// that missing blocks are empty collections, rather than null.
-	// First we need to CoerceValue to ensure that all object types match.
-	v, err = configBlock.CoerceValue(v)
+	// Unmarshal to get back the underlying type.
+	var rawState cty.Value
+	if resp.UpgradedState.JSON != nil {
+		rawState, err = ctyjson.Unmarshal(resp.UpgradedState.JSON, res.CoreConfigSchema().ImpliedType())
+	}
+	if resp.UpgradedState.MsgPack != nil {
+		rawState, err = msgpack.Unmarshal(resp.UpgradedState.MsgPack, res.CoreConfigSchema().ImpliedType())
+	}
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to unmarshal: %w", err)
 	}
 
-	// Normalize the value and fill in any missing blocks.
-	v = schema.NormalizeObjectFromLegacySDK(v, configBlock)
-
-	// Convert the value back to an InstanceState.
-	newState, err := res.ShimInstanceStateFromValue(v)
+	newState, err := res.ShimInstanceStateFromValue(rawState)
 	if err != nil {
 		return nil, err
 	}
+
 	newState.RawConfig = instanceState.RawConfig
 
 	// Copy the original ID and meta to the new state and stamp in the new version.
 	newState.ID = instanceState.ID
 
 	// If state upgraders have modified the ID, respect the modification.
-	if updatedID, ok := findID(v); ok {
+	if updatedID, ok := findID(rawState); ok {
 		newState.ID = updatedID
 	}
 
@@ -90,7 +102,7 @@ func upgradeResourceState(ctx context.Context, p *schema.Provider, res *schema.R
 		if newState.Meta == nil {
 			newState.Meta = map[string]interface{}{}
 		}
-		newState.Meta["schema_version"] = strconv.Itoa(version)
+		newState.Meta["schema_version"] = strconv.Itoa(int(version))
 	}
 	return newState, nil
 }
@@ -108,3 +120,8 @@ func findID(v cty.Value) (string, bool) {
 	}
 	return id.AsString(), true
 }
+
+func has[K comparable, V any](m map[K]V, k K) bool {
+	_, ok := m[k]
+	return ok
+}

pkg/tfshim/sdk-v2/upgrade_state_test.go

@@ -0,0 +1,53 @@
+package sdkv2
+
+import (
+	"context"
+	// "math/big"
+	"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 TestUpgradeResourceState(t *testing.T) {
+	state := 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
+	}
+
+	rSchema := &schema.Resource{
+		Schema: map[string]*schema.Schema{
+			"x": {Type: schema.TypeInt, Optional: true},
+		},
+	}
+
+	require.NoError(t, rSchema.InternalValidate(rSchema.Schema, true))
+
+	t.Logf(`Attributes["x"]: %s`, state().Attributes["x"])
+
+	const tfToken = "test_token"
+
+	actual, err := upgradeResourceState(context.Background(), tfToken, &schema.Provider{
+		ResourcesMap: map[string]*schema.Resource{
+			tfToken: {
+				UseJSONNumber: true,
+				Schema: map[string]*schema.Schema{
+					"x": {Type: schema.TypeInt, Optional: true},
+				},
+			},
+		},
+	}, rSchema, state())
+
+	require.NoError(t, err)
+	assert.Equal(t, state().Attributes, actual.Attributes)
+}