From 1ca631bda0f3986f8ff4a563cfd6c5bca51baf3d Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Thu, 16 Mar 2023 20:46:40 -0400
Subject: [PATCH] remove planned objects from state on error

When planning encounters an error we were returning early without
cleaning out any planed data sources which cannot be serialized. Move
the cleanup to the common walkPlan method where the PriorState is
assigned so that it cannot be missed.
---
 internal/terraform/context_plan.go            | 22 +++---------
 internal/terraform/context_plan2_test.go      | 34 +++++++++++++++++++
 .../data-source-read-with-plan-error/main.tf  | 12 +++++++
 3 files changed, 51 insertions(+), 17 deletions(-)
 create mode 100644 internal/terraform/testdata/data-source-read-with-plan-error/main.tf

diff --git a/internal/terraform/context_plan.go b/internal/terraform/context_plan.go
index fc2e0be616..0e97f7cb21 100644
--- a/internal/terraform/context_plan.go
+++ b/internal/terraform/context_plan.go
@@ -286,18 +286,6 @@ func (c *Context) plan(config *configs.Config, prevRunState *states.State, opts
 
 	plan, walkDiags := c.planWalk(config, prevRunState, opts)
 	diags = diags.Append(walkDiags)
-	if diags.HasErrors() {
-		// Non-nil plan along with errors indicates a non-applyable partial
-		// plan that's only suitable to be shown to the user as extra context
-		// to help understand the errors.
-		return plan, diags
-	}
-
-	// The refreshed state ends up with some placeholder objects in it for
-	// objects pending creation. We only really care about those being in
-	// the working state, since that's what we're going to use when applying,
-	// so we'll prune them all here.
-	plan.PriorState.SyncWrapper().RemovePlannedResourceInstanceObjects()
 
 	return plan, diags
 }
@@ -339,10 +327,6 @@ func (c *Context) refreshOnlyPlan(config *configs.Config, prevRunState *states.S
 		))
 	}
 
-	// Prune out any placeholder objects we put in the state to represent
-	// objects that would need to be created.
-	plan.PriorState.SyncWrapper().RemovePlannedResourceInstanceObjects()
-
 	// We don't populate RelevantResources for a refresh-only plan, because
 	// they never have any planned actions and so no resource can ever be
 	// "relevant" per the intended meaning of that field.
@@ -580,9 +564,13 @@ func (c *Context) planWalk(config *configs.Config, prevRunState *states.State, o
 	// we encountered errors, which we'll return as part of a non-nil plan
 	// so that e.g. the UI can show what was planned so far in case that extra
 	// context helps the user to understand the error messages we're returning.
-
 	prevRunState = walker.PrevRunState.Close()
+
+	// The refreshed state may have data resource objects which were deferred
+	// to apply and cannot be serialized.
+	walker.RefreshState.RemovePlannedResourceInstanceObjects()
 	priorState := walker.RefreshState.Close()
+
 	driftedResources, driftDiags := c.driftedResources(config, prevRunState, priorState, moveResults)
 	diags = diags.Append(driftDiags)
 
diff --git a/internal/terraform/context_plan2_test.go b/internal/terraform/context_plan2_test.go
index 634144dee0..62311e7314 100644
--- a/internal/terraform/context_plan2_test.go
+++ b/internal/terraform/context_plan2_test.go
@@ -3947,3 +3947,37 @@ output "out" {
 	})
 	assertNoErrors(t, diags)
 }
+
+// Make sure the data sources in the prior state are serializeable even if
+// there were an error in the plan.
+func TestContext2Plan_dataSourceReadPlanError(t *testing.T) {
+	m, snap := testModuleWithSnapshot(t, "data-source-read-with-plan-error")
+	awsProvider := testProvider("aws")
+	testProvider := testProvider("test")
+
+	testProvider.PlanResourceChangeFn = func(req providers.PlanResourceChangeRequest) (resp providers.PlanResourceChangeResponse) {
+		resp.PlannedState = req.ProposedNewState
+		resp.Diagnostics = resp.Diagnostics.Append(errors.New("oops"))
+		return resp
+	}
+
+	state := states.NewState()
+
+	ctx := testContext2(t, &ContextOpts{
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("aws"):  testProviderFuncFixed(awsProvider),
+			addrs.NewDefaultProvider("test"): testProviderFuncFixed(testProvider),
+		},
+	})
+
+	plan, diags := ctx.Plan(m, state, DefaultPlanOpts)
+	if !diags.HasErrors() {
+		t.Fatalf("expected plan error")
+	}
+
+	// make sure we can serialize the plan even if there were an error
+	_, _, _, err := contextOptsForPlanViaFile(t, snap, plan)
+	if err != nil {
+		t.Fatalf("failed to round-trip through planfile: %s", err)
+	}
+}
diff --git a/internal/terraform/testdata/data-source-read-with-plan-error/main.tf b/internal/terraform/testdata/data-source-read-with-plan-error/main.tf
new file mode 100644
index 0000000000..2559406f7a
--- /dev/null
+++ b/internal/terraform/testdata/data-source-read-with-plan-error/main.tf
@@ -0,0 +1,12 @@
+resource "aws_instance" "foo" {
+}
+
+// this will be postponed until apply
+data "aws_data_source" "foo" {
+  foo = aws_instance.foo.id
+}
+
+// this will cause an error in the final plan
+resource "test_instance" "bar" {
+  foo = "error"
+}