refactoring

cmd/tfstate-lookup/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
@@ -34,12 +33,9 @@ func main() {
 func _main() error {
 	var (
 		stateFile string
-		raw       bool
 	)
 	flag.StringVar(&stateFile, "state", "terraform.tfstate", "tfstate file path")
 	flag.StringVar(&stateFile, "s", "terraform.tfstate", "tfstate file path")
-	flag.BoolVar(&raw, "raw", false, "raw output")
-	flag.BoolVar(&raw, "r", false, "raw output")
 	flag.Parse()
 	if len(flag.Args()) == 0 {
 		flag.Usage()
@@ -50,6 +46,7 @@ func _main() error {
 	if err != nil {
 		return err
 	}
+	defer f.Close()
 
 	s, err := tfstate.Read(f)
 	if err != nil {
@@ -59,17 +56,6 @@ func _main() error {
 	if err != nil {
 		return err
 	}
-	if raw {
-		switch res.(type) {
-		case string, float64:
-			fmt.Fprintln(os.Stdout, res)
-		default:
-			json.NewEncoder(os.Stdout).Encode(res)
-		}
-	} else {
-		enc := json.NewEncoder(os.Stdout)
-		enc.SetIndent("", "  ")
-		enc.Encode(res)
-	}
+	fmt.Println(res.String())
 	return nil
 }

tfstate/lookup.go

@@ -14,6 +14,40 @@ import (
 	"github.com/pkg/errors"
 )
 
+// Attribute represents tfstate resource attributes
+type Attribute struct {
+	Value interface{}
+}
+
+func (a Attribute) String() string {
+	switch v := a.Value; v.(type) {
+	case string, float64:
+		return fmt.Sprint(v)
+	default:
+		b, _ := json.Marshal(v)
+		return string(b)
+	}
+}
+
+func (a *Attribute) Query(query string) (*Attribute, error) {
+	jq, err := gojq.Parse(query)
+	if err != nil {
+		return nil, err
+	}
+	iter := jq.Run(a.Value)
+	for {
+		v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if err, ok := v.(error); ok {
+			return nil, err
+		}
+		return &Attribute{Value: v}, nil
+	}
+	return nil, fmt.Errorf("%s is not found in attributes", query)
+}
+
 // TFState represents a tfstate
 type TFState struct {
 	statefile *statefile.File
@@ -31,21 +65,16 @@ func Read(src io.Reader) (*TFState, error) {
 }
 
 // Lookup lookups attributes of the specified key in tfstate
-func (s *TFState) Lookup(key string) (interface{}, error) {
-	b, query, err := lookupAttrs(s.statefile, key)
+func (s *TFState) Lookup(key string) (*Attribute, error) {
+	attr, query, err := lookupAttrs(s.statefile, key)
 	if err != nil {
 		return nil, err
 	}
-	log.Println("[debug] attrs", string(b))
-
-	var obj interface{}
-	if err := json.Unmarshal(b, &obj); err != nil {
-		return nil, err
-	}
-	return queryObj(obj, query)
+	log.Println("[debug] attrs", attr.String())
+	return attr.Query(query)
 }
 
-func lookupAttrs(file *statefile.File, key string) ([]byte, string, error) {
+func lookupAttrs(file *statefile.File, key string) (*Attribute, string, error) {
 	name := key
 	var module *states.Module
 	nameParts := strings.Split(name, ".")
@@ -105,24 +134,10 @@ func lookupAttrs(file *statefile.File, key string) ([]byte, string, error) {
 	if instance == nil || instance.Current == nil {
 		return nil, query, fmt.Errorf("%s is not found in state file", key)
 	}
-	return instance.Current.AttrsJSON, query, nil
-}
 
-func queryObj(obj interface{}, query string) (interface{}, error) {
-	jq, err := gojq.Parse(query)
-	if err != nil {
-		return nil, err
+	var value interface{}
+	if err := json.Unmarshal(instance.Current.AttrsJSON, &value); err != nil {
+		return nil, query, errors.Wrap(err, "invalid json attributes")
 	}
-	iter := jq.Run(obj)
-	for {
-		v, ok := iter.Next()
-		if !ok {
-			break
-		}
-		if err, ok := v.(error); ok {
-			return nil, err
-		}
-		return v, nil
-	}
-	return nil, fmt.Errorf("%s is not found in attributes", query)
+	return &Attribute{Value: value}, query, nil
 }

tfstate/lookup_test.go

@@ -71,7 +71,7 @@ func TestLookupOK(t *testing.T) {
 		if err != nil {
 			t.Error(err)
 		}
-		if diff := cmp.Diff(res, ts.Result); diff != "" {
+		if diff := cmp.Diff(res.Value, ts.Result); diff != "" {
 			t.Errorf("%s unexpected result %s", ts.Key, diff)
 		}
 	}