From 168d84b3c43b9edbeba9c5256e484b82f041db8d Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Tue, 27 Nov 2018 15:30:18 -0800
Subject: [PATCH] core: Make resource type schema versions visible to callers

Previously we were fetching these from the provider but then immediately
discarding the version numbers because the schema API had nowhere to put
them.

To avoid a late-breaking change to the internal structure of
terraform.ProviderSchema (which is constructed directly all over the
tests) we're retaining the resource type schemas in a new map alongside
the existing one with the same keys, rather than just switching to
using the providers.Schema struct directly there.

The methods that return resource type schemas now return two arguments,
intentionally creating a little API friction here so each new caller can
be reminded to think about whether they need to do something with the
schema version, though it can be ignored by many callers.

Since this was a breaking change to the Schemas API anyway, this also
fixes another API wart where there was a separate method for fetching
managed vs. data resource types and thus every caller ended up having a
switch statement on "mode". Now we just accept mode as an argument and
do the switch statement within the single SchemaForResourceType method.
---
 backend/local/backend_plan.go        |  2 +-
 terraform/eval_apply.go              |  2 +-
 terraform/eval_diff.go               | 10 ++---
 terraform/eval_read_data.go          |  6 +--
 terraform/eval_refresh.go            |  2 +-
 terraform/eval_state.go              | 41 +++++++-----------
 terraform/eval_validate.go           |  8 ++--
 terraform/eval_validate_selfref.go   | 14 +------
 terraform/evaluate.go                | 12 +-----
 terraform/evaluate_valid.go          |  9 +---
 terraform/node_resource_abstract.go  |  8 ++--
 terraform/schemas.go                 | 62 +++++++++++++---------------
 terraform/transform_attach_schema.go | 13 ++----
 13 files changed, 70 insertions(+), 119 deletions(-)

diff --git a/backend/local/backend_plan.go b/backend/local/backend_plan.go
index ac1a7701b7..c007edb358 100644
--- a/backend/local/backend_plan.go
+++ b/backend/local/backend_plan.go
@@ -245,7 +245,7 @@ func (b *Local) renderPlan(plan *plans.Plan, schemas *terraform.Schemas) {
 			b.CLI.Output(fmt.Sprintf("(schema missing for %s)\n", rcs.ProviderAddr))
 			continue
 		}
