Ensure nested data blocks execute last of all Terraform resources (#33301)

* Ensure nested data blocks execute last of all Terraform resources

* add test, execute only during apply

* address comments

internal/terraform/graph_builder_apply.go

@@ -148,6 +148,10 @@ func (b *ApplyGraphBuilder) Steps() []GraphTransformer {
 		&ReferenceTransformer{},
 		&AttachDependenciesTransformer{},
 
+		// Nested data blocks should be loaded after every other resource has
+		// done its thing.
+		&checkStartTransformer{Config: b.Config, Operation: b.Operation},
+
 		// Detect when create_before_destroy must be forced on for a particular
 		// node due to dependency edges, to avoid graph cycles during apply.
 		&ForcedCBDTransformer{},

internal/terraform/graph_builder_apply_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/plans"
+	"github.com/hashicorp/terraform/internal/providers"
 	"github.com/hashicorp/terraform/internal/states"
 )
 
@@ -702,6 +703,91 @@ func TestApplyGraphBuilder_orphanedWithProvider(t *testing.T) {
 	testGraphNotContains(t, g, "provider.test")
 }
 
+func TestApplyGraphBuilder_withChecks(t *testing.T) {
+	awsProvider := mockProviderWithResourceTypeSchema("aws_instance", simpleTestSchema())
+
+	changes := &plans.Changes{
+		Resources: []*plans.ResourceInstanceChangeSrc{
+			{
+				Addr: mustResourceInstanceAddr("aws_instance.foo"),
+				ChangeSrc: plans.ChangeSrc{
+					Action: plans.Create,
+				},
+			},
+			{
+				Addr: mustResourceInstanceAddr("aws_instance.baz"),
+				ChangeSrc: plans.ChangeSrc{
+					Action: plans.Create,
+				},
+			},
+			{
+				Addr: mustResourceInstanceAddr("data.aws_data_source.bar"),
+				ChangeSrc: plans.ChangeSrc{
+					Action: plans.Read,
+				},
+				ActionReason: plans.ResourceInstanceReadBecauseCheckNested,
+			},
+		},
+	}
+
+	plugins := newContextPlugins(map[addrs.Provider]providers.Factory{
+		addrs.NewDefaultProvider("aws"): providers.FactoryFixed(awsProvider),
+	}, nil)
+
+	b := &ApplyGraphBuilder{
+		Config:    testModule(t, "apply-with-checks"),
+		Changes:   changes,
+		Plugins:   plugins,
+		State:     states.NewState(),
+		Operation: walkApply,
+	}
+
+	g, err := b.Build(addrs.RootModuleInstance)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if g.Path.String() != addrs.RootModuleInstance.String() {
+		t.Fatalf("wrong path %q", g.Path.String())
+	}
+
+	got := strings.TrimSpace(g.String())
+	// We're especially looking for the edge here, where aws_instance.bat
+	// has a dependency on aws_instance.boo
+	want := strings.TrimSpace(testPlanWithCheckGraphBuilderStr)
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Fatalf("\ngot:\n%s\n\nwant:\n%s\n\ndiff:\n%s", got, want, diff)
+	}
+
+}
+
+const testPlanWithCheckGraphBuilderStr = `
+(execute checks)
+  aws_instance.baz
+aws_instance.baz
+  aws_instance.baz (expand)
+  aws_instance.foo
+aws_instance.baz (expand)
+  provider["registry.terraform.io/hashicorp/aws"]
+aws_instance.foo
+  aws_instance.foo (expand)
+aws_instance.foo (expand)
+  provider["registry.terraform.io/hashicorp/aws"]
+check.my_check (expand)
+  data.aws_data_source.bar
+data.aws_data_source.bar
+  (execute checks)
+  data.aws_data_source.bar (expand)
+data.aws_data_source.bar (expand)
+  provider["registry.terraform.io/hashicorp/aws"]
+provider["registry.terraform.io/hashicorp/aws"]
+provider["registry.terraform.io/hashicorp/aws"] (close)
+  data.aws_data_source.bar
+root
+  check.my_check (expand)
+  provider["registry.terraform.io/hashicorp/aws"] (close)
+`
+
 const testApplyGraphBuilderStr = `
 module.child (close)
   module.child.test_object.other

internal/terraform/node_check.go

@@ -183,3 +183,22 @@ func (n *nodeCheckAssert) Execute(ctx EvalContext, _ walkOperation) tfdiags.Diag
 func (n *nodeCheckAssert) Name() string {
 	return n.addr.String() + " (assertions)"
 }
+
+var (
+	_ GraphNodeExecutable = (*nodeCheckStart)(nil)
+)
+
+// We need to ensure that any nested data sources execute after all other
+// resource changes have been applied. This node acts as a single point of
+// dependency that can enforce this ordering.
+type nodeCheckStart struct{}
+
+func (n *nodeCheckStart) Execute(context EvalContext, operation walkOperation) tfdiags.Diagnostics {
+	// This node doesn't actually do anything, except simplify the underlying
+	// graph structure.
+	return nil
+}
+
+func (n *nodeCheckStart) Name() string {
+	return "(execute checks)"
+}

internal/terraform/testdata/apply-with-checks/main.tf

@@ -0,0 +1,20 @@
+
+resource "aws_instance" "foo" {
+  test_string = "Hello, world!"
+}
+
+resource "aws_instance" "baz" {
+  test_string = aws_instance.foo.test_string
+}
+
+check "my_check" {
+  data "aws_data_source" "bar" {
+    id = "UI098L"
+  }
+
+  assert {
+    condition = data.aws_data_source.bar.foo == "valid value"
+    error_message = "invalid value"
+  }
+
+}

internal/terraform/transform_check.go

@@ -93,7 +93,10 @@ func (t *checkTransformer) transform(g *Graph, cfg *configs.Config, allNodes []d
 							// Make sure we report our checks before we execute any
 							// embedded data resource.
 							g.Connect(dag.BasicEdge(other, report))
-							continue
+
+							// There's at most one embedded data source, and
+							// we've found it so stop looking.
+							break
 						}
 					}
 				}

internal/terraform/transform_check_starter.go

@@ -0,0 +1,126 @@
+package terraform
+
+import (
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/configs"
+	"github.com/hashicorp/terraform/internal/dag"
+)
+
+var _ GraphTransformer = (*checkStartTransformer)(nil)
+
+// checkStartTransformer checks if the configuration has any data blocks nested
+// within check blocks, and if it does then it introduces a nodeCheckStart
+// vertex that ensures all resources have been applied before it starts loading
+// the nested data sources.
+type checkStartTransformer struct {
+	// Config for the entire module.
+	Config *configs.Config
+
+	// Operation is the current operation this node will be part of.
+	Operation walkOperation
+}
+
+func (s *checkStartTransformer) Transform(graph *Graph) error {
+	if s.Operation != walkApply && s.Operation != walkPlan {
+		// We only actually execute the checks during plan apply operations
+		// so if we are doing something else we can just skip this and
+		// leave the graph alone.
+		return nil
+	}
+
+	var resources []dag.Vertex
+	var nested []dag.Vertex
+
+	// We're going to step through all the vertices and pull out the relevant
+	// resources and data sources.
+	for _, vertex := range graph.Vertices() {
+		if node, isResource := vertex.(GraphNodeCreator); isResource {
+			addr := node.CreateAddr()
+
+			if addr.Resource.Resource.Mode == addrs.ManagedResourceMode {
+				// This is a resource, so we want to make sure it executes
+				// before any nested data sources.
+
+				// We can reduce the number of additional edges we write into
+				// the graph by only including "leaf" resources, that is
+				// resources that aren't referenced by other resources. If a
+				// resource is referenced by another resource then we know that
+				// it will execute before that resource so we only need to worry
+				// about the referencing resource.
+
+				leafResource := true
+				for _, other := range graph.UpEdges(vertex) {
+					if otherResource, isResource := other.(GraphNodeCreator); isResource {
+						otherAddr := otherResource.CreateAddr()
+						if otherAddr.Resource.Resource.Mode == addrs.ManagedResourceMode {
+							// Then this resource is being referenced so skip
+							// it.
+							leafResource = false
+							break
+						}
+					}
+				}
+
+				if leafResource {
+					resources = append(resources, vertex)
+				}
+
+				// We've handled the resource so move to the next vertex.
+				continue
+			}
+
+			// Now, we know we are processing a data block.
+
+			config := s.Config
+			if !addr.Module.IsRoot() {
+				config = s.Config.Descendent(addr.Module.Module())
+			}
+			if config == nil {
+				// might have been deleted, so it won't be subject to any checks
+				// anyway.
+				continue
+			}
+
+			resource := config.Module.ResourceByAddr(addr.Resource.Resource)
+			if resource == nil {
+				// might have been deleted, so it won't be subject to any checks
+				// anyway.
+				continue
+			}
+
+			if _, ok := resource.Container.(*configs.Check); ok {
+				// Then this is a data source within a check block, so let's
+				// make a note of it.
+				nested = append(nested, vertex)
+			}
+
+			// Otherwise, it's just a normal data source. From a check block we
+			// don't really care when Terraform is loading non-nested data
+			// sources so we'll just forget about it and move on.
+		}
+	}
+
+	if len(nested) > 0 {
+
+		// We don't need to do any of this if we don't have any nested data
+		// sources, so we check that first.
+		//
+		// Otherwise we introduce a vertex that can act as a pauser between
+		// our nested data sources and leaf resources.
+
+		check := &nodeCheckStart{}
+		graph.Add(check)
+
+		// Finally, connect everything up so it all executes in order.
+
+		for _, vertex := range nested {
+			graph.Connect(dag.BasicEdge(vertex, check))
+		}
+
+		for _, vertex := range resources {
+			graph.Connect(dag.BasicEdge(check, vertex))
+		}
+	}
+
+	return nil
+}