Addressed PR commenta

rainkwan · rainkwan · commit c8ff5badaef8 · 2025-05-05T13:26:48.000-04:00
diff --git a/internal/fwserver/server_upgraderesourceidentity.go b/internal/fwserver/server_upgraderesourceidentity.go
@@ -69,51 +69,16 @@ func (s *Server) UpgradeResourceIdentity(ctx context.Context, req *UpgradeResour
 		},
 	}
 
-	// TODO: Maybe throw error if the schemas are the same, it is a bug in core
-
-	// Terraform CLI can call UpgradeResourceIdentity even if the stored Identity
-	// version matches the current schema. Presumably this is to account for
-	// the previous terraform-plugin-sdk implementation, which handled some
-	// Identity fixups on behalf of Terraform CLI. When this happens, we do not
-	// want to return errors for a missing ResourceWithUpgradeIdentity
-	// implementation or an undefined version within an existing
-	// ResourceWithUpgradeIdentity implementation as that would be confusing
-	// detail for provider developers. Instead, the framework will attempt to
-	// roundtrip the prior RawState to a Identity matching the current Schema.
-	//
-	// TODO: To prevent provider developers from accidentally implementing
-	// ResourceWithUpgradeIdentity with a version matching the current schema
-	// version which would never get called, the framework can introduce a
-	// unit test helper.
-	// Reference: https://github.com/hashicorp/terraform-plugin-framework/issues/113
-	//
-	// UnmarshalWithOpts allows optionally ignoring instances in which elements being
-	// do not have a corresponding attribute within the schema.
-	/*	if req.Version == req.IdentitySchema.GetVersion() {
-		logging.FrameworkTrace(ctx, "UpgradeResourceIdentity request version matches current Schema version, using framework defined passthrough implementation")
-
-		identitySchemaType := req.IdentitySchema.Type().TerraformType(ctx)
-
-		rawStateValue, err := req.RawState.UnmarshalWithOpts(identitySchemaType, unmarshalOpts)
-
-		if err != nil {
-			resp.Diagnostics.AddError(
-				"Unable to Read Previously Saved Identity for UpgradeResourceIdentity",
-				"There was an error reading the saved resource Identity using the current resource schema.\n\n"+
-					"If this resource Identity was last refreshed with Terraform CLI 0.11 and earlier, it must be refreshed or applied with an older provider version first. "+
-					"If you manually modified the resource Identity, you will need to manually modify it to match the current resource schema. "+
-					"Otherwise, please report this to the provider developer:\n\n"+err.Error(),
-			)
-			return
-		}
-
-		resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
-			Schema: req.IdentitySchema,
-			Raw:    rawStateValue,
-		}
-
+	if req.Version == req.IdentitySchema.GetVersion() {
+		resp.Diagnostics.AddError(
+			"Unable to Read Previously Saved Identity for UpgradeResourceIdentity",
+			"There was an error reading the saved resource identity using the current resource identity schema.\n\n"+
+				"If this resource Identity was last refreshed with Terraform CLI 0.11 and earlier, it must be refreshed or applied with an older provider version first. "+
+				"If you manually modified the resource Identity, you will need to manually modify it to match the current resource identity schema. "+
+				"Otherwise, please report this to the provider developer:\n\n",
+		)
 		return
-	}*/
+	}
 
 	if resourceWithConfigure, ok := req.Resource.(resource.ResourceWithConfigure); ok {
 		logging.FrameworkTrace(ctx, "Resource implements ResourceWithConfigure")
@@ -178,7 +143,7 @@ func (s *Server) UpgradeResourceIdentity(ctx context.Context, req *UpgradeResour
 
 		priorSchemaType := resourceIdentityUpgrader.PriorSchema.Type().TerraformType(ctx)
 
-		_, err := req.RawState.UnmarshalWithOpts(priorSchemaType, unmarshalOpts)
+		rawIdentityValue, err := req.RawState.UnmarshalWithOpts(priorSchemaType, unmarshalOpts)
 
 		if err != nil {
 			resp.Diagnostics.AddError(
@@ -189,12 +154,15 @@ func (s *Server) UpgradeResourceIdentity(ctx context.Context, req *UpgradeResour
 			return
 		}
 
-		upgradeResourceIdentityRequest.RawState = req.RawState
+		upgradeResourceIdentityRequest.Identity = &tfsdk.ResourceIdentity{
+			Raw:    rawIdentityValue, // from the output of req.RawState.UnmarshalWithOpts
+			Schema: *resourceIdentityUpgrader.PriorSchema,
+		}
 
 	}
 
 	upgradeResourceIdentityResponse := resource.UpgradeResourceIdentityResponse{
-		UpgradedIdentity: &tfsdk.ResourceIdentity{
+		Identity: &tfsdk.ResourceIdentity{
 			Schema: req.IdentitySchema,
 			// Raw is intentionally not set.
 		},
@@ -215,7 +183,7 @@ func (s *Server) UpgradeResourceIdentity(ctx context.Context, req *UpgradeResour
 		return
 	}
 
-	if upgradeResourceIdentityResponse.UpgradedIdentity.Raw.Type() == nil || upgradeResourceIdentityResponse.UpgradedIdentity.Raw.IsNull() {
+	if upgradeResourceIdentityResponse.Identity.Raw.Type() == nil || upgradeResourceIdentityResponse.Identity.Raw.IsNull() {
 		resp.Diagnostics.AddError(
 			"Missing Upgraded Resource Identity",
 			fmt.Sprintf("After attempting a resource Identity upgrade to version %d, the provider did not return any Identity data. ", req.Version)+
@@ -225,16 +193,16 @@ func (s *Server) UpgradeResourceIdentity(ctx context.Context, req *UpgradeResour
 		return
 	}
 
-	if upgradeResourceIdentityResponse.UpgradedIdentity != nil {
-		logging.FrameworkTrace(ctx, "UpgradeResourceIdentityResponse Raw State set, overriding State")
+	if upgradeResourceIdentityResponse.Identity.Raw.Type() != nil {
+		logging.FrameworkTrace(ctx, "UpgradeResourceIdentityResponse RawState set, overriding State")
 
 		resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
 			Schema: req.IdentitySchema,
-			Raw:    upgradeResourceIdentityResponse.UpgradedIdentity.Raw,
+			Raw:    upgradeResourceIdentityResponse.Identity.Raw,
 		}
 
 		return
 	}
 
-	resp.UpgradedIdentity = upgradeResourceIdentityResponse.UpgradedIdentity
+	resp.UpgradedIdentity = upgradeResourceIdentityResponse.Identity
 }
diff --git a/internal/fwserver/server_upgraderesourceidentity_test.go b/internal/fwserver/server_upgraderesourceidentity_test.go
@@ -139,7 +139,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 										Id: id,
 									}
 
-									resp.Diagnostics.Append(resp.UpgradedIdentity.Set(ctx, upgradedIdentityData)...)
+									resp.Diagnostics.Append(resp.Identity.Set(ctx, upgradedIdentityData)...)
 								},
 							},
 						}
@@ -221,7 +221,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 										return
 									}
 
-									resp.UpgradedIdentity = ResourceIdentity
+									resp.Identity = ResourceIdentity
 								},
 							},
 						}
