Fix framework identity interceptor

Author: YakDriver

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,