From 874b3339624b624559a0075549d2bb751038ea84 Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Tue, 19 Feb 2019 14:18:47 -0800
Subject: [PATCH] command/format: fix an issue where data resources were not
 displaying (#20386)

Fixes #20245
---
 command/format/state.go      | 16 ++++++++++++----
 command/format/state_test.go | 22 +++++++++++++++++++++-
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/command/format/state.go b/command/format/state.go
index 69ecb06f3..f411ef9c6 100644
--- a/command/format/state.go
+++ b/command/format/state.go
@@ -116,7 +116,12 @@ func formatStateModule(p blockBodyDiffPrinter, m *states.Module, schemas *terraf
 
 			switch addr.Mode {
 			case addrs.ManagedResourceMode:
-				if _, exists := schemas.Providers[provider].ResourceTypes[addr.Type]; !exists {
+				schema, _ = schemas.ResourceTypeConfig(
+					provider,
+					addr.Mode,
+					addr.Type,
+				)
+				if schema == nil {
 					p.buf.WriteString(fmt.Sprintf(
 						"# missing schema for provider %q resource type %s\n\n", provider, addr.Type))
 					continue
@@ -127,9 +132,13 @@ func formatStateModule(p blockBodyDiffPrinter, m *states.Module, schemas *terraf
 					addr.Type,
 					addr.Name,
 				))
-				schema = schemas.Providers[provider].ResourceTypes[addr.Type]
 			case addrs.DataResourceMode:
-				if _, exists := schemas.Providers[provider].ResourceTypes[addr.Type]; !exists {
+				schema, _ = schemas.ResourceTypeConfig(
+					provider,
+					addr.Mode,
+					addr.Type,
+				)
+				if schema == nil {
 					p.buf.WriteString(fmt.Sprintf(
 						"# missing schema for provider %q data source %s\n\n", provider, addr.Type))
 					continue
@@ -140,7 +149,6 @@ func formatStateModule(p blockBodyDiffPrinter, m *states.Module, schemas *terraf
 					addr.Type,
 					addr.Name,
 				))
-				schema = schemas.Providers[provider].DataSources[addr.Type]
 			default:
 				// should never happen, since the above is exhaustive
 				p.buf.WriteString(addr.String())
diff --git a/command/format/state_test.go b/command/format/state_test.go
index 9bdfb81b3..b39a15c62 100644
--- a/command/format/state_test.go
+++ b/command/format/state_test.go
@@ -42,6 +42,21 @@ func TestState(t *testing.T) {
 			Type: "test",
 		}.Absolute(addrs.RootModuleInstance),
 	)
+	rootModule.SetResourceInstanceCurrent(
+		addrs.Resource{
+			Mode: addrs.DataResourceMode,
+			Type: "test_data_source",
+			Name: "data",
+		}.Instance(addrs.NoKey),
+		&states.ResourceInstanceObjectSrc{
+			Status:        states.ObjectReady,
+			SchemaVersion: 1,
+			AttrsJSON:     []byte(`{"compute":"sure"}`),
+		},
+		addrs.ProviderConfig{
+			Type: "test",
+		}.Absolute(addrs.RootModuleInstance),
+	)
 
 	tests := []struct {
 		State *StateOpts
@@ -142,7 +157,12 @@ func testSchemas() *terraform.Schemas {
 	}
 }
 
-const TestOutput = `# test_resource.baz[0]: 
+const TestOutput = `# data.test_data_source.data: 
+data "test_data_source" "data" {
+    compute = "sure"
+}
+
+# test_resource.baz[0]: 
 resource "test_resource" "baz" {
     woozles = "confuzles"
 }