@@ -273,7 +273,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 										Schema: testIdentitySchema,
 									}
 
-									resp.UpgradedIdentity = &ResourceIdentity
+									resp.Identity = &ResourceIdentity
 								},
 							},
 						}
@@ -449,7 +449,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 										Id: id,
 									}
 
-									resp.Diagnostics.Append(resp.UpgradedIdentity.Set(ctx, upgradedIdentityData)...)
+									resp.Diagnostics.Append(resp.Identity.Set(ctx, upgradedIdentityData)...)
 								},
 							},
 						}
@@ -500,115 +500,6 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 				},
 			},
 		},
-		/*"Version-current-flatmap": { // TODO: See if we need to add these tests back
-			server: &fwserver.Server{
-				Provider: &testprovider.Provider{
-					ResourcesMethod: func(_ context.Context) []func() resource.Resource {
-						return []func() resource.Resource{
-							func() resource.Resource {
-								return &testprovider.ResourceWithUpgradeIdentity{
-
-									Resource: &testprovider.Resource{},
-									UpgradeResourceIdentityMethod: func(ctx context.Context) map[int64]resource.IdentityUpgrader {
-										return map[int64]resource.IdentityUpgrader{
-											0: {
-												IdentityUpgrader: func(ctx context.Context, req resource.UpgradeResourceIdentityRequest, resp *resource.UpgradeResourceIdentityResponse) {
-													// Purposfully not setting resp.ResourceIdentity or resp.UpgradedIdentity
-												},
-											},
-										}
-									},
-								}
-							},
-						}
-					},
-				},
-			},
-			request: &fwserver.UpgradeResourceIdentityRequest{
-				RawState: &tfprotov6.RawState{
-					Flatmap: map[string]string{
-						"flatmap": "is not supported",
-					},
-				},
-				IdentitySchema: testIdentitySchema,
-				Resource:       &testprovider.ResourceWithUpgradeIdentity{},
-				Version:        1, // Must match current tfsdk.Schema version to trigger framework implementation
-			},
-			expectedResponse: &fwserver.UpgradeResourceIdentityResponse{
-				Diagnostics: diag.Diagnostics{
-					diag.NewErrorDiagnostic(
-						"Unable to Read Previously Saved Identity for UpgradeResourceIdentity",
-						"There was an error reading the saved resource Identity using the current resource schema.\n\n"+
-							"If this resource Identity was last refreshed with Terraform CLI 0.11 and earlier, it must be refreshed or applied with an older provider version first. "+
-							"If you manually modified the resource Identity, you will need to manually modify it to match the current resource schema. "+
-							"Otherwise, please report this to the provider developer:\n\n"+
-							"flatmap Identitys cannot be unmarshaled, only Identitys written by Terraform 0.12 and higher can be unmarshaled",
-					),
-				},
-			},
-		},
-		"Version-current-json-match": {
-			server: &fwserver.Server{
-				Provider: &testprovider.Provider{
-					ResourcesMethod: func(_ context.Context) []func() resource.Resource {
-						return []func() resource.Resource{
-							func() resource.Resource {
-								return &testprovider.ResourceWithUpgradeIdentity{
-									Resource: &testprovider.Resource{},
-									UpgradeResourceIdentityMethod: func(ctx context.Context) map[int64]resource.IdentityUpgrader {
-										return map[int64]resource.IdentityUpgrader{
-											0: {
-												IdentityUpgrader: func(ctx context.Context, req resource.UpgradeResourceIdentityRequest, resp *resource.UpgradeResourceIdentityResponse) {
-													// Purposfully not setting resp.ResourceIdentity or resp.UpgradedIdentity
-												},
-											},
-										}
-									},
-								}
-							},
-						}
-					},
-				},
-			},
-			request: &fwserver.UpgradeResourceIdentityRequest{
-				RawState: testNewRawState(t, map[string]interface{}{
-					"id": "test-id-value",
-				}),
-				IdentitySchema: testIdentitySchema,
-				Resource:       &testprovider.ResourceWithUpgradeIdentity{},
-				Version:        1, // Must match current tfsdk.Schema version to trigger framework implementation
-			},
-			expectedResponse: &fwserver.UpgradeResourceIdentityResponse{
-				UpgradedIdentity: &tfsdk.ResourceIdentity{
-					Raw: tftypes.NewValue(schemaIdentityType, map[string]tftypes.Value{
-						"id": tftypes.NewValue(tftypes.String, "test-id-value"),
-					}),
-					Schema: testIdentitySchema,
-				},
-			},
-		},
-		"Version-current-json-mismatch": {
-			server: &fwserver.Server{
-				Provider: &testprovider.Provider{},
-			},
-			request: &fwserver.UpgradeResourceIdentityRequest{
-				RawState: testNewRawState(t, map[string]interface{}{
-					"id":                    "test-id-value",
-					"nonexistent_attribute": "value",
-				}),
-				IdentitySchema: testIdentitySchema,
-				Resource:       &testprovider.ResourceWithUpgradeIdentity{},
-				Version:        1, // Must match current tfsdk.IdentitySchema version to trigger framework implementation
-			},
-			expectedResponse: &fwserver.UpgradeResourceIdentityResponse{
-				UpgradedIdentity: &tfsdk.ResourceIdentity{
-					Raw: tftypes.NewValue(schemaIdentityType, map[string]tftypes.Value{
-						"id": tftypes.NewValue(tftypes.String, "test-id-value"),
-					}),
-					Schema: testIdentitySchema,
-				},
-			},
-		},*/
 		"Version-not-implemented": {
 			server: &fwserver.Server{
 				Provider: &testprovider.Provider{},
diff --git a/internal/proto5server/server_upgraderesourceidentity.go b/internal/proto5server/server_upgraderesourceidentity.go
@@ -15,7 +15,6 @@ import (
 
 // UpgradeResourceIdentity satisfies the tfprotov5.ProviderServer interface.
 func (s *Server) UpgradeResourceIdentity(ctx context.Context, proto5Req *tfprotov5.UpgradeResourceIdentityRequest) (*tfprotov5.UpgradeResourceIdentityResponse, error) {
-	// panic("unimplemented") // TODO:ResourceIdentity: implement
 	ctx = s.registerContext(ctx)
 	ctx = logging.InitContext(ctx)
 
diff --git a/internal/proto5server/server_upgraderesourceidentity_test.go b/internal/proto5server/server_upgraderesourceidentity_test.go
@@ -42,6 +42,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 				RequiredForImport: true,
 			},
 		},
+		Version: 1, // Must be above 0
 	}
 
 	ctx := context.Background()
@@ -91,7 +92,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 														}
 
 														// Prevent Missing Upgraded Resource Identity error
