From 5978502fd93fba7d49604cb7c3266299b79ce8f5 Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Fri, 6 Dec 2024 16:10:27 +0000
Subject: [PATCH] fix partial init errors

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

diff --git a/pkg/tests/schema_pulumi_test.go b/pkg/tests/schema_pulumi_test.go
index 903d48476..37113570a 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..b8f2cc0d2 100644
--- a/pkg/tfshim/sdk-v2/provider2.go
+++ b/pkg/tfshim/sdk-v2/provider2.go
@@ -349,7 +349,7 @@ func (p *planResourceChangeImpl) Apply(
 
 	resp, err := p.server.ApplyResourceChange(ctx, t, ty, cfg, st, pl, priv, meta)
 	if err != nil {
-		return nil, err
+		return resp, err
 	}
 	return &v2InstanceState2{
 		resourceType: t,
@@ -662,25 +662,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 newState, err
+			}
 		}
 	}
+	returnErr := handleDiagnostics(ctx, resp.Diagnostics, applyErr)
 	return &v2InstanceState2{
 		resourceType: typeName,
 		stateValue:   newState,
 		meta:         meta,
-	}, nil
+	}, returnErr
 }
 
 func (s *grpcServer) ReadResource(