Handle diags from gRPC state upgrader (#2053)

Isolating this from an AWS test failure; handling diags turns a panic
into a regular error. Looks like the error message is still unfortunate.
The test is isolated from
pulumi/pulumi-aws#4015

Author: t0yv0

pkg/tfshim/sdk-v2/provider2.go

@@ -353,7 +353,7 @@ type grpcServer struct {
 // This will return an error if any of the diagnostics are error-level, or a given err is non-nil.
 // It will also logs the diagnostics into TF loggers, so they appear when debugging with the bridged
 // provider with TF_LOG=TRACE or similar.
-func (s *grpcServer) handle(ctx context.Context, diags []*tfprotov5.Diagnostic, err error) error {
+func handleDiagnostics(ctx context.Context, diags []*tfprotov5.Diagnostic, err error) error {
 	var dd diag.Diagnostics
 	for _, d := range diags {
 		if d == nil {
@@ -430,7 +430,7 @@ func (s *grpcServer) PlanResourceChange(
 		req.ProviderMeta = &tfprotov5.DynamicValue{MsgPack: providerMetaVal}
 	}
 	resp, err := s.gserver.PlanResourceChangeExtra(ctx, req)
-	if err := s.handle(ctx, resp.Diagnostics, err); err != nil {
+	if err := handleDiagnostics(ctx, resp.Diagnostics, err); err != nil {
 		return nil, err
 	}
 	// Ignore resp.UnsafeToUseLegacyTypeSystem - does not matter for Pulumi bridged providers.
@@ -508,7 +508,7 @@ func (s *grpcServer) ApplyResourceChange(
 		req.ProviderMeta = &tfprotov5.DynamicValue{MsgPack: providerMetaVal}
 	}
 	resp, err := s.gserver.ApplyResourceChange(ctx, req)
-	if err := s.handle(ctx, resp.Diagnostics, err); err != nil {
+	if err := handleDiagnostics(ctx, resp.Diagnostics, err); err != nil {
 		return nil, err
 	}
 	newState, err := msgpack.Unmarshal(resp.NewState.MsgPack, ty)
@@ -559,7 +559,7 @@ func (s *grpcServer) ReadResource(
 		req.ProviderMeta = &tfprotov5.DynamicValue{MsgPack: providerMetaVal}
 	}
 	resp, err := s.gserver.ReadResource(ctx, req)
-	if err := s.handle(ctx, resp.Diagnostics, err); err != nil {
+	if err := handleDiagnostics(ctx, resp.Diagnostics, err); err != nil {
 		return nil, err
 	}
 	newState, err := msgpack.Unmarshal(resp.NewState.MsgPack, ty)
@@ -590,7 +590,7 @@ func (s *grpcServer) ImportResourceState(
 		ID:       id,
 	}
 	resp, err := s.gserver.ImportResourceState(ctx, req)
-	if err := s.handle(ctx, resp.Diagnostics, err); err != nil {
+	if err := handleDiagnostics(ctx, resp.Diagnostics, err); err != nil {
 		return nil, err
 	}
 	out := []v2InstanceState2{}

pkg/tfshim/sdk-v2/provider2_test.go

@@ -14,10 +14,11 @@ import (
 func TestProvider2UpgradeResourceState(t *testing.T) {
 	const tfToken = "test_token"
 	for _, tc := range []struct {
-		name     string
-		state    cty.Value
-		rSchema  *schema.Resource
-		expected cty.Value
+		name      string
+		state     cty.Value
+		rSchema   *schema.Resource
+		expected  cty.Value
+		expectErr func(*testing.T, error)
 	}{
 		{
 			name: "no upgrade",
@@ -213,6 +214,58 @@ func TestProvider2UpgradeResourceState(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "handle errors",
+			state: func() cty.Value {
+				return cty.ObjectVal(map[string]cty.Value{
+					"compute_resources": cty.ObjectVal(map[string]cty.Value{
+						"ec2_configuration": cty.ObjectVal(map[string]cty.Value{
+							"image_id_override": cty.StringVal("override"),
+						}),
+					}),
+					"id": cty.StringVal("id"),
+				})
+			}(),
+			expectErr: func(t *testing.T, err error) {
+				require.Error(t, err)
+				require.ErrorContains(t, err, "missing expected [")
+			},
+			rSchema: &schema.Resource{
+				Schema: map[string]*schema.Schema{
+					"compute_resources": {
+						Type:     schema.TypeList,
+						Optional: true,
+						ForceNew: true,
+						MinItems: 0,
+						MaxItems: 1,
+						Elem: &schema.Resource{
+							Schema: map[string]*schema.Schema{
+								"ec2_configuration": {
+									Type:     schema.TypeList,
+									Optional: true,
+									Computed: true,
+									ForceNew: true,
+									MaxItems: 2,
+									Elem: &schema.Resource{
+										Schema: map[string]*schema.Schema{
+											"image_id_override": {
+												Type:     schema.TypeString,
+												Optional: true,
+												Computed: true,
+											},
+											"image_type": {
+												Type:     schema.TypeString,
+												Optional: true,
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			tf := &schema.Provider{
@@ -227,9 +280,12 @@ func TestProvider2UpgradeResourceState(t *testing.T) {
 				resourceType: tfToken,
 				stateValue:   tc.state,
 			})
-			require.NoError(t, err)
-
-			assert.Equal(t, tc.expected, actual.(*v2InstanceState2).stateValue)
+			if tc.expectErr != nil {
+				tc.expectErr(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.expected, actual.(*v2InstanceState2).stateValue)
+			}
 		})
 	}
 }

pkg/tfshim/sdk-v2/upgrade_state.go

@@ -49,8 +49,8 @@ func upgradeResourceStateGRPC(
 		RawState: &tfprotov5.RawState{JSON: jsonBytes},
 		Version:  version,
 	})
-	if err != nil {
-		return cty.Value{}, nil, err
+	if uerr := handleDiagnostics(ctx, resp.Diagnostics, err); uerr != nil {
+		return cty.Value{}, nil, uerr
 	}
 
 	newState, err := msgpack.Unmarshal(resp.UpgradedState.MsgPack, res.CoreConfigSchema().ImpliedType())