Switch upgradeResourceState to a grcp method

This ensures that our behavior aligns with TF's behavior.

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

@@ -94,7 +94,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 {
@@ -115,12 +115,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(
@@ -139,7 +139,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)
 	}
@@ -149,7 +149,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(
@@ -175,7 +175,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

@@ -343,7 +343,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
 	}
@@ -380,7 +380,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

@@ -2,17 +2,19 @@ package sdkv2
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 
 	"github.com/hashicorp/go-cty/cty"
+	"github.com/hashicorp/go-cty/cty/msgpack"
+	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
 	"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
 	}
@@ -22,7 +24,7 @@ func upgradeResourceState(ctx context.Context, p *schema.Provider, res *schema.R
 	// Ensure that we have an ID in the attributes.
 	m["id"] = instanceState.ID
 
-	version, hasVersion := 0, false
+	version, hasVersion := int64(0), false
 	if versionValue, ok := instanceState.Meta["schema_version"]; ok {
 		versionString, ok := versionValue.(string)
 		if !ok {
@@ -32,56 +34,26 @@ func upgradeResourceState(ctx context.Context, p *schema.Provider, res *schema.R
 		if err != nil {
 			return nil, 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())
+	rawState, err := upgradeResourceStateRPC(ctx, typeName, m, p, res, version)
 	if err != nil {
 		return nil, err
 	}
 
-	// Next, migrate the JSON state up to the current version.
-	json, err = schema.UpgradeJSONState(ctx, version, json, res, p.Meta())
-	if err != nil {
-		return nil, 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)
+	newState, err := res.ShimInstanceStateFromValue(rawState)
 	if err != nil {
 		return nil, 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)
-	if err != nil {
-		return nil, 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)
-	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 +62,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 +80,39 @@ func findID(v cty.Value) (string, bool) {
 	}
 	return id.AsString(), true
 }
+
+// Perform the UpgradeResourceState operation by invoking TF's underlying gRPC server.
+func upgradeResourceStateRPC(
+	ctx context.Context, typeName string, m map[string]string,
+	p *schema.Provider, res *schema.Resource,
+	version int64,
+) (cty.Value, error) {
+	// Call SDKv2's underlying UpgradeResourceState.
+	resp, err := schema.NewGRPCProviderServer(p).
+		UpgradeResourceState(ctx, &tfprotov5.UpgradeResourceStateRequest{
+			TypeName: typeName,
+			Version:  version,
+			RawState: &tfprotov5.RawState{Flatmap: m},
+		})
+	if err != nil {
+		return cty.Value{}, fmt.Errorf("upgrade resource state: %w", err)
+	}
+
+	// Handle returned diagnostics.
+	for _, d := range resp.Diagnostics {
+		msg := fmt.Sprintf("%s: %s", d.Summary, d.Detail)
+		switch d.Severity {
+		case tfprotov5.DiagnosticSeverityError:
+			err = errors.Join(err, d.Attribute.NewError(fmt.Errorf("%s", msg)))
+		case tfprotov5.DiagnosticSeverityWarning:
+			// Accessing the logger (GetLogger) requires an import cycle on
+			// tfbridge, so ignore for now.
+		}
+	}
+	if err != nil {
+		return cty.Value{}, err
+	}
+
+	// Unmarshal to get back the underlying type.
+	return msgpack.Unmarshal(resp.UpgradedState.MsgPack, res.CoreConfigSchema().ImpliedType())
+}