command/format: fix an issue where data resources were not displaying (#20386)

Fixes #20245

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())

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"
 }