TFECO-9166: Add support for ResourceIdentity#SchemaFunc to allow providers to save RAM (#1459)

* Add support for ResourceIdentity#SchemaFunc to allow providers to save RAM

* handle nil ResourceIdentity for SchemaMap() and simplify code

Author: ansgarm

helper/schema/core_schema.go

@@ -389,11 +389,11 @@ func (r *Resource) CoreIdentitySchema() (*configschema.Block, error) {
 }
 
 func (r *Resource) coreIdentitySchema() (*configschema.Block, error) {
-	if r.Identity == nil || r.Identity.Schema == nil {
+	if r.Identity.SchemaMap() == nil {
 		return nil, fmt.Errorf("resource does not have an identity schema")
 	}
 	// while there is schemaMapWithIdentity, we don't need to use it here
 	// as we're only interested in the existing CoreConfigSchema() method
 	// to convert our schema
-	return schemaMap(r.Identity.Schema).CoreConfigSchema(), nil
+	return schemaMap(r.Identity.SchemaMap()).CoreConfigSchema(), nil
 }

helper/schema/grpc_provider_test.go

@@ -3319,25 +3319,27 @@ func TestGRPCProviderServerGetResourceIdentitySchemas(t *testing.T) {
 					},
 					"test_resource2": {
 						Identity: &ResourceIdentity{
-							Schema: map[string]*Schema{
-								"test2": {
-									Type:              TypeString,
-									RequiredForImport: false,
-									OptionalForImport: true,
-									Description:       "test resource 2",
-								},
-								"test2-2": {
-									Type:              TypeList,
-									RequiredForImport: false,
-									OptionalForImport: true,
-									Description:       "test resource 2-2",
-								},
-								"test2-3": {
-									Type:              TypeInt,
-									RequiredForImport: false,
-									OptionalForImport: true,
-									Description:       "test resource 2-3",
-								},
+							SchemaFunc: func() map[string]*Schema {
+								return map[string]*Schema{
+									"test2": {
+										Type:              TypeString,
+										RequiredForImport: false,
+										OptionalForImport: true,
+										Description:       "test resource 2",
+									},
+									"test2-2": {
+										Type:              TypeList,
+										RequiredForImport: false,
+										OptionalForImport: true,
+										Description:       "test resource 2-2",
+									},
+									"test2-3": {
+										Type:              TypeInt,
+										RequiredForImport: false,
+										OptionalForImport: true,
+										Description:       "test resource 2-3",
+									},
+								}
 							},
 						},
 					},

helper/schema/resource.go

@@ -728,11 +728,7 @@ func (r *Resource) ShimInstanceStateFromValue(state cty.Value) (*terraform.Insta
 
 	// We now rebuild the state through the ResourceData, so that the set indexes
 	// match what helper/schema expects.
-	var identity map[string]*Schema
-	if r.Identity != nil {
-		identity = r.Identity.Schema
-	}
-	data, err := schemaMapWithIdentity{r.SchemaMap(), identity}.Data(s, nil)
+	data, err := schemaMapWithIdentity{r.SchemaMap(), r.Identity.SchemaMap()}.Data(s, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -905,11 +901,7 @@ func (r *Resource) Apply(
 	s *terraform.InstanceState,
 	d *terraform.InstanceDiff,
 	meta interface{}) (*terraform.InstanceState, diag.Diagnostics) {
-	var identity map[string]*Schema
-	if r.Identity != nil {
-		identity = r.Identity.Schema
-	}
-	schema := schemaMapWithIdentity{r.SchemaMap(), identity}
+	schema := schemaMapWithIdentity{r.SchemaMap(), r.Identity.SchemaMap()}
 	data, err := schema.Data(s, d)
 	if err != nil {
 		return s, diag.FromErr(err)
@@ -1033,12 +1025,8 @@ func (r *Resource) SimpleDiff(
 	c *terraform.ResourceConfig,
 	meta interface{}) (*terraform.InstanceDiff, error) {
 
-	var identity map[string]*Schema
-	if r.Identity != nil {
-		identity = r.Identity.Schema
-	}
 	// TODO: figure out if it makes sense to be able to set identity in CustomizeDiff at all
-	instanceDiff, err := schemaMapWithIdentity{r.SchemaMap(), identity}.Diff(ctx, s, c, r.CustomizeDiff, meta, false)
+	instanceDiff, err := schemaMapWithIdentity{r.SchemaMap(), r.Identity.SchemaMap()}.Diff(ctx, s, c, r.CustomizeDiff, meta, false)
 	if err != nil {
 		return instanceDiff, err
 	}
@@ -1127,11 +1115,7 @@ func (r *Resource) RefreshWithoutUpgrade(
 		}
 	}
 
-	var identity map[string]*Schema
-	if r.Identity != nil {
-		identity = r.Identity.Schema
-	}
-	schema := schemaMapWithIdentity{r.SchemaMap(), identity}
+	schema := schemaMapWithIdentity{r.SchemaMap(), r.Identity.SchemaMap()}
 
 	if r.Exists != nil {
 		// Make a copy of data so that if it is modified it doesn't
@@ -1442,7 +1426,7 @@ func (r *Resource) Data(s *terraform.InstanceState) *ResourceData {
 func (r *Resource) TestResourceData() *ResourceData {
 	return &ResourceData{
 		schema:         r.SchemaMap(),
-		identitySchema: r.Identity.Schema,
+		identitySchema: r.Identity.SchemaMap(),
 	}
 }
 
@@ -1489,11 +1473,11 @@ func (r *ResourceIdentity) InternalIdentityValidate() error {
 		return fmt.Errorf(`The resource identity is empty`)
 	}
 
-	if len(r.Schema) == 0 {
+	if len(r.SchemaMap()) == 0 {
 		return fmt.Errorf(`The resource identity schema is empty`)
 	}
 
-	for k, v := range r.Schema {
+	for k, v := range r.SchemaMap() {
 		if !v.OptionalForImport && !v.RequiredForImport {
 			return fmt.Errorf(`OptionalForImport or RequiredForImport must be set for resource identity`)
 		}

helper/schema/resource_identity.go

@@ -37,6 +37,12 @@ type ResourceIdentity struct {
 	// previous resource schemas.
 	Schema map[string]*Schema
 
+	// SchemaFunc is an optional function that returns the schema for the
+	// identity. Use this field instead of Schema on resource identity
+	// declarations to prevent storing all identity schema information in
+	// memory for the lifecycle of a provider.
+	SchemaFunc func() map[string]*Schema
+
 	// New struct, will be similar to (Resource).StateUpgraders
 	IdentityUpgraders []IdentityUpgrader
 }
@@ -70,3 +76,18 @@ type ResourceIdentity struct {
 // identity schema version data for a managed resource instance. Values must
 // align to the typing mentioned above.
 type ResourceIdentityUpgradeFunc func(ctx context.Context, rawState map[string]interface{}, meta interface{}) (map[string]interface{}, error)
+
+// SchemaMap returns the schema information for this Resource whether it is
+// defined via the SchemaFunc field or Schema field. The SchemaFunc field, if
+// defined, takes precedence over the Schema field.
+func (ri *ResourceIdentity) SchemaMap() map[string]*Schema {
+	if ri == nil {
+		return nil
+	}
+
+	if ri.SchemaFunc != nil {
+		return ri.SchemaFunc()
+	}
+
+	return ri.Schema
+}

helper/schema/resource_identity_test.go

@@ -0,0 +1,13 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package schema
+
+import "testing"
+
+func TestResourceIdentity_SchemaMap_handles_nil_identity(t *testing.T) {
+	var ri *ResourceIdentity
+	if ri.SchemaMap() != nil {
+		t.Fatal("expected nil schema map")
+	}
+}

helper/schema/shims.go

@@ -43,11 +43,7 @@ func diffFromValues(ctx context.Context, prior, planned, config cty.Value, res *
 	removeConfigUnknowns(cfg.Config)
 	removeConfigUnknowns(cfg.Raw)
 
-	var identity map[string]*Schema
-	if res.Identity != nil {
-		identity = res.Identity.Schema
-	}
-	diff, err := schemaMapWithIdentity{res.SchemaMap(), identity}.Diff(ctx, instanceState, cfg, cust, nil, false)
+	diff, err := schemaMapWithIdentity{res.SchemaMap(), res.Identity.SchemaMap()}.Diff(ctx, instanceState, cfg, cust, nil, false)
 	if err != nil {
 		return nil, err
 	}