helper/schema: Add SchemaFunc field to Resource type

.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 {