helper/schema: Add SchemaFunc field to Resource type (#1218)

Reference: #1217

This change introduces a new `SchemaFunc` field and `SchemaMap` method to the `Resource` type. Currently, the `Schema` field data of all `Resource` is held in memory for the lifecycle of the provider server, which is problematic for providers with many resources and/or larger schemas in resources. The new field enables provider developers to swap pieces of resident memory usage for the slight additional computation time of reconstituting the data when necessary.

Callers directly referencing the exported `Schema` field should switch to referencing the `SchemaMap` method, which returns the result of `SchemaFunc` or `Schema` in that preference order. To ensure internal usage was migrated to the new method, this change was performed by temporarily commenting out the `Schema` field itself with broken references in non-testing code migrated to the method. The `Schema` field is not marked as deprecated via Go documentation comment as this would introduce a major ecosystem burden to migrate with generally little benefit for most use cases.

The `Resource` type `InternalValidate` method has been updated to return an error if both `Schema` and `SchemaFunc` are defined.

Provider developers are encouraged to migrate resources to the terraform-plugin-framework, as it already behaves in a manner similar to `SchemaFunc` by nature of resource schema data being behind a method call, amongst many of the other benefits outlined at https://developer.hashicorp.com/terraform/plugin/framework-benefits.

.changes/unreleased/ENHANCEMENTS-20230623-102412.yaml

@@ -0,0 +1,6 @@
+kind: ENHANCEMENTS
+body: 'helper/schema: Added `Resource` type `SchemaFunc` field and `SchemaMap` method,
+  which can reduce resident memory usage with large schemas'
+time: 2023-06-23T10:24:12.025356-04:00
+custom:
+  Issue: "1217"

.changes/unreleased/NOTES-20230623-102528.yaml

@@ -0,0 +1,7 @@
+kind: NOTES
+body: 'helper/schema: Consumers directly referencing the `Resource` type `Schema`
+  field should switch to the `SchemaMap` method to ensure new `SchemaFunc` field
+  data is properly retrieved'
+time: 2023-06-23T10:25:28.864812-04:00
+custom:
+  Issue: "1217"

helper/schema/core_schema.go

@@ -367,5 +367,5 @@ func (r *Resource) CoreConfigSchema() *configschema.Block {
 }
 
 func (r *Resource) coreConfigSchema() *configschema.Block {
-	return schemaMap(r.Schema).CoreConfigSchema()
+	return schemaMap(r.SchemaMap()).CoreConfigSchema()
 }

helper/schema/field_reader.go

@@ -84,7 +84,7 @@ func addrToSchema(addr []string, schemaMap map[string]*Schema) []*Schema {
 			case *Resource:
 				current = &Schema{
 					Type: typeObject,
-					Elem: v.Schema,
+					Elem: v.SchemaMap(),
 				}
 			case *Schema:
 				current = v

helper/schema/field_reader_config.go

@@ -303,7 +303,7 @@ func (r *ConfigFieldReader) hasComputedSubKeys(key string, schema *Schema) bool
 
 	switch t := schema.Elem.(type) {
 	case *Resource:
-		for k, schema := range t.Schema {
+		for k, schema := range t.SchemaMap() {
 			if r.Config.IsComputed(prefix + k) {
 				return true
 			}

helper/schema/grpc_provider.go

@@ -1303,7 +1303,7 @@ func stripResourceModifiers(r *Resource) *Resource {
 	newResource.CustomizeDiff = nil
 	newResource.Schema = map[string]*Schema{}
 
-	for k, s := range r.Schema {
+	for k, s := range r.SchemaMap() {
 		newResource.Schema[k] = stripSchema(s)
 	}
 

helper/schema/grpc_provider_test.go

@@ -807,6 +807,23 @@ func TestApplyResourceChange(t *testing.T) {
 				},
 			},
 		},
+		{
+			Description: "CreateContext_SchemaFunc",
+			TestResource: &Resource{
+				SchemaFunc: func() map[string]*Schema {
+					return map[string]*Schema{
+						"id": {
+							Type:     TypeString,
+							Computed: true,
+						},
+					}
+				},
+				CreateContext: func(_ context.Context, rd *ResourceData, _ interface{}) diag.Diagnostics {
+					rd.SetId("bar") // expected in response
+					return nil
+				},
+			},
+		},
 	}
 
 	for _, testCase := range testCases {
@@ -1254,6 +1271,47 @@ func TestReadDataSource(t *testing.T) {
 				},
 			},
 		},
+		"SchemaFunc": {
+			server: NewGRPCProviderServer(&Provider{
+				DataSourcesMap: map[string]*Resource{
+					"test": {
+						SchemaFunc: func() map[string]*Schema {
+							return map[string]*Schema{
+								"id": {
+									Type:     TypeString,
+									Computed: true,
+								},
+							}
+						},
+						ReadContext: func(ctx context.Context, d *ResourceData, meta interface{}) diag.Diagnostics {
+							d.SetId("test-id")
+							return nil
+						},
+					},
+				},
+			}),
+			req: &tfprotov5.ReadDataSourceRequest{
+				TypeName: "test",
+				Config: &tfprotov5.DynamicValue{
+					MsgPack: mustMsgpackMarshal(
+						cty.EmptyObject,
+						cty.NullVal(cty.EmptyObject),
+					),
+				},
+			},
+			expected: &tfprotov5.ReadDataSourceResponse{
+				State: &tfprotov5.DynamicValue{
+					MsgPack: mustMsgpackMarshal(
+						cty.Object(map[string]cty.Type{
+							"id": cty.String,
+						}),
+						cty.ObjectVal(map[string]cty.Value{
+							"id": cty.StringVal("test-id"),
+						}),
+					),
+				},
+			},
+		},
 		"null-object": {
 			server: NewGRPCProviderServer(&Provider{
 				DataSourcesMap: map[string]*Resource{

helper/schema/resource.go

@@ -56,13 +56,25 @@ var ReservedResourceFields = []string{
 // being implemented.
 type Resource struct {
 	// Schema is the structure and type information for this component. This
-	// field is required for all Resource concepts.
+	// field, or SchemaFunc, is required for all Resource concepts. To prevent
+	// storing all schema information in memory for the lifecycle of a provider,
+	// use SchemaFunc instead.
 	//
 	// The keys of this map are the names used in a practitioner configuration,
 	// such as the attribute or block name. The values describe the structure
 	// and type information of that attribute or block.
 	Schema map[string]*Schema
 
+	// SchemaFunc is the structure and type information for this component. This
+	// field, or Schema, is required for all Resource concepts. Use this field
+	// instead of Schema on top level Resource declarations to prevent storing
+	// all schema information in memory for the lifecycle of a provider.
+	//
+	// The keys of this map are the names used in a practitioner configuration,
+	// such as the attribute or block name. The values describe the structure
+	// and type information of that attribute or block.
+	SchemaFunc func() map[string]*Schema
+
 	// SchemaVersion is the version number for this resource's Schema
 	// definition. This field is only valid when the Resource is a managed
 	// resource.
@@ -585,6 +597,17 @@ type Resource struct {
 	UseJSONNumber bool
 }
 
+// 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 (r *Resource) SchemaMap() map[string]*Schema {
+	if r.SchemaFunc != nil {
+		return r.SchemaFunc()
+	}
+
+	return r.Schema
+}
+
 // ShimInstanceStateFromValue converts a cty.Value to a
 // terraform.InstanceState.
 func (r *Resource) ShimInstanceStateFromValue(state cty.Value) (*terraform.InstanceState, error) {
@@ -594,7 +617,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.
-	data, err := schemaMap(r.Schema).Data(s, nil)
+	data, err := schemaMap(r.SchemaMap()).Data(s, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -767,7 +790,8 @@ func (r *Resource) Apply(
 	s *terraform.InstanceState,
 	d *terraform.InstanceDiff,
 	meta interface{}) (*terraform.InstanceState, diag.Diagnostics) {
-	data, err := schemaMap(r.Schema).Data(s, d)
+	schema := schemaMap(r.SchemaMap())
+	data, err := schema.Data(s, d)
 	if err != nil {
 		return s, diag.FromErr(err)
 	}
@@ -824,7 +848,7 @@ func (r *Resource) Apply(
 		}
 
 		// Reset the data to be stateless since we just destroyed
-		data, err = schemaMap(r.Schema).Data(nil, d)
+		data, err = schema.Data(nil, d)
 		if err != nil {
 			return nil, append(diags, diag.FromErr(err)...)
 		}
@@ -868,7 +892,7 @@ func (r *Resource) Diff(
 		return nil, fmt.Errorf("[ERR] Error decoding timeout: %s", err)
 	}
 
-	instanceDiff, err := schemaMap(r.Schema).Diff(ctx, s, c, r.CustomizeDiff, meta, true)
+	instanceDiff, err := schemaMap(r.SchemaMap()).Diff(ctx, s, c, r.CustomizeDiff, meta, true)
 	if err != nil {
 		return instanceDiff, err
 	}
@@ -890,7 +914,7 @@ func (r *Resource) SimpleDiff(
 	c *terraform.ResourceConfig,
 	meta interface{}) (*terraform.InstanceDiff, error) {
 
-	instanceDiff, err := schemaMap(r.Schema).Diff(ctx, s, c, r.CustomizeDiff, meta, false)
+	instanceDiff, err := schemaMap(r.SchemaMap()).Diff(ctx, s, c, r.CustomizeDiff, meta, false)
 	if err != nil {
 		return instanceDiff, err
 	}
@@ -915,7 +939,7 @@ func (r *Resource) SimpleDiff(
 
 // Validate validates the resource configuration against the schema.
 func (r *Resource) Validate(c *terraform.ResourceConfig) diag.Diagnostics {
-	diags := schemaMap(r.Schema).Validate(c)
+	diags := schemaMap(r.SchemaMap()).Validate(c)
 
 	if r.DeprecationMessage != "" {
 		diags = append(diags, diag.Diagnostic{
@@ -937,7 +961,7 @@ func (r *Resource) ReadDataApply(
 ) (*terraform.InstanceState, diag.Diagnostics) {
 	// Data sources are always built completely from scratch
 	// on each read, so the source state is always nil.
-	data, err := schemaMap(r.Schema).Data(nil, d)
+	data, err := schemaMap(r.SchemaMap()).Data(nil, d)
 	if err != nil {
 		return nil, diag.FromErr(err)
 	}
@@ -978,10 +1002,12 @@ func (r *Resource) RefreshWithoutUpgrade(
 		}
 	}
 
+	schema := schemaMap(r.SchemaMap())
+
 	if r.Exists != nil {
 		// Make a copy of data so that if it is modified it doesn't
 		// affect our Read later.
-		data, err := schemaMap(r.Schema).Data(s, nil)
+		data, err := schema.Data(s, nil)
 		if err != nil {
 			return s, diag.FromErr(err)
 		}
@@ -1004,7 +1030,7 @@ func (r *Resource) RefreshWithoutUpgrade(
 		}
 	}
 
-	data, err := schemaMap(r.Schema).Data(s, nil)
+	data, err := schema.Data(s, nil)
 	if err != nil {
 		return s, diag.FromErr(err)
 	}
@@ -1023,7 +1049,7 @@ func (r *Resource) RefreshWithoutUpgrade(
 		state = nil
 	}
 
-	schemaMap(r.Schema).handleDiffSuppressOnRefresh(ctx, s, state)
+	schema.handleDiffSuppressOnRefresh(ctx, s, state)
 	return r.recordCurrentSchemaVersion(state), diags
 }
 
@@ -1069,13 +1095,14 @@ func (r *Resource) InternalValidate(topSchemaMap schemaMap, writable bool) error
 		}
 	}
 
+	schema := schemaMap(r.SchemaMap())
 	tsm := topSchemaMap
 
 	if r.isTopLevel() && writable {
 		// All non-Computed attributes must be ForceNew if Update is not defined
 		if !r.updateFuncSet() {
 			nonForceNewAttrs := make([]string, 0)
-			for k, v := range r.Schema {
+			for k, v := range schema {
 				if !v.ForceNew && !v.Computed {
 					nonForceNewAttrs = append(nonForceNewAttrs, k)
 				}
@@ -1086,19 +1113,19 @@ func (r *Resource) InternalValidate(topSchemaMap schemaMap, writable bool) error
 			}
 		} else {
 			nonUpdateableAttrs := make([]string, 0)
-			for k, v := range r.Schema {
+			for k, v := range schema {
 				if v.ForceNew || v.Computed && !v.Optional {
 					nonUpdateableAttrs = append(nonUpdateableAttrs, k)
 				}
 			}
-			updateableAttrs := len(r.Schema) - len(nonUpdateableAttrs)
+			updateableAttrs := len(schema) - len(nonUpdateableAttrs)
 			if updateableAttrs == 0 {
 				return fmt.Errorf(
 					"All fields are ForceNew or Computed w/out Optional, Update is superfluous")
 			}
 		}
 
-		tsm = schemaMap(r.Schema)
+		tsm = schema
 
 		// Destroy, and Read are required
 		if !r.readFuncSet() {
@@ -1157,14 +1184,18 @@ func (r *Resource) InternalValidate(topSchemaMap schemaMap, writable bool) error
 
 	// Data source
 	if r.isTopLevel() && !writable {
-		tsm = schemaMap(r.Schema)
+		tsm = schema
 		for k := range tsm {
 			if isReservedDataSourceFieldName(k) {
 				return fmt.Errorf("%s is a reserved field name", k)
 			}
 		}
 	}
 
+	if r.SchemaFunc != nil && r.Schema != nil {
+		return fmt.Errorf("SchemaFunc and Schema should not both be set")
+	}
+
 	// check context funcs are not set alongside their nonctx counterparts
 	if r.CreateContext != nil && r.Create != nil {
 		return fmt.Errorf("CreateContext and Create should not both be set")
@@ -1207,7 +1238,7 @@ func (r *Resource) InternalValidate(topSchemaMap schemaMap, writable bool) error
 		return fmt.Errorf("Delete and DeleteWithoutTimeout should not both be set")
 	}
 
-	return schemaMap(r.Schema).InternalValidate(tsm)
+	return schema.InternalValidate(tsm)
 }
 
 func isReservedDataSourceFieldName(name string) bool {
@@ -1254,7 +1285,7 @@ func isReservedResourceFieldName(name string) bool {
 //
 // This function is useful for unit tests and ResourceImporter functions.
 func (r *Resource) Data(s *terraform.InstanceState) *ResourceData {
-	result, err := schemaMap(r.Schema).Data(s, nil)
+	result, err := schemaMap(r.SchemaMap()).Data(s, nil)
 	if err != nil {
 		// At the time of writing, this isn't possible (Data never returns
 		// non-nil errors). We panic to find this in the future if we have to.
@@ -1281,7 +1312,7 @@ func (r *Resource) Data(s *terraform.InstanceState) *ResourceData {
 // TODO: May be able to be removed with the above ResourceData function.
 func (r *Resource) TestResourceData() *ResourceData {
 	return &ResourceData{
-		schema: r.Schema,
+		schema: r.SchemaMap(),
 	}
 }
 

helper/schema/resource_test.go

@@ -1080,6 +1080,51 @@ func TestResourceInternalValidate(t *testing.T) {
 			true,
 			true,
 		},
+		27: { // Non-Writable SchemaFunc and Schema should not both be set
+			In: &Resource{
+				Schema: map[string]*Schema{
+					"test": {
+						Type:     TypeString,
+						Required: true,
+					},
+				},
+				SchemaFunc: func() map[string]*Schema {
+					return map[string]*Schema{
+						"test": {
+							Type:     TypeString,
+							Required: true,
+						},
+					}
+				},
+				Read: Noop,
+			},
+			Writable: false,
+			Err:      true,
+		},
+		28: { // Writable SchemaFunc and Schema should not both be set
+			In: &Resource{
+				Schema: map[string]*Schema{
+					"test": {
+						Type:     TypeString,
+						Required: true,
+					},
+				},
+				SchemaFunc: func() map[string]*Schema {
+					return map[string]*Schema{
+						"test": {
+							Type:     TypeString,
+							Required: true,
+						},
+					}
+				},
+				Create: Noop,
+				Read:   Noop,
+				Update: Noop,
+				Delete: Noop,
+			},
+			Writable: true,
+			Err:      true,
+		},
 	}
 
 	for i, tc := range cases {