Fix framework identity interceptor

.changelog/44518.txt

@@ -0,0 +1,6 @@
+```release-note:bug
+provider: Fix `Missing Resource Identity After Update` errors for non-refreshed and failed updates of Plugin Framework based resources
+```
+```release-note:bug
+provider: Fix `Unexpected Identity Change` errors when fully-null identity values in state are updated to valid values for Plugin Framework based resources
+```

internal/provider/framework/identity_interceptor.go

@@ -9,6 +9,8 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/attr"
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
+	"github.com/hashicorp/terraform-plugin-framework/tfsdk"
+	"github.com/hashicorp/terraform-plugin-framework/types"
 	inttypes "github.com/hashicorp/terraform-provider-aws/internal/types"
 	"github.com/hashicorp/terraform-provider-aws/names"
 )
@@ -56,6 +58,41 @@ func (r identityInterceptor) create(ctx context.Context, opts interceptorOptions
 				}
 			}
 		}
+	case OnError:
+		identity := response.Identity
+		if identity == nil {
+			break
+		}
+
+		if identityIsFullyNull(ctx, identity, r.attributes) {
+			for _, att := range r.attributes {
+				switch att.Name() {
+				case names.AttrAccountID:
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.AccountID(ctx))...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+				case names.AttrRegion:
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.Region(ctx))...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+				default:
+					var attrVal attr.Value
+					opts.response.Diagnostics.Append(response.State.GetAttribute(ctx, path.Root(att.ResourceAttributeName()), &attrVal)...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), attrVal)...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+				}
+			}
+		}
 	}
 }
 
@@ -103,11 +140,109 @@ func (r identityInterceptor) read(ctx context.Context, opts interceptorOptions[r
 }
 
 func (r identityInterceptor) update(ctx context.Context, opts interceptorOptions[resource.UpdateRequest, resource.UpdateResponse]) {
+	awsClient := opts.c
+
+	switch response, when := opts.response, opts.when; when {
+	case After:
+		if response.State.Raw.IsNull() {
+			break
+		}
+		identity := response.Identity
+		if identity == nil {
+			break
+		}
+
+		for _, att := range r.attributes {
+			switch att.Name() {
+			case names.AttrAccountID:
+				opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.AccountID(ctx))...)
+				if opts.response.Diagnostics.HasError() {
+					return
+				}
+
+			case names.AttrRegion:
+				opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.Region(ctx))...)
+				if opts.response.Diagnostics.HasError() {
+					return
+				}
+
+			default:
+				var attrVal attr.Value
+				opts.response.Diagnostics.Append(response.State.GetAttribute(ctx, path.Root(att.ResourceAttributeName()), &attrVal)...)
+				if opts.response.Diagnostics.HasError() {
+					return
+				}
+
+				opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), attrVal)...)
+				if opts.response.Diagnostics.HasError() {
+					return
+				}
+			}
+		}
+	case OnError:
+		if response.State.Raw.IsNull() {
+			break
+		}
+		identity := response.Identity
+		if identity == nil {
+			break
+		}
+
+		if identityIsFullyNull(ctx, identity, r.attributes) {
+			for _, att := range r.attributes {
+				switch att.Name() {
+				case names.AttrAccountID:
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.AccountID(ctx))...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+				case names.AttrRegion:
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), awsClient.Region(ctx))...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+				default:
+					var attrVal attr.Value
+					opts.response.Diagnostics.Append(response.State.GetAttribute(ctx, path.Root(att.ResourceAttributeName()), &attrVal)...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+
+					opts.response.Diagnostics.Append(identity.SetAttribute(ctx, path.Root(att.Name()), attrVal)...)
+					if opts.response.Diagnostics.HasError() {
+						return
+					}
+				}
+			}
+		}
+	}
 }
 
 func (r identityInterceptor) delete(ctx context.Context, opts interceptorOptions[resource.DeleteRequest, resource.DeleteResponse]) {
 }
 