-		rSchema := providerSchema.SchemaForResourceAddr(rcs.Addr.Resource.Resource)
+		rSchema, _ := providerSchema.SchemaForResourceAddr(rcs.Addr.Resource.Resource)
 		if rSchema == nil {
 			// Should never happen
 			b.CLI.Output(fmt.Sprintf("(schema missing for %s)\n", rcs.Addr))
diff --git a/terraform/eval_apply.go b/terraform/eval_apply.go
index e2674cd277..8c09e997a9 100644
--- a/terraform/eval_apply.go
+++ b/terraform/eval_apply.go
@@ -46,7 +46,7 @@ func (n *EvalApply) Eval(ctx EvalContext) (interface{}, error) {
 		state = &states.ResourceInstanceObject{}
 	}
 
-	schema := (*n.ProviderSchema).ResourceTypes[n.Addr.Resource.Type]
+	schema, _ := (*n.ProviderSchema).SchemaForResourceType(n.Addr.Resource.Mode, n.Addr.Resource.Type)
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
diff --git a/terraform/eval_diff.go b/terraform/eval_diff.go
index 50bb5ea3b7..e0f0b04442 100644
--- a/terraform/eval_diff.go
+++ b/terraform/eval_diff.go
@@ -43,10 +43,10 @@ func (n *EvalCheckPlannedChange) Eval(ctx EvalContext) (interface{}, error) {
 	plannedChange := *n.Planned
 	actualChange := *n.Actual
 
-	schema := providerSchema.ResourceTypes[n.Addr.Resource.Type]
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
-		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
+		return nil, fmt.Errorf("provider does not support %q", n.Addr.Resource.Type)
 	}
 
 	var diags tfdiags.Diagnostics
@@ -129,7 +129,7 @@ func (n *EvalDiff) Eval(ctx EvalContext) (interface{}, error) {
 	var diags tfdiags.Diagnostics
 
 	// Evaluate the configuration
-	schema := providerSchema.ResourceTypes[n.Addr.Resource.Type]
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
@@ -833,7 +833,7 @@ func (n *EvalReadDiff) Eval(ctx EvalContext) (interface{}, error) {
 	changes := ctx.Changes()
 	addr := n.Addr.Absolute(ctx.Path())
 
-	schema := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
@@ -894,7 +894,7 @@ func (n *EvalWriteDiff) Eval(ctx EvalContext) (interface{}, error) {
 		panic("inconsistent address and/or deposed key in EvalWriteDiff")
 	}
 
-	schema := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
diff --git a/terraform/eval_read_data.go b/terraform/eval_read_data.go
index 773ac452a0..e9096909eb 100644
--- a/terraform/eval_read_data.go
+++ b/terraform/eval_read_data.go
@@ -83,7 +83,7 @@ func (n *EvalReadData) Eval(ctx EvalContext) (interface{}, error) {
 	config := *n.Config
 	provider := *n.Provider
 	providerSchema := *n.ProviderSchema
-	schema := providerSchema.DataSources[n.Addr.Resource.Type]
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider %q does not support data source %q", n.ProviderAddr.ProviderConfig.Type, n.Addr.Resource.Type)
@@ -322,7 +322,7 @@ func (n *EvalReadDataDiff) Eval(ctx EvalContext) (interface{}, error) {
 	} else {
 		config := *n.Config
 		providerSchema := *n.ProviderSchema
-		schema := providerSchema.DataSources[n.Addr.Resource.Type]
+		schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 		if schema == nil {
 			// Should be caught during validation, so we don't bother with a pretty error here
 			return nil, fmt.Errorf("provider does not support data source %q", n.Addr.Resource.Type)
@@ -440,7 +440,7 @@ func (n *EvalReadDataApply) Eval(ctx EvalContext) (interface{}, error) {
 		return nil, diags.Err()
 	}
 
-	schema := providerSchema.DataSources[n.Addr.Resource.Type]
+	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support data source %q", n.Addr.Resource.Type)
diff --git a/terraform/eval_refresh.go b/terraform/eval_refresh.go
index babf2309fd..ec3822cece 100644
--- a/terraform/eval_refresh.go
+++ b/terraform/eval_refresh.go
@@ -36,7 +36,7 @@ func (n *EvalRefresh) Eval(ctx EvalContext) (interface{}, error) {
 		return nil, diags.ErrWithWarnings()
 	}
 
-	schema := (*n.ProviderSchema).ResourceTypes[n.Addr.Resource.Type]
+	schema, _ := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		return nil, fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Type)
diff --git a/terraform/eval_state.go b/terraform/eval_state.go
index 09ad55190a..9c2856301e 100644
--- a/terraform/eval_state.go
+++ b/terraform/eval_state.go
@@ -49,16 +49,11 @@ func (n *EvalReadState) Eval(ctx EvalContext) (interface{}, error) {
 		return nil, nil
 	}
 
-	// TODO: Update n.ResourceTypeSchema to be a providers.Schema and then
-	// check the version number here and upgrade if necessary.
-	/*
-		if src.SchemaVersion < n.ResourceTypeSchema.Version {
-			// TODO: Implement schema upgrades
-			return nil, fmt.Errorf("schema upgrading is not yet implemented to take state from version %d to version %d", src.SchemaVersion, n.ResourceTypeSchema.Version)
-		}
-	*/
-
-	schema := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
+	schema, currentVersion := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
+	if src.SchemaVersion < currentVersion {
+		// TODO: Implement schema upgrades
+		return nil, fmt.Errorf("schema upgrading is not yet implemented to take state from version %d to version %d", src.SchemaVersion, currentVersion)
+	}
 
 	obj, err := src.Decode(schema.ImpliedType())
 	if err != nil {
@@ -107,16 +102,12 @@ func (n *EvalReadStateDeposed) Eval(ctx EvalContext) (interface{}, error) {
 		return nil, nil
 	}
 
-	// TODO: Update n.ResourceTypeSchema to be a providers.Schema and then
-	// check the version number here and upgrade if necessary.
-	/*
-		if src.SchemaVersion < n.ResourceTypeSchema.Version {
-			// TODO: Implement schema upgrades
-			return nil, fmt.Errorf("schema upgrading is not yet implemented to take state from version %d to version %d", src.SchemaVersion, n.ResourceTypeSchema.Version)
-		}
-	*/
+	schema, currentVersion := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
+	if src.SchemaVersion < currentVersion {
+		// TODO: Implement schema upgrades
+		return nil, fmt.Errorf("schema upgrading is not yet implemented to take state from version %d to version %d", src.SchemaVersion, currentVersion)
+	}
 
-	schema := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
 	obj, err := src.Decode(schema.ImpliedType())
 	if err != nil {
 		return nil, err
@@ -216,16 +207,14 @@ func (n *EvalWriteState) Eval(ctx EvalContext) (interface{}, error) {
 		log.Printf("[TRACE] EvalWriteState: removing current state object for %s", absAddr)
 	}
 
-	// TODO: Update this to use providers.Schema and populate the real
-	// schema version in the second argument to Encode below.
-	schema := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
+	schema, currentVersion := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// It shouldn't be possible to get this far in any real scenario
 		// without a schema, but we might end up here in contrived tests that
 		// fail to set up their world properly.
 		return nil, fmt.Errorf("failed to encode %s in state: no resource type schema available", absAddr)
 	}
-	src, err := obj.Encode(schema.ImpliedType(), 0)
+	src, err := obj.Encode(schema.ImpliedType(), currentVersion)
 	if err != nil {
 		return nil, fmt.Errorf("failed to encode %s in state: %s", absAddr, err)
 	}
@@ -282,16 +271,14 @@ func (n *EvalWriteStateDeposed) Eval(ctx EvalContext) (interface{}, error) {
 		panic("EvalWriteStateDeposed used with no ProviderSchema object")
 	}
 
-	// TODO: Update this to use providers.Schema and populate the real
-	// schema version in the second argument to Encode below.
-	schema := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
+	schema, currentVersion := (*n.ProviderSchema).SchemaForResourceAddr(n.Addr.ContainingResource())
 	if schema == nil {
 		// It shouldn't be possible to get this far in any real scenario
 		// without a schema, but we might end up here in contrived tests that
 		// fail to set up their world properly.
 		return nil, fmt.Errorf("failed to encode %s in state: no resource type schema available", absAddr)
 	}
-	src, err := obj.Encode(schema.ImpliedType(), 0)
+	src, err := obj.Encode(schema.ImpliedType(), currentVersion)
 	if err != nil {
 		return nil, fmt.Errorf("failed to encode %s in state: %s", absAddr, err)
 	}
diff --git a/terraform/eval_validate.go b/terraform/eval_validate.go
index c26a9b90a8..3d79f7989f 100644
--- a/terraform/eval_validate.go
+++ b/terraform/eval_validate.go
@@ -373,8 +373,8 @@ func (n *EvalValidateResource) Eval(ctx EvalContext) (interface{}, error) {
 	// in the provider abstraction.
 	switch mode {
 	case addrs.ManagedResourceMode:
-		schema, exists := schema.ResourceTypes[cfg.Type]
-		if !exists {
+		schema, _ := schema.SchemaForResourceType(mode, cfg.Type)
+		if schema == nil {
 			diags = diags.Append(&hcl.Diagnostic{
 				Severity: hcl.DiagError,
 				Summary:  "Invalid resource type",
@@ -403,8 +403,8 @@ func (n *EvalValidateResource) Eval(ctx EvalContext) (interface{}, error) {
 		}
 
 	case addrs.DataResourceMode:
-		schema, exists := schema.DataSources[cfg.Type]
-		if !exists {
+		schema, _ := schema.SchemaForResourceType(mode, cfg.Type)
+		if schema == nil {
 			diags = diags.Append(&hcl.Diagnostic{
 				Severity: hcl.DiagError,
 				Summary:  "Invalid data source",
diff --git a/terraform/eval_validate_selfref.go b/terraform/eval_validate_selfref.go
index 2121f5dacd..edd604fa77 100644
--- a/terraform/eval_validate_selfref.go
+++ b/terraform/eval_validate_selfref.go
@@ -40,19 +40,9 @@ func (n *EvalValidateSelfRef) Eval(ctx EvalContext) (interface{}, error) {
 	var schema *configschema.Block
 	switch tAddr := addr.(type) {
 	case addrs.Resource:
-		switch tAddr.Mode {
-		case addrs.ManagedResourceMode:
-			schema = providerSchema.ResourceTypes[tAddr.Type]
-		case addrs.DataResourceMode:
-			schema = providerSchema.DataSources[tAddr.Type]
-		}
+		schema, _ = providerSchema.SchemaForResourceAddr(tAddr)
 	case addrs.ResourceInstance:
-		switch tAddr.Resource.Mode {
-		case addrs.ManagedResourceMode:
-			schema = providerSchema.ResourceTypes[tAddr.Resource.Type]
-		case addrs.DataResourceMode:
-			schema = providerSchema.DataSources[tAddr.Resource.Type]
-		}
+		schema, _ = providerSchema.SchemaForResourceAddr(tAddr.ContainingResource())
 	}
 
 	if schema == nil {
diff --git a/terraform/evaluate.go b/terraform/evaluate.go
index 918df2e875..37839865f2 100644
--- a/terraform/evaluate.go
+++ b/terraform/evaluate.go
@@ -798,17 +798,9 @@ func (d *evaluationStateData) getResourceInstancesAll(addr addrs.Resource, rng t
 
 func (d *evaluationStateData) getResourceSchema(addr addrs.Resource, providerAddr addrs.AbsProviderConfig) *configschema.Block {
 	providerType := providerAddr.ProviderConfig.Type
-	typeName := addr.Type
 	schemas := d.Evaluator.Schemas
-	switch addr.Mode {
-	case addrs.ManagedResourceMode:
-		return schemas.ResourceTypeConfig(providerType, typeName)
-	case addrs.DataResourceMode:
-		return schemas.DataSourceConfig(providerType, typeName)
-	default:
-		log.Printf("[WARN] Don't know how to fetch schema for resource %s", providerAddr)
-		return nil
-	}
+	schema, _ := schemas.ResourceTypeConfig(providerType, addr.Mode, addr.Type)
+	return schema
 }
 
 // coerceInstanceKey attempts to convert the given key to the type expected
diff --git a/terraform/evaluate_valid.go b/terraform/evaluate_valid.go
index 284cc548e2..7fc34a967c 100644
--- a/terraform/evaluate_valid.go
+++ b/terraform/evaluate_valid.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/hashicorp/terraform/addrs"
 	"github.com/hashicorp/terraform/configs"
-	"github.com/hashicorp/terraform/configs/configschema"
 	"github.com/hashicorp/terraform/tfdiags"
 )
 
@@ -112,13 +111,7 @@ func (d *evaluationStateData) staticValidateResourceReference(modCfg *configs.Co
 	// account provider inheritance, etc but it's okay here because we're only
 	// paying attention to the type anyway.
 	providerType := cfg.ProviderConfigAddr().Type
-	var schema *configschema.Block
-	switch addr.Mode {
-	case addrs.ManagedResourceMode:
-		schema = d.Evaluator.Schemas.ResourceTypeConfig(providerType, addr.Type)
-	case addrs.DataResourceMode:
-		schema = d.Evaluator.Schemas.DataSourceConfig(providerType, addr.Type)
-	}
+	schema, _ := d.Evaluator.Schemas.ResourceTypeConfig(providerType, addr.Mode, addr.Type)
 
 	if schema == nil {
 		// Prior validation should've taken care of a resource block with an
diff --git a/terraform/node_resource_abstract.go b/terraform/node_resource_abstract.go
index 4284a9b596..613e1b4c8b 100644
--- a/terraform/node_resource_abstract.go
+++ b/terraform/node_resource_abstract.go
@@ -47,8 +47,9 @@ type NodeAbstractResource struct {
 	// interfaces if you're running those transforms, but also be explicitly
 	// set if you already have that information.
 
-	Schema *configschema.Block // Schema for processing the configuration body
-	Config *configs.Resource   // Config is the resource in the config
+	Schema        *configschema.Block // Schema for processing the configuration body
+	SchemaVersion uint64              // Schema version of "Schema", as decided by the provider
+	Config        *configs.Resource   // Config is the resource in the config
 
 	ProvisionerSchemas map[string]*configschema.Block
 
@@ -423,8 +424,9 @@ func (n *NodeAbstractResource) AttachResourceConfig(c *configs.Resource) {
 }
 
 // GraphNodeAttachResourceSchema impl
-func (n *NodeAbstractResource) AttachResourceSchema(schema *configschema.Block) {
+func (n *NodeAbstractResource) AttachResourceSchema(schema *configschema.Block, version uint64) {
 	n.Schema = schema
+	n.SchemaVersion = version
 }
 
 // GraphNodeDotter impl.
diff --git a/terraform/schemas.go b/terraform/schemas.go
index 5475fe74b5..21d803b538 100644
--- a/terraform/schemas.go
+++ b/terraform/schemas.go
@@ -50,31 +50,12 @@ func (ss *Schemas) ProviderConfig(typeName string) *configschema.Block {
 // a resource using the "provider" meta-argument. Therefore it's important to
 // always pass the correct provider name, even though it many cases it feels
 // redundant.
-func (ss *Schemas) ResourceTypeConfig(providerType string, resourceType string) *configschema.Block {
+func (ss *Schemas) ResourceTypeConfig(providerType string, resourceMode addrs.ResourceMode, resourceType string) (block *configschema.Block, schemaVersion uint64) {
 	ps := ss.ProviderSchema(providerType)
 	if ps == nil || ps.ResourceTypes == nil {
-		return nil
+		return nil, 0
 	}
-
-	return ps.ResourceTypes[resourceType]
-}
-
-// DataSourceConfig returns the schema for the configuration of a given
-// data source belonging to a given provider type, or nil of no such
-// schema is available.
-//
-// In many cases the provider type is inferrable from the data source name,
-// but this is not always true because users can override the provider for
-// a resource using the "provider" meta-argument. Therefore it's important to
-// always pass the correct provider name, even though it many cases it feels
-// redundant.
-func (ss *Schemas) DataSourceConfig(providerType string, dataSource string) *configschema.Block {
-	ps := ss.ProviderSchema(providerType)
-	if ps == nil || ps.DataSources == nil {
-		return nil
-	}
-
-	return ps.DataSources[dataSource]
+	return ps.SchemaForResourceType(resourceMode, resourceType)
 }
 
 // ProvisionerConfig returns the schema for the configuration of a given
@@ -146,10 +127,13 @@ func loadProviderSchemas(schemas map[string]*ProviderSchema, config *configs.Con
 			Provider:      resp.Provider.Block,
 			ResourceTypes: make(map[string]*configschema.Block),
 			DataSources:   make(map[string]*configschema.Block),
+
+			ResourceTypeSchemaVersions: make(map[string]uint64),
 		}
 
 		for t, r := range resp.ResourceTypes {
 			s.ResourceTypes[t] = r.Block
+			s.ResourceTypeSchemaVersions[t] = r.Version
 		}
 
 		for t, d := range resp.DataSources {
@@ -241,22 +225,32 @@ type ProviderSchema struct {
 	Provider      *configschema.Block
 	ResourceTypes map[string]*configschema.Block
 	DataSources   map[string]*configschema.Block
+
+	ResourceTypeSchemaVersions map[string]uint64
+}
+
+// SchemaForResourceType attempts to find a schema for the given mode and type.
+// Returns nil if no such schema is available.
+func (ps *ProviderSchema) SchemaForResourceType(mode addrs.ResourceMode, typeName string) (schema *configschema.Block, version uint64) {
+	var m map[string]providers.Schema
+	switch mode {
+	case addrs.ManagedResourceMode:
+		return ps.ResourceTypes[typeName], ps.ResourceTypeSchemaVersions[typeName]
+	case addrs.DataResourceMode:
+		// Data resources don't have schema versions right now, since state is discarded for each refresh
+		return ps.DataSources[typeName], 0
+	default:
+		// Shouldn't happen, because the above cases are comprehensive.
+		return nil, 0
+	}
+	s := m[typeName]
+	return s.Block, s.Version
 }
 
 // SchemaForResourceAddr attempts to find a schema for the mode and type from
 // the given resource address. Returns nil if no such schema is available.
-func (ps *ProviderSchema) SchemaForResourceAddr(addr addrs.Resource) *configschema.Block {
-	var m map[string]*configschema.Block
-	switch addr.Mode {
-	case addrs.ManagedResourceMode:
-		m = ps.ResourceTypes
-	case addrs.DataResourceMode:
-		m = ps.DataSources
-	default:
-		// Shouldn't happen, because the above cases are comprehensive.
-		return nil
-	}
-	return m[addr.Type]
+func (ps *ProviderSchema) SchemaForResourceAddr(addr addrs.Resource) (schema *configschema.Block, version uint64) {
+	return ps.SchemaForResourceType(addr.Mode, addr.Type)
 }
 
 // ProviderSchemaRequest is used to describe to a ResourceProvider which
diff --git a/terraform/transform_attach_schema.go b/terraform/transform_attach_schema.go
index 2eceb1255d..c7695dd4ee 100644
--- a/terraform/transform_attach_schema.go
+++ b/terraform/transform_attach_schema.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"log"
 
-	"github.com/hashicorp/terraform/addrs"
 	"github.com/hashicorp/terraform/configs/configschema"
 	"github.com/hashicorp/terraform/dag"
 )
@@ -15,7 +14,7 @@ type GraphNodeAttachResourceSchema interface {
 	GraphNodeResource
 	GraphNodeProviderConsumer
 
-	AttachResourceSchema(*configschema.Block)
+	AttachResourceSchema(schema *configschema.Block, version uint64)
 }
 
 // GraphNodeAttachProviderConfigSchema is an interface implemented by node types
@@ -62,19 +61,13 @@ func (t *AttachSchemaTransformer) Transform(g *Graph) error {
 			providerAddr, _ := tv.ProvidedBy()
 			providerType := providerAddr.ProviderConfig.Type
 
-			var schema *configschema.Block
-			switch mode {
-			case addrs.ManagedResourceMode:
-				schema = t.Schemas.ResourceTypeConfig(providerType, typeName)
-			case addrs.DataResourceMode:
-				schema = t.Schemas.DataSourceConfig(providerType, typeName)
-			}
+			schema, version := t.Schemas.ResourceTypeConfig(providerType, mode, typeName)
 			if schema == nil {
 				log.Printf("[ERROR] AttachSchemaTransformer: No resource schema available for %s", addr)
 				continue
 			}
 			log.Printf("[TRACE] AttachSchemaTransformer: attaching resource schema to %s", dag.VertexName(v))
-			tv.AttachResourceSchema(schema)
+			tv.AttachResourceSchema(schema, version)
 		}
 
 		if tv, ok := v.(GraphNodeAttachProviderConfigSchema); ok {