ensure status, private, warnings, etc are retained

Various pieces of the state and/or warnings were dropped when the
provider returns an error. Do a little cleanup of `.apply` to make the
logic easier to follow.

terraform/context_apply_test.go

@@ -12510,3 +12510,69 @@ func TestContext2Apply_errorRestorePrivateData(t *testing.T) {
 		t.Fatal("missing private data in state")
 	}
 }
+
+func TestContext2Apply_errorRestoreStatus(t *testing.T) {
+	// empty config to remove our resource
+	m := testModuleInline(t, map[string]string{
+		"main.tf": "",
+	})
+
+	p := simpleMockProvider()
+	p.ApplyResourceChangeFn = func(req providers.ApplyResourceChangeRequest) (resp providers.ApplyResourceChangeResponse) {
+		// We error during apply, but return the current object state.
+		resp.Diagnostics = resp.Diagnostics.Append(errors.New("oops"))
+		// return a warning too to make sure it isn't dropped
+		resp.Diagnostics = resp.Diagnostics.Append(tfdiags.SimpleWarning("warned"))
+		resp.NewState = req.PriorState
+		resp.Private = req.PlannedPrivate
+		return resp
+	}
+
+	addr := mustResourceInstanceAddr("test_object.a")
+
+	state := states.BuildState(func(s *states.SyncState) {
+		s.SetResourceInstanceCurrent(addr, &states.ResourceInstanceObjectSrc{
+			Status:    states.ObjectTainted,
+			AttrsJSON: []byte(`{"test_string":"foo"}`),
+			Private:   []byte("private"),
+		}, mustProviderConfig(`provider["registry.terraform.io/hashicorp/test"]`))
+	})
+
+	ctx := testContext2(t, &ContextOpts{
+		Config: m,
+		State:  state,
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("test"): testProviderFuncFixed(p),
+		},
+	})
+
+	_, diags := ctx.Plan()
+	if diags.HasErrors() {
+		t.Fatal(diags.Err())
+	}
+
+	state, diags = ctx.Apply()
+
+	if len(diags) != 2 {
+		t.Fatal("expected 1 error and 1 warning")
+	}
+
+	errString := diags.ErrWithWarnings().Error()
+	if !strings.Contains(errString, "oops") || !strings.Contains(errString, "warned") {
+		t.Fatalf("error missing expected info: %q", errString)
+	}
+
+	res := state.ResourceInstance(addr)
+	if res == nil {
+		t.Fatal("resource was removed from state")
+	}
+
+	if res.Current.Status != states.ObjectTainted {
+		t.Fatal("resource should still be tainted in the state")
+	}
+
+	if string(res.Current.Private) != "private" {
+		t.Fatalf("incorrect private data, got %q", res.Current.Private)
+	}
+
+}

terraform/node_resource_abstract_instance.go

@@ -1870,7 +1870,6 @@ func (n *NodeAbstractResourceInstance) apply(
 	unmarkedBefore, beforePaths := change.Before.UnmarkDeepWithPaths()
 	unmarkedAfter, afterPaths := change.After.UnmarkDeepWithPaths()
 
-	var newState *states.ResourceInstanceObject
 	// If we have an Update action, our before and after values are equal,
 	// and only differ on their sensitivity, the newVal is the after val
 	// and we should not communicate with the provider. We do need to update
@@ -1880,7 +1879,7 @@ func (n *NodeAbstractResourceInstance) apply(
 	eq := eqV.IsKnown() && eqV.True()
 	if change.Action == plans.Update && eq && !marksEqual(beforePaths, afterPaths) {
 		// Copy the previous state, changing only the value
-		newState = &states.ResourceInstanceObject{
+		newState := &states.ResourceInstanceObject{
 			CreateBeforeDestroy: state.CreateBeforeDestroy,
 			Dependencies:        state.Dependencies,
 			Private:             state.Private,
@@ -1956,7 +1955,7 @@ func (n *NodeAbstractResourceInstance) apply(
 		// Bail early in this particular case, because an object that doesn't
 		// conform to the schema can't be saved in the state anyway -- the
 		// serializer will reject it.
-		return newState, diags, applyError
+		return nil, diags, applyError
 	}
 
 	// After this point we have a type-conforming result object and so we
@@ -2072,35 +2071,40 @@ func (n *NodeAbstractResourceInstance) apply(
 		}
 	}
 
-	// Sometimes providers return a null value when an operation fails for some
+	switch {
-	// reason, but we'd rather keep the prior state so that the error can be
+	case diags.HasErrors() && newVal.IsNull():
-	// corrected on a subsequent run. We must only do this for null new value
+		// Sometimes providers return a null value when an operation fails for
-	// though, or else we may discard partial updates the provider was able to
+		// some reason, but we'd rather keep the prior state so that the error
-	// complete.
+		// can be corrected on a subsequent run. We must only do this for null
-	if diags.HasErrors() && newVal.IsNull() {
+		// new value though, or else we may discard partial updates the
-		// Otherwise, we'll continue but using the prior state as the new value,
+		// provider was able to complete. Otherwise, we'll continue using the
-		// making this effectively a no-op. If the item really _has_ been
+		// prior state as the new value, making this effectively a no-op.  If
-		// deleted then our next refresh will detect that and fix it up.
+		// the item really _has_ been deleted then our next refresh will detect
-		// If change.Action is Create then change.Before will also be null,
+		// that and fix it up.
-		// which is fine.
+		return state.DeepCopy(), nil, diags.Err()
-		newState = state.DeepCopy()
-	}
 
-	if !newVal.IsNull() { // null value indicates that the object is deleted, so we won't set a new state in that case
+	case diags.HasErrors() && !newVal.IsNull():
-		newState = &states.ResourceInstanceObject{
+		// if we have an error, make sure we restore the object status in the new state
+		newState := &states.ResourceInstanceObject{
+			Status:              state.Status,
+			Value:               newVal,
+			Private:             resp.Private,
+			CreateBeforeDestroy: createBeforeDestroy,
+		}
+		return newState, nil, diags.Err()
+
+	case !newVal.IsNull():
+		// Non error case with a new state
+		newState := &states.ResourceInstanceObject{
 			Status:              states.ObjectReady,
 			Value:               newVal,
 			Private:             resp.Private,
 			CreateBeforeDestroy: createBeforeDestroy,
 		}
-	}
+		return newState, diags, nil
 
-	if diags.HasErrors() {
+	default:
-		// At this point, if we have an error in diags (and hadn't already returned), we return it as an error and clear the diags.
+		// Non error case, were the object was deleted
-		applyError = diags.Err()
+		return nil, diags, nil
-		log.Printf("[DEBUG] %s: apply errored", n.Addr)
-		return newState, nil, applyError
 	}
-
-	return newState, diags, applyError
 }