+// identityIsFullyNull returns true if a resource supports identity and
+// all attributes are set to null values
+func identityIsFullyNull(ctx context.Context, identity *tfsdk.ResourceIdentity, attributes []inttypes.IdentityAttribute) bool {
+	if identity == nil {
+		return true
+	}
+
+	for _, attr := range attributes {
+		var attrVal types.String
+		if diags := identity.GetAttribute(ctx, path.Root(attr.Name()), &attrVal); diags.HasError() {
+			return false
+		}
+		if !attrVal.IsNull() && attrVal.ValueString() != "" {
+			return false
+		}
+	}
+
+	return true
+}
+
 func newIdentityInterceptor(attributes []inttypes.IdentityAttribute) identityInterceptor {
 	return identityInterceptor{
 		attributes: attributes,

internal/provider/framework/identity_interceptor_test.go

@@ -269,3 +269,96 @@ func (c mockClient) ValidateInContextRegionInPartition(ctx context.Context) erro
 func (c mockClient) AwsConfig(context.Context) aws.Config { // nosemgrep:ci.aws-in-func-name
 	panic("not implemented") //lintignore:R009
 }
+
+func TestIdentityIsFullyNull(t *testing.T) {
+	t.Parallel()
+
+	attributes := []inttypes.IdentityAttribute{
+		inttypes.StringIdentityAttribute("account_id", false),
+		inttypes.StringIdentityAttribute("region", false),
+		inttypes.StringIdentityAttribute("bucket", true),
+	}
+
+	testCases := map[string]struct {
+		identityValues map[string]string
+		expectNull     bool
+		description    string
+	}{
+		"all_null": {
+			identityValues: map[string]string{},
+			expectNull:     true,
+			description:    "All attributes null should return true",
+		},
+		"some_null": {
+			identityValues: map[string]string{
+				"account_id": "123456789012",
+				// region and bucket remain null
+			},
+			expectNull:  false,
+			description: "Some attributes set should return false",
+		},
+		"all_set": {
+			identityValues: map[string]string{
+				"account_id": "123456789012",
+				"region":     "us-west-2", // lintignore:AWSAT003
+				"bucket":     "test-bucket",
+			},
+			expectNull:  false,
+			description: "All attributes set should return false",
+		},
+		"empty_string_values": {
+			identityValues: map[string]string{
+				"account_id": "",
+				"region":     "",
+				"bucket":     "",
+			},
+			expectNull:  true,
+			description: "Empty string values should be treated as null",
+		},
+		"nil_identity": {
+			identityValues: nil, // This will result in nil identity
+			expectNull:     true,
+			description:    "Nil identity should return true",
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			ctx := context.Background()
+
+			// Create identity schema
+			identitySchema := &identityschema.Schema{
+				Attributes: map[string]identityschema.Attribute{
+					"account_id": identityschema.StringAttribute{},
+					"region":     identityschema.StringAttribute{},
+					"bucket":     identityschema.StringAttribute{},
+				},
+			}
+
+			var identity *tfsdk.ResourceIdentity
+
+			// Handle nil identity case
+			if tc.identityValues == nil {
+				identity = nil
+			} else {
+				// Create identity with values
+				identity = emtpyIdentityFromSchema(ctx, identitySchema)
+				for attrName, value := range tc.identityValues {
+					if value != "" {
+						diags := identity.SetAttribute(ctx, path.Root(attrName), value)
+						if diags.HasError() {
+							t.Fatalf("unexpected error setting %s in identity: %s", attrName, fwdiag.DiagnosticsError(diags))
+						}
+					}
+				}
+			}
+
+			result := identityIsFullyNull(ctx, identity, attributes)
+			if result != tc.expectNull {
+				t.Errorf("%s: expected identityIsFullyNull to return %v, got %v",
+					tc.description, tc.expectNull, result)
+			}
+		})
+	}
+}