Fix crash when nested data blocks are mixed with the import command

internal/terraform/context_import_test.go

@@ -8,11 +8,13 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/zclconf/go-cty/cty"
+
 	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/configs/configschema"
 	"github.com/hashicorp/terraform/internal/providers"
 	"github.com/hashicorp/terraform/internal/states"
-	"github.com/zclconf/go-cty/cty"
 )
 
 func TestContextImport_basic(t *testing.T) {
@@ -988,12 +990,62 @@ resource "test_resource" "test" {
 	}
 }
 
+func TestContextImport_33572(t *testing.T) {
+	p := testProvider("aws")
+	m := testModule(t, "issue-33572")
+
+	ctx := testContext2(t, &ContextOpts{
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("aws"): testProviderFuncFixed(p),
+		},
+	})
+
+	p.ImportResourceStateResponse = &providers.ImportResourceStateResponse{
+		ImportedResources: []providers.ImportedResource{
+			{
+				TypeName: "aws_instance",
+				State: cty.ObjectVal(map[string]cty.Value{
+					"id": cty.StringVal("foo"),
+				}),
+			},
+		},
+	}
+
+	state, diags := ctx.Import(m, states.NewState(), &ImportOpts{
+		Targets: []*ImportTarget{
+			{
+				Addr: addrs.RootModuleInstance.ResourceInstance(
+					addrs.ManagedResourceMode, "aws_instance", "foo", addrs.NoKey,
+				),
+				ID: "bar",
+			},
+		},
+	})
+	if diags.HasErrors() {
+		t.Fatalf("unexpected errors: %s", diags.Err())
+	}
+	actual := strings.TrimSpace(state.String())
+	expected := strings.TrimSpace(testImportStrWithDataSource)
+	if diff := cmp.Diff(actual, expected); len(diff) > 0 {
+		t.Fatalf("wrong final state\ngot:\n%s\nwant:\n%s\ndiff:\n%s", actual, expected, diff)
+	}
+}
+
 const testImportStr = `
 aws_instance.foo:
   ID = foo
   provider = provider["registry.terraform.io/hashicorp/aws"]
 `
 
+const testImportStrWithDataSource = `
+data.aws_data_source.bar:
+  ID = baz
+  provider = provider["registry.terraform.io/hashicorp/aws"]
+aws_instance.foo:
+  ID = foo
+  provider = provider["registry.terraform.io/hashicorp/aws"]
+`
+
 const testImportCountIndexStr = `
 aws_instance.foo.0:
   ID = foo

internal/terraform/testdata/issue-33572/main.tf

@@ -0,0 +1,14 @@
+provider "aws" {}
+
+resource "aws_instance" "foo" {}
+
+check "aws_instance_exists" {
+  data "aws_data_source" "bar" {
+    id = "baz"
+  }
+
+  assert {
+    condition     = data.aws_data_source.bar.foo == "Hello, world!"
+    error_message = "incorrect value"
+  }
+}

internal/terraform/transform_check.go

@@ -116,10 +116,14 @@ func (t *checkTransformer) transform(g *Graph, cfg *configs.Config, allNodes []d
 // ReportChecks returns true if this operation should report any check blocks
 // that it is about to execute.
 //
-// This is generally only true for planning operations, as apply operations
-// recreate the expected checks from the plan.
+// This is true for planning operations, as apply operations recreate the
+// expected checks from the plan.
+//
+// We'll also report the checks during an import operation. We still execute
+// our check blocks during an import operation so they need to be reported
+// first.
 func (t *checkTransformer) ReportChecks() bool {
-	return t.Operation == walkPlan
+	return t.Operation == walkPlan || t.Operation == walkImport
 }
 
 // ExecuteChecks returns true if this operation should actually execute any
@@ -129,7 +133,7 @@ func (t *checkTransformer) ReportChecks() bool {
 // graph, but they will only validate things like references and syntax.
 func (t *checkTransformer) ExecuteChecks() bool {
 	switch t.Operation {
-	case walkPlan, walkApply:
+	case walkPlan, walkApply, walkImport:
 		// We only actually execute the checks for plan and apply operations.
 		return true
 	default: