From dc668dff386b70a8c151ded6cd1af5ff32867956 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Tue, 8 Mar 2022 13:06:30 -0500
Subject: [PATCH] ensure UI hooks are called for data sources

The UI hooks for data source reads were missed during planning. Move the
hook calls to immediatley before and after the ReadDataSource calls to
ensure they are called during both plan and apply.
---
 internal/terraform/context_apply_test.go      |  9 ++++++++
 .../node_resource_abstract_instance.go        | 22 +++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/internal/terraform/context_apply_test.go b/internal/terraform/context_apply_test.go
index 46dcbd58b8..5963b22426 100644
--- a/internal/terraform/context_apply_test.go
+++ b/internal/terraform/context_apply_test.go
@@ -1428,7 +1428,9 @@ func TestContext2Apply_dataBasic(t *testing.T) {
 		}),
 	}
 
+	hook := new(MockHook)
 	ctx := testContext2(t, &ContextOpts{
+		Hooks: []Hook{hook},
 		Providers: map[addrs.Provider]providers.Factory{
 			addrs.NewDefaultProvider("null"): testProviderFuncFixed(p),
 		},
@@ -1449,6 +1451,13 @@ func TestContext2Apply_dataBasic(t *testing.T) {
 	if actual != expected {
 		t.Fatalf("wrong result\n\ngot:\n%s\n\nwant:\n%s", actual, expected)
 	}
+
+	if !hook.PreApplyCalled {
+		t.Fatal("PreApply not called for data source read")
+	}
+	if !hook.PostApplyCalled {
+		t.Fatal("PostApply not called for data source read")
+	}
 }
 
 func TestContext2Apply_destroyData(t *testing.T) {
diff --git a/internal/terraform/node_resource_abstract_instance.go b/internal/terraform/node_resource_abstract_instance.go
index fb431df53f..67298b65a2 100644
--- a/internal/terraform/node_resource_abstract_instance.go
+++ b/internal/terraform/node_resource_abstract_instance.go
@@ -1381,6 +1381,13 @@ func (n *NodeAbstractResourceInstance) readDataSource(ctx EvalContext, configVal
 	// to actually call the provider to read the data.
 	log.Printf("[TRACE] readDataSource: %s configuration is complete, so reading from provider", n.Addr)
 
+	diags = diags.Append(ctx.Hook(func(h Hook) (HookAction, error) {
+		return h.PreApply(n.Addr, states.CurrentGen, plans.Read, cty.NullVal(configVal.Type()), configVal)
+	}))
+	if diags.HasErrors() {
+		return newVal, diags
+	}
+
 	resp := provider.ReadDataSource(providers.ReadDataSourceRequest{
 		TypeName:     n.Addr.ContainingResource().Resource.Type,
 		Config:       configVal,
@@ -1445,6 +1452,10 @@ func (n *NodeAbstractResourceInstance) readDataSource(ctx EvalContext, configVal
 		newVal = newVal.MarkWithPaths(pvm)
 	}
 
+	diags = diags.Append(ctx.Hook(func(h Hook) (HookAction, error) {
+		return h.PostApply(n.Addr, states.CurrentGen, newVal, diags.Err())
+	}))
+
 	return newVal, diags
 }
 
@@ -1703,13 +1714,6 @@ func (n *NodeAbstractResourceInstance) applyDataSource(ctx EvalContext, planned
 		return nil, keyData, diags
 	}
 
-	diags = diags.Append(ctx.Hook(func(h Hook) (HookAction, error) {
-		return h.PreApply(n.Addr, states.CurrentGen, planned.Action, planned.Before, planned.After)
-	}))
-	if diags.HasErrors() {
-		return nil, keyData, diags
-	}
-
 	config := *n.Config
 	schema, _ := providerSchema.SchemaForResourceAddr(n.Addr.ContainingResource().Resource)
 	if schema == nil {
@@ -1751,10 +1755,6 @@ func (n *NodeAbstractResourceInstance) applyDataSource(ctx EvalContext, planned
 		Status: states.ObjectReady,
 	}
 
-	diags = diags.Append(ctx.Hook(func(h Hook) (HookAction, error) {
-		return h.PostApply(n.Addr, states.CurrentGen, newVal, diags.Err())
-	}))
-
 	return state, keyData, diags
 }