-														resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
+														resp.Identity = &tfsdk.ResourceIdentity{
 															Raw: tftypes.NewValue(testIdentityType, map[string]tftypes.Value{
 																"test_id": tftypes.NewValue(tftypes.String, "test-id-value"),
 															}),
@@ -241,7 +242,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 											return map[int64]resource.IdentityUpgrader{
 												0: {
 													IdentityUpgrader: func(_ context.Context, _ resource.UpgradeResourceIdentityRequest, resp *resource.UpgradeResourceIdentityResponse) {
-														resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
+														resp.Identity = &tfsdk.ResourceIdentity{
 															Raw: tftypes.NewValue(testIdentityType, map[string]tftypes.Value{
 																"test_id": tftypes.NewValue(tftypes.String, "test-id-value"),
 															}),
diff --git a/internal/proto6server/server_upgraderesourceidentity.go b/internal/proto6server/server_upgraderesourceidentity.go
@@ -15,7 +15,6 @@ import (
 
 // UpgradeResourceIdentity satisfies the tfprotov6.ProviderServer interface.
 func (s *Server) UpgradeResourceIdentity(ctx context.Context, proto6Req *tfprotov6.UpgradeResourceIdentityRequest) (*tfprotov6.UpgradeResourceIdentityResponse, error) {
-	// panic("unimplemented") // TODO:ResourceIdentity: implement
 	ctx = s.registerContext(ctx)
 	ctx = logging.InitContext(ctx)
 
diff --git a/internal/proto6server/server_upgraderesourceidentity_test.go b/internal/proto6server/server_upgraderesourceidentity_test.go
@@ -42,6 +42,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 				RequiredForImport: true,
 			},
 		},
+		Version: 1, // Must be above 0
 	}
 
 	ctx := context.Background()
@@ -91,7 +92,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 														}
 
 														// Prevent Missing Upgraded Resource Identity error
-														resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
+														resp.Identity = &tfsdk.ResourceIdentity{
 															Raw: tftypes.NewValue(testIdentityType, map[string]tftypes.Value{
 																"test_id": tftypes.NewValue(tftypes.String, "test-id-value"),
 															}),
@@ -241,7 +242,7 @@ func TestServerUpgradeResourceIdentity(t *testing.T) {
 											return map[int64]resource.IdentityUpgrader{
 												0: {
 													IdentityUpgrader: func(_ context.Context, _ resource.UpgradeResourceIdentityRequest, resp *resource.UpgradeResourceIdentityResponse) {
-														resp.UpgradedIdentity = &tfsdk.ResourceIdentity{
+														resp.Identity = &tfsdk.ResourceIdentity{
 															Raw: tftypes.NewValue(testIdentityType, map[string]tftypes.Value{
 																"test_id": tftypes.NewValue(tftypes.String, "test-id-value"),
 															}),
diff --git a/resource/identity_upgrader.go b/resource/identity_upgrader.go
@@ -9,41 +9,39 @@ import (
 	"github.com/hashicorp/terraform-plugin-go/tftypes"
 )
 
-// Implementation handler for a UpgradeState operation.
+// Implementation handler for an UpgradeIdentity operation.
 //
-// This is used to encapsulate all upgrade logic from a prior state to the
-// current schema version when a Resource implements the
-// ResourceWithUpgradeState interface.
+// This is used to encapsulate all upgrade logic from a prior identity to the
+// current version when a Resource implements the
+// ResourceWithUpgradeIdentity interface.
 type IdentityUpgrader struct {
-	// Schema information for the prior state version. While not required,
-	// setting this will populate the UpgradeStateRequest type State
+	// Schema information for the prior identity version. While not required,
+	// setting this will populate the UpgradeResourceIdentityRequest type Identity
 	// field similar to other Resource data types. This allows for easier data
 	// handling such as calling Get() or GetAttribute().
 	//
-	// If not set, prior state data is available in the
-	// UpgradeResourceStateRequest type RawState field.
+	// If not set, prior identity data is available in the
+	// UpgradeResourceIdentityRequest type RawIdentity field.
 	PriorSchema *identityschema.Schema
 
 	// Version is the version schema that this Upgrader will handle, converting
 	// it to Version+1.
 	Version int64
 
-	// Type describes the schema that this function can upgrade. Type is
-	// required to decode the schema if the state was stored in a legacy
-	// flatmap format.
+	// Type describes the schema that this function can upgrade.
 	Type tftypes.Type
 
-	// Provider defined logic for upgrading a resource state from the prior
-	// state version to the current schema version.
+	// Provider defined logic for upgrading a resource identity from the prior
+	// identity version to the current schema version.
 	//
 	// The context.Context parameter contains framework-defined loggers and
 	// supports request cancellation.
 	//
-	// The UpgradeStateRequest parameter contains the prior state data.
-	// If PriorSchema was set, the State field will be available. Otherwise,
-	// the RawState must be used.
+	// The UpgradeResourceIdentityRequest parameter contains the prior identity data.
+	// If PriorSchema was set, the Identity field will be available. Otherwise,
+	// the RawIdentity must be used.
 	//
-	// The UpgradeStateResponse parameter should contain the upgraded
-	// state data and can be used to signal any logic warnings or errors.
+	// The UpgradeResourceIdentityResponse parameter should contain the upgraded
+	// identity data and can be used to signal any logic warnings or errors.
 	IdentityUpgrader func(context.Context, UpgradeResourceIdentityRequest, *UpgradeResourceIdentityResponse)
 }
diff --git a/resource/resource.go b/resource/resource.go
diff --git a/resource/upgrade_identity.go b/resource/upgrade_identity.go
