feat: loading in plan state data only loads once per module

module state is loaded into root context with index merging as
appropriate

Author: Emyrk

hclext/merge.go

@@ -24,3 +24,48 @@ func MergeObjects(a, b cty.Value) cty.Value {
 func isNotEmptyObject(val cty.Value) bool {
 	return !val.IsNull() && val.IsKnown() && val.Type().IsObjectType() && val.LengthInt() > 0
 }
+
+func MergeWithTupleElement(list cty.Value, idx int, val cty.Value) cty.Value {
+	if list.IsNull() ||
+		!list.Type().IsTupleType() ||
+		list.LengthInt() <= idx {
+		return InsertTupleElement(list, idx, val)
+	}
+
+	existingElement := list.Index(cty.NumberIntVal(int64(idx)))
+	merged := MergeObjects(existingElement, val)
+	return InsertTupleElement(list, idx, merged)
+}
+
+// InsertTupleElement inserts a value into a tuple at the specified index.
+// If the idx is outside the bounds of the list, it grows the tuple to
+// the new size, and fills in `cty.NilVal` for the missing elements.
+//
+// This function will not panic. If the list value is not a list, it will
+// be replaced with an empty list.
+func InsertTupleElement(list cty.Value, idx int, val cty.Value) cty.Value {
+	if list.IsNull() || !list.Type().IsTupleType() {
+		// better than a panic
+		list = cty.EmptyTupleVal
+	}
+
+	if idx < 0 {
+		// Nothing to do?
+		return list
+	}
+
+	newList := make([]cty.Value, max(idx+1, list.LengthInt()))
+	for i := 0; i < len(newList); i++ {
+		newList[i] = cty.NilVal // Always insert a nil by default
+
+		if i < list.LengthInt() { // keep the original
+			newList[i] = list.Index(cty.NumberIntVal(int64(i)))
+		}
+
+		if i == idx { // add the new value
+			newList[i] = val
+		}
+	}
+
+	return cty.TupleVal(newList)
+}

plan.go

@@ -15,6 +15,8 @@ import (
 	tfjson "github.com/hashicorp/terraform-json"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/gocty"
+
+	"github.com/coder/preview/hclext"
 )
 
 func PlanJSONHook(dfs fs.FS, input Input) (func(ctx *tfcontext.Context, blocks terraform.Blocks, inputVars map[string]cty.Value), error) {
@@ -31,29 +33,45 @@ func PlanJSONHook(dfs fs.FS, input Input) (func(ctx *tfcontext.Context, blocks t
 		return nil, fmt.Errorf("unable to parse plan JSON: %w", err)
 	}
 
-	var _ = plan
 	return func(ctx *tfcontext.Context, blocks terraform.Blocks, inputVars map[string]cty.Value) {
+		loaded := make(map[*tfjson.StateModule]bool)
+
 		// Do not recurse to child blocks.
 		// TODO: Only load into the single parent context for the module.
+		// And do not load context for a module more than once
 		for _, block := range blocks {
+			// TODO: Maybe switch to the 'configuration' block
 			planMod := priorPlanModule(plan, block)
 			if planMod == nil {
 				continue
 			}
 
-			// TODO: Nested blocks might have an issue here with the correct context.
-			// We want the "module context", which is the parent of the top level
-			// block. Maybe there is a way to discover what that is via some
-			// var set in the context.
-			err = loadResourcesToContext(block.Context().Parent(), planMod.Resources)
+			if loaded[planMod] {
+				// No need to load this module into state again
+				continue
+			}
+
+			rootCtx := block.Context()
+			for {
+				if rootCtx.Parent() != nil {
+					rootCtx = rootCtx.Parent()
+					continue
+				}
+				break
+			}
+
+			// Load state into the context
+			err := loadResourcesToContext(rootCtx, planMod.Resources)
 			if err != nil {
 				// TODO: Somehow handle this error
 				panic(fmt.Sprintf("unable to load resources to context: %v", err))
 			}
+			loaded[planMod] = true
 		}
 	}, nil
 }
 
+// priorPlanModule returns the state data of the module a given block is in.
 func priorPlanModule(plan *tfjson.Plan, block *terraform.Block) *tfjson.StateModule {
 	if !block.InModule() {
 		return plan.PriorState.Values.RootModule
@@ -85,6 +103,22 @@ func priorPlanModule(plan *tfjson.Plan, block *terraform.Block) *tfjson.StateMod
 	return current
 }
 
+func matchingBlock(block *terraform.Block, planMod *tfjson.StateModule) *tfjson.StateResource {
+	ref := block.Reference()
+	matchKey := keyMatcher(ref.RawKey())
+
+	for _, resource := range planMod.Resources {
+		if ref.BlockType().ShortName() == string(resource.Mode) &&
+			ref.TypeLabel() == resource.Type &&
+			ref.NameLabel() == resource.Name &&
+			matchKey(resource.Index) {
+
+			return resource
+		}
+	}
+	return nil
+}
+
 func loadResourcesToContext(ctx *tfcontext.Context, resources []*tfjson.StateResource) error {
 	for _, resource := range resources {
 		if resource.Mode != "data" {
@@ -96,12 +130,35 @@ func loadResourcesToContext(ctx *tfcontext.Context, resources []*tfjson.StateRes
 			continue
 		}
 
+		path := []string{string(resource.Mode), resource.Type, resource.Name}
+
+		// Always merge with any existing values
+		existing := ctx.Get(path...)
+
 		val, err := toCtyValue(resource.AttributeValues)
 		if err != nil {
 			return fmt.Errorf("unable to determine value of resource %q: %w", resource.Address, err)
 		}
 
-		ctx.Set(val, string(resource.Mode), resource.Type, resource.Name)
+		var merged cty.Value
+		switch resource.Index.(type) {
+		case int, int32, int64, float32, float64:
+			asInt, ok := toInt(resource.Index)
+			if !ok {
+				return fmt.Errorf("unable to convert index '%v' to int", resource.Index)
+			}
+
+			if !existing.Type().IsTupleType() {
+				continue
+			}
+			merged = hclext.MergeWithTupleElement(existing, int(asInt), val)
+		case nil:
+			merged = hclext.MergeObjects(existing, val)
+		default:
+			return fmt.Errorf("unsupported index type %T", resource.Index)
+		}
+
+		ctx.Set(merged, string(resource.Mode), resource.Type, resource.Name)
 	}
 	return nil
 }
@@ -172,3 +229,51 @@ func TrivyParsePlanJSON(reader io.Reader) (*tfjson.Plan, error) {
 
 	return nil, err
 }
+
+func keyMatcher(key cty.Value) func(to any) bool {
+	switch {
+	case key.Type().Equals(cty.Number):
+		idx, _ := key.AsBigFloat().Int64()
+		return func(to any) bool {
+			asInt, ok := toInt(to)
+			return ok && asInt == idx
+		}
+
+	case key.Type().Equals(cty.String):
+		// TODO: handle key strings
+	}
+
+	return func(to any) bool {
+		return true
+	}
+}
+
+func toInt(to any) (int64, bool) {
+	switch typed := to.(type) {
+	case uint:
+		return int64(typed), true
+	case uint8:
+		return int64(typed), true
+	case uint16:
+		return int64(typed), true
+	case uint32:
+		return int64(typed), true
+	case uint64:
+		return int64(typed), true
+	case int:
+		return int64(typed), true
+	case int8:
+		return int64(typed), true
+	case int16:
+		return int64(typed), true
+	case int32:
+		return int64(typed), true
+	case int64:
+		return typed, true
+	case float32:
+		return int64(typed), true
+	case float64:
+		return int64(typed), true
+	}
+	return 0, false
+}

preview_test.go

@@ -156,11 +156,18 @@ func Test_Extract(t *testing.T) {
 			dir:         "module",
 			expTags:     map[string]string{},
 			unknownTags: []string{},
-			input:       preview.Input{},
+			input: preview.Input{
+				PlanJSONPath: "plan.json",
+				ParameterValues: map[string]string{
+					"extra": "foobar",
+				},
+			},
 			params: map[string]assertParam{
 				"jetbrains_ide": ap().
 					optVals("CL", "GO", "IU", "PY", "WS").
 					value("GO"),
+				"extra": ap().
+					value("foobar"),
 			},
 		},
 		{
@@ -285,16 +292,28 @@ func Test_Extract(t *testing.T) {
 			dir:  "demo_flat",
 			expTags: map[string]string{
 				"cluster": "confidential",
+				"hash":    "52bb4d943694f2f5867a251780f85e5a68906787b4ffa3157e29b9ef510b1a97",
 			},
 			input: preview.Input{
-				PlanJSONPath:    "",
-				ParameterValues: map[string]string{},
+				PlanJSONPath: "plan.json",
+				ParameterValues: map[string]string{
+					"hash": "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9",
+				},
 				Owner: types.WorkspaceOwner{
 					Groups: []string{"admin"},
 				},
 			},
 			unknownTags: []string{},
-			params:      map[string]assertParam{},
+			params: map[string]assertParam{
+				"hash": ap().
+					value("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"),
+				"security_level": ap(),
+				"region":         ap(),
+				"cpu":            ap(),
+				"browser":        ap(),
+				"team":           ap().optVals("frontend", "backend", "fullstack"),
+				"jetbrains_ide":  ap(),
+			},
 		},
 		{
 			name:    "count",
@@ -367,8 +386,6 @@ func Test_Extract(t *testing.T) {
 			}
 
 			dirFs := os.DirFS(filepath.Join("testdata", tc.dir))
-			//a, b := fs.ReadDir(dirFs, ".")
-			//fmt.Println(a, b)
 
 			output, diags := preview.Preview(context.Background(), tc.input, dirFs)
 			if tc.failPreview {

testdata/demo_flat/jetbrains_ide/main.tf

@@ -253,14 +253,15 @@ locals {
     }
   }
 
-  icon          = local.jetbrains_ides[data.coder_parameter.jetbrains_ide.value].icon
-  json_data     = var.latest ? jsondecode(data.http.jetbrains_ide_versions[data.coder_parameter.jetbrains_ide.value].response_body) : {}
+  first_val = jsondecode(data.coder_parameter.jetbrains_ide.value)[0]
+  icon          = local.jetbrains_ides[local.first_val].icon
+  json_data     = var.latest ? jsondecode(data.http.jetbrains_ide_versions[local.first_val].response_body) : {}
   key           = var.latest ? keys(local.json_data)[0] : ""
-  display_name  = local.jetbrains_ides[data.coder_parameter.jetbrains_ide.value].name
-  identifier    = data.coder_parameter.jetbrains_ide.value
-  download_link = var.latest ? local.json_data[local.key][0].downloads.linux.link : local.jetbrains_ides[data.coder_parameter.jetbrains_ide.value].download_link
-  build_number  = var.latest ? local.json_data[local.key][0].build : local.jetbrains_ides[data.coder_parameter.jetbrains_ide.value].build_number
-  version       = var.latest ? local.json_data[local.key][0].version : var.jetbrains_ide_versions[data.coder_parameter.jetbrains_ide.value].version
+  display_name  = local.jetbrains_ides[local.first_val].name
+  identifier    = local.first_val
+  download_link = var.latest ? local.json_data[local.key][0].downloads.linux.link : local.jetbrains_ides[local.first_val].download_link
+  build_number  = var.latest ? local.json_data[local.key][0].build : local.jetbrains_ides[local.first_val].build_number
+  version       = var.latest ? local.json_data[local.key][0].version : var.jetbrains_ide_versions[local.first_val].version
 }
 
 data "coder_parameter" "jetbrains_ide" {

testdata/demo_flat/main.tf

@@ -22,7 +22,6 @@ data coder_workspace_owner "me" {}
 module "jetbrains_gateway" {
   count          = 1
   source         = "./jetbrains_ide"
-  version        = "1.0.28"
   agent_id       = "random"
   folder         = "/home/coder/example"
   jetbrains_ides = local.teams[data.coder_parameter.team.value].codes

testdata/demo_flat/parameters.tf

@@ -1,5 +1,6 @@
 data "coder_parameter" "team" {
-  name        = "Team"
+  name        = "team"
+  display_name = "Team"
   description = "Which team are you on?"
   type        = "string"
   default     = "fullstack"
@@ -120,7 +121,7 @@ data "docker_registry_image" "coder" {
 }
 
 data "coder_parameter" "region" {
-  name         = "Region"
+  name         = "region"
   display_name = "Region"
   description  = "What region are you in?"
   form_type    = "dropdown"

testdata/module/main.tf

@@ -3,9 +3,14 @@ terraform {
     coder = {
       source = "coder/coder"
     }
+    docker = {
+      source  = "kreuzwerker/docker"
+      version = "3.0.2"
+    }
   }
 }
 
+
 module "jetbrains_gateway" {
   count          = 1
   source         = "registry.coder.com/modules/jetbrains-gateway/coder"
@@ -17,8 +22,22 @@ module "jetbrains_gateway" {
   default        = "GO"
 }
 
+data "coder_parameter" "extra" {
+  count = 1
+  name = "extra"
+  display_name = "Extra Param"
+  description = "A param to throw into the mix."
+  type        = "string"
+  default     = trimprefix(data.docker_registry_image.coder[1].sha256_digest, "sha256:")
+}
+
 data "coder_workspace" "me" {}
 resource "coder_agent" "main" {
   arch = "amd64"
   os   = "linux"
 }
+
+data "docker_registry_image" "coder" {
+  count = 2
+  name = count.index == 0 ? "ghcr.io/coder/coder:latest": "ghcr.io/coder/coder:v2.20.1"
+}