From 0ef10a3f9d602696291a325d78c4bb9270d5b1c6 Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Tue, 10 Dec 2024 16:33:19 +0000
Subject: [PATCH] Revert "Revert "Pass state back to the engine if Apply
 encountered an error" (#2707)"

---
 pkg/tests/schema_pulumi_test.go | 51 +++++++++++++++++++++++++++++++++
 pkg/tfshim/sdk-v2/provider2.go  | 37 +++++++++++-------------
 2 files changed, 67 insertions(+), 21 deletions(-)

diff --git a/pkg/tests/schema_pulumi_test.go b/pkg/tests/schema_pulumi_test.go
index ca2b8d485..791f331e6 100644
--- a/pkg/tests/schema_pulumi_test.go
+++ b/pkg/tests/schema_pulumi_test.go
@@ -189,3 +189,54 @@ func TestMakeTerraformResultNilVsEmptyMap(t *testing.T) {
 		assert.True(t, props["test"].DeepEquals(emptyMap))
 	})
 }
+
+func TestResourceInitFailure(t *testing.T) {
+	t.Parallel()
+
+	resMap := map[string]*schema.Resource{
+		"prov_test": {
+			Schema: map[string]*schema.Schema{
+				"test": {
+					Type:     schema.TypeString,
+					Required: true,
+				},
+			},
+			CreateContext: func(ctx context.Context, rd *schema.ResourceData, i interface{}) diag.Diagnostics {
+				rd.SetId("1")
+				return diag.Errorf("INIT TEST ERROR")
+			},
+		},
+	}
+	prov := &schema.Provider{ResourcesMap: resMap}
+	bridgedProvider := pulcheck.BridgedProvider(t, "prov", prov)
+
+	pt := pulcheck.PulCheck(t, bridgedProvider, `
+name: test
+runtime: yaml
+resources:
+  mainRes:
+    type: prov:index:Test
+	properties:
+	  test: "hello"
+`)
+
+	_, err := pt.CurrentStack().Up(pt.Context())
+	require.Error(t, err)
+	require.ErrorContains(t, err, "INIT TEST ERROR")
+
+	stack := pt.ExportStack(t)
+
+	data, err := stack.Deployment.MarshalJSON()
+	require.NoError(t, err)
+
+	var stateMap map[string]interface{}
+	err = json.Unmarshal(data, &stateMap)
+	require.NoError(t, err)
+
+	resourcesList := stateMap["resources"].([]interface{})
+	require.Len(t, resourcesList, 3)
+	mainResState := resourcesList[2].(map[string]interface{}) // stack, provider, resource
+	initErrors := mainResState["initErrors"].([]interface{})
+	require.Len(t, initErrors, 1)
+	require.Contains(t, initErrors[0], "INIT TEST ERROR")
+}
diff --git a/pkg/tfshim/sdk-v2/provider2.go b/pkg/tfshim/sdk-v2/provider2.go
index ef6e33438..347496e15 100644
--- a/pkg/tfshim/sdk-v2/provider2.go
+++ b/pkg/tfshim/sdk-v2/provider2.go
@@ -347,15 +347,7 @@ func (p *planResourceChangeImpl) Apply(
 		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
-	}
-	return &v2InstanceState2{
-		resourceType: t,
-		stateValue:   resp.stateValue,
-		meta:         resp.meta,
-	}, nil
+	return p.server.ApplyResourceChange(ctx, t, ty, cfg, st, pl, priv, meta)
 }
 
 // This method is called to service `pulumi refresh` requests and maps naturally to the TF
@@ -662,25 +654,28 @@ func (s *grpcServer) ApplyResourceChange(
 		}
 		req.ProviderMeta = &tfprotov5.DynamicValue{MsgPack: providerMetaVal}
 	}
-	resp, err := s.gserver.ApplyResourceChange(ctx, req)
-	if err := handleDiagnostics(ctx, resp.Diagnostics, err); err != nil {
-		return nil, err
-	}
-	newState, err := msgpack.Unmarshal(resp.NewState.MsgPack, ty)
-	if err != nil {
-		return nil, err
-	}
+	resp, applyErr := s.gserver.ApplyResourceChange(ctx, req)
+	newState := cty.Value{}
 	var meta map[string]interface{}
-	if resp.Private != nil {
-		if err := json.Unmarshal(resp.Private, &meta); err != nil {
-			return nil, err
+	if resp != nil {
+		if resp.NewState != nil {
+			newState, err = msgpack.Unmarshal(resp.NewState.MsgPack, ty)
+			if err != nil {
+				return nil, err
+			}
+		}
+		if resp.Private != nil {
+			if err := json.Unmarshal(resp.Private, &meta); err != nil {
+				return nil, err
+			}
 		}
 	}
+	returnErr := handleDiagnostics(ctx, resp.Diagnostics, applyErr)
 	return &v2InstanceState2{
 		resourceType: typeName,
 		stateValue:   newState,
 		meta:         meta,
-	}, nil
+	}, returnErr
 }
 
 func (s *grpcServer) ReadResource(