Repository: lyraproj/hiera
Pull request #87 (2020-11-29 18:16:14)

Support looking up multiple keys
Adds a new option `--multi-lookup` which will look up all of the keys and output the results as a map. This is drastically faster when using expensive plugins like Terraform remote state and you have a lot of keys to look up.


Patch:
diff --git a/cli/hiera.go b/cli/hiera.go
index 49e9d4a..6c463bb 100644
--- a/cli/hiera.go
+++ b/cli/hiera.go
@@ -96,7 +96,7 @@ func NewCommand() *cobra.Command {
 	flags.Var(&dflt, `default`,
 		`a value to return if Hiera can't find a value in data`)
 	flags.StringVar(&cmdOpts.Type, `type`, ``,
-		`assert that the value has the specified type`)
+		`assert that the value has the specified type (if using --all this must be a map)`)
 	flags.StringVar(&dialect, `dialect`, `pcore`,
 		`dialect to use for rich data serialization and parsing of types pcore|dgo'`)
 	flags.StringVar(&cmdOpts.RenderAs, `render-as`, ``,
@@ -111,6 +111,8 @@ func NewCommand() *cobra.Command {
 		`a key:value or key=value where value is literal expressed using Puppet DSL`)
 	flags.StringArrayVar(&cmdOpts.FactPaths, `facts`, nil,
 		`like --vars but will also make variables available under the "facts" (for compatibility with Puppet's ruby version of Hiera)`)
+	flags.BoolVar(&cmdOpts.LookupAll, `all`, false,
+		`lookup all of the keys and output the results as a map`)
 
 	cmd.SetHelpTemplate(helpTemplate)
 	return cmd
diff --git a/hiera/hiera.go b/hiera/hiera.go
index ef13222..d67129d 100644
--- a/hiera/hiera.go
+++ b/hiera/hiera.go
@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"github.com/lyraproj/dgo/dgo"
+	"github.com/lyraproj/dgo/streamer"
 	"github.com/lyraproj/dgo/typ"
 	"github.com/lyraproj/dgo/util"
 	"github.com/lyraproj/dgo/vf"
@@ -51,6 +52,8 @@ type CommandOptions struct {
 
 	// ExplainOptions should be set to true to explain how lookup options were found for the lookup
 	ExplainOptions bool
+
+	LookupAll bool
 }
 
 // Lookup performs a lookup using the given parameters.
@@ -112,6 +115,45 @@ func Lookup2(
 	return nil
 }
 
+// LookupAll performs a lookup using the given parameters for all of the names passed in.
+//
+// ic - The lookup invocation
+//
+// names[] - The name or names to lookup
+//
+// valueType - Optional expected type of the found value
+//
+// override - Optional map to use as override. Values found here are returned immediately (no merge)
+//
+// defaultValuesHash - Optional map to use as the last resort
+//
+// options - Optional map with merge strategy and options
+func LookupAll(
+	ic api.Invocation,
+	names []string,
+	valueType dgo.StructMapType,
+	override dgo.Map,
+	defaultValuesHash dgo.Map,
+	options dgo.Map) dgo.Value {
+	response := vf.MutableMap()
+	for _, name := range names {
+		a := []string{name}
+		if v := lookupInMap(a, override); v != nil {
+			response.Put(name, ensureTypeFromMap(valueType, name, v))
+			continue
+		}
+		if v := ic.Lookup(api.NewKey(name), options); v != nil {
+			response.Put(name, ensureTypeFromMap(valueType, name, v))
+			continue
+		}
+		if v := lookupInMap(a, defaultValuesHash); v != nil {
+			response.Put(name, ensureTypeFromMap(valueType, name, v))
+			continue
+		}
+	}
+	return response
+}
+
 func lookupInMap(names []string, m dgo.Map) dgo.Value {
 	if m != nil && m.Len() > 0 {
 		for _, name := range names {
@@ -123,6 +165,16 @@ func lookupInMap(names []string, m dgo.Map) dgo.Value {
 	return nil
 }
 
+func ensureTypeFromMap(t dgo.StructMapType, k string, v dgo.Value) dgo.Value {
+	if t == nil {
+		return v
+	}
+	if e := t.Get(k); e != nil {
+		return ensureType(e.Value().(dgo.Type), v)
+	}
+	panic(fmt.Errorf("key '%s' was not found in the type map", k))
+}
+
 func ensureType(t dgo.Type, v dgo.Value) dgo.Value {
 	if t == nil || t.Instance(v) {
 		return v
@@ -159,11 +211,7 @@ var needParsePrefix = []string{`{`, `[`, `"`, `'`}
 // LookupAndRender performs a lookup using the given command options and arguments and renders the result on the given
 // io.Writer in accordance with the `RenderAs` option.
 func LookupAndRender(c api.Session, opts *CommandOptions, args []string, out io.Writer) bool {
-	tp := typ.Any
-	dl := c.Dialect()
-	if opts.Type != `` {
-		tp = dl.ParseType(nil, vf.String(opts.Type))
-	}
+	tp := parseType(opts.Type, c.Dialect())
 
 	var options dgo.Map
 	if !(opts.Merge == `` || opts.Merge == `first`) {
@@ -185,7 +233,17 @@ func LookupAndRender(c api.Session, opts *CommandOptions, args []string, out io.
 		explainer = explain.NewExplainer(opts.ExplainOptions, opts.ExplainOptions && !opts.ExplainData)
 	}
 
-	found := Lookup2(c.Invocation(createScope(c, opts), explainer), args, tp, dv, nil, nil, options, nil)
+	var found dgo.Value
+	invocation := c.Invocation(createScope(c, opts), explainer)
+	if opts.LookupAll {
+		stp, ok := tp.(dgo.StructMapType)
+		if !ok && opts.Type != `` {
+			panic(fmt.Errorf("type must be a map"))
+		}
+		found = LookupAll(invocation, args, stp, nil, nil, options)
+	} else {
+		found = Lookup2(invocation, args, tp, dv, nil, nil, options, nil)
+	}
 	if explainer != nil {
 		renderAs := Text
 		if opts.RenderAs != `` {
@@ -207,6 +265,14 @@ func LookupAndRender(c api.Session, opts *CommandOptions, args []string, out io.
 	return true
 }
 
+func parseType(t string, dl streamer.Dialect) dgo.Type {
+	tp := typ.Any
+	if t != `` {
+		tp = dl.ParseType(nil, vf.String(t))
+	}
+	return tp
+}
+
 func parseCommandLineValue(c api.Session, vs string) dgo.Value {
 	vs = strings.TrimSpace(vs)
 	for _, pfx := range needParsePrefix {


Test patch:
diff --git a/lookup/lookup_test.go b/lookup/lookup_test.go
index 6235e65..6d87286 100644
--- a/lookup/lookup_test.go
+++ b/lookup/lookup_test.go
@@ -435,6 +435,62 @@ func TestLookup_issue75(t *testing.T) {
 	}
 }
 
+func TestLookup_all_json(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		result, err := cli.ExecuteLookup(`hash`, `array`, `--render-as`, `json`, `--all`)
+		require.NoError(t, err)
+		require.Equal(t, `{"hash":{"one":1,"two":"two","three":{"a":"A","c":"C"}},"array":["one","two","three"]}
+`, string(result))
+	})
+}
+
+func TestLookup_all_simple(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		result, err := cli.ExecuteLookup(`simple`, `--render-as`, `s`, `--all`)
+		require.NoError(t, err)
+		require.Equal(t, `{"simple":"value"}
+`, string(result))
+	})
+}
+
+func TestLookup_all_not_there(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		result, err := cli.ExecuteLookup(`simple`, `not_there`, `--render-as`, `s`, `--all`)
+		require.NoError(t, err)
+		require.Equal(t, `{"simple":"value"}
+`, string(result))
+	})
+}
+
+func TestLookup_all_type(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		result, err := cli.ExecuteLookup(`stringkey`, `intkey`, `--all`, `--dialect`, `dgo`, `--render-as`, `s`, `--type`, `{"stringkey":string,"intkey":int}`)
+		require.NoError(t, err)
+		require.Equal(t, `{"stringkey":"stringvalue","intkey":1}
+`, string(result))
+	})
+}
+
+func TestLookup_all_invalid_type(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		_, err := cli.ExecuteLookup(`stringkey`, `intkey`, `--all`, `--dialect`, `dgo`, `--render-as`, `s`, `--type`, `{"stringkey":int,"intkey":int}`)
+		require.Error(t, err, `the value 'stringvalue' cannot be converted to an int`)
+	})
+}
+
+func TestLookup_all_invalid_type_map(t *testing.T) {
+	ensureTestPlugin(t)
+	inTestdata(func() {
+		_, err := cli.ExecuteLookup(`stringkey`, `intkey`, `--all`, `--dialect`, `dgo`, `--render-as`, `s`, `--type`, `string`)
+		require.Error(t, err, `type must be a map`)
+	})
+}
+
 /*
 func TestDataHash_refuseToDie(t *testing.T) {
 	ensureTestPlugin(t)
diff --git a/lookup/testdata/hiera/common.yaml b/lookup/testdata/hiera/common.yaml
index dc65b3f..d7f67da 100644
--- a/lookup/testdata/hiera/common.yaml
+++ b/lookup/testdata/hiera/common.yaml
@@ -34,3 +34,9 @@ lookup_options:
     merge: deep
   sense:
     convert_to: Sensitive
+
+simple: value
+
+stringkey: stringvalue
+
+intkey: 1