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.

Author: jbardin

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

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
 }