Add coll.JQ using gojq library (hairyhenderson#1585)

* feat: add coll.JQ using gojq library * fix: jq function naming (gojq -> jq) * test: add tests (take from jsonpath_test.go) * chore: add TODO for nil values (are they allowed?) * refactor: use fmt.Errorf instead of errors.Wrapf Co-authored-by: Dave Henderson <[email protected]> * fix: wrong alias for coll.JQ Co-authored-by: Dave Henderson <[email protected]> * docs: add links to JQ Co-authored-by: Dave Henderson <[email protected]> * test: add assertions after json marshal/unmarshal Co-authored-by: Dave Henderson <[email protected]> * refactor: use fmt.Errorf instead of errors.Wrapf Co-authored-by: Dave Henderson <[email protected]> * fix: test syntax and null handling * docs: improve documentation * docs: add blank line * Support cancellation Signed-off-by: Dave Henderson <[email protected]> * Support (almost) all types, not just map[string]interface{} and []interface{} Signed-off-by: Dave Henderson <[email protected]> * add an integration test for coll.JQ Signed-off-by: Dave Henderson <[email protected]> Signed-off-by: Dave Henderson <[email protected]> Co-authored-by: Andreas Hochsteger <[email protected]> Co-authored-by: Dave Henderson <[email protected]>
flanksource · Apr 1, 2023 · 56cbb95 · 56cbb95
1 parent 475dd7e
commit 56cbb95
Show file tree

Hide file tree

Showing 7 changed files with 453 additions and 0 deletions.
diff --git a/coll/jq.go b/coll/jq.go
@@ -0,0 +1,96 @@
+package coll
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"reflect"
+
+	"github.com/itchyny/gojq"
+)
+
+// JQ -
+func JQ(ctx context.Context, jqExpr string, in interface{}) (interface{}, error) {
+	query, err := gojq.Parse(jqExpr)
+	if err != nil {
+		return nil, fmt.Errorf("jq parsing expression %q: %w", jqExpr, err)
+	}
+
+	// convert input to a supported type, if necessary
+	in, err = jqConvertType(in)
+	if err != nil {
+		return nil, fmt.Errorf("jq type conversion: %w", err)
+	}
+
+	iter := query.RunWithContext(ctx, in)
+	var out interface{}
+	a := []interface{}{}
+	for {
+		v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if err, ok := v.(error); ok {
+			return nil, fmt.Errorf("jq execution: %w", err)
+		}
+		a = append(a, v)
+	}
+	if len(a) == 1 {
+		out = a[0]
+	} else {
+		out = a
+	}
+
+	return out, nil
+}
+
+// jqConvertType converts the input to a map[string]interface{}, []interface{},
+// or other supported primitive JSON types.
+func jqConvertType(in interface{}) (interface{}, error) {
+	// if it's already a supported type, pass it through
+	switch in.(type) {
+	case map[string]interface{}, []interface{},
+		string, []byte,
+		nil, bool,
+		int, int8, int16, int32, int64,
+		uint, uint8, uint16, uint32, uint64,
+		float32, float64:
+		return in, nil
+	}
+
+	inType := reflect.TypeOf(in)
+	value := reflect.ValueOf(in)
+
+	// pointers need to be dereferenced first
+	if inType.Kind() == reflect.Ptr {
+		inType = inType.Elem()
+		value = value.Elem()
+	}
+
+	mapType := reflect.TypeOf(map[string]interface{}{})
+	sliceType := reflect.TypeOf([]interface{}{})
+	// if it can be converted to a map or slice, do that
+	if inType.ConvertibleTo(mapType) {
+		return value.Convert(mapType).Interface(), nil
+	} else if inType.ConvertibleTo(sliceType) {
+		return value.Convert(sliceType).Interface(), nil
+	}
+
+	// if it's a struct, the simplest (though not necessarily most efficient)
+	// is to JSON marshal/unmarshal it
+	if inType.Kind() == reflect.Struct {
+		b, err := json.Marshal(in)
+		if err != nil {
+			return nil, fmt.Errorf("json marshal struct: %w", err)
+		}
+		var m map[string]interface{}
+		err = json.Unmarshal(b, &m)
+		if err != nil {
+			return nil, fmt.Errorf("json unmarshal struct: %w", err)
+		}
+		return m, nil
+	}
+
+	// we maybe don't need to convert the value, so return it as-is
+	return in, nil
+}
diff --git a/coll/jq_test.go b/coll/jq_test.go
@@ -0,0 +1,236 @@
+package coll
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestJQ(t *testing.T) {
+	ctx := context.Background()
+	in := map[string]interface{}{
+		"store": map[string]interface{}{
+			"book": []interface{}{
+				map[string]interface{}{
+					"category": "reference",
+					"author":   "Nigel Rees",
+					"title":    "Sayings of the Century",
+					"price":    8.95,
+				},
+				map[string]interface{}{
+					"category": "fiction",
+					"author":   "Evelyn Waugh",
+					"title":    "Sword of Honour",
+					"price":    12.99,
+				},
+				map[string]interface{}{
+					"category": "fiction",
+					"author":   "Herman Melville",
+					"title":    "Moby Dick",
+					"isbn":     "0-553-21311-3",
+					"price":    8.99,
+				},
+				map[string]interface{}{
+					"category": "fiction",
+					"author":   "J. R. R. Tolkien",
+					"title":    "The Lord of the Rings",
+					"isbn":     "0-395-19395-8",
+					"price":    22.99,
+				},
+			},
+			"bicycle": map[string]interface{}{
+				"color": "red",
+				"price": 19.95,
+			},
+		},
+	}
+	out, err := JQ(ctx, ".store.bicycle.color", in)
+	assert.NoError(t, err)
+	assert.Equal(t, "red", out)
+
+	out, err = JQ(ctx, ".store.bicycle.price", in)
+	assert.NoError(t, err)
+	assert.Equal(t, 19.95, out)
+
+	out, err = JQ(ctx, ".store.bogus", in)
+	assert.NoError(t, err)
+	assert.Nil(t, out)
+
+	_, err = JQ(ctx, "{.store.unclosed", in)
+	assert.Error(t, err)
+
+	out, err = JQ(ctx, ".store", in)
+	assert.NoError(t, err)
+	assert.EqualValues(t, in["store"], out)
+
+	out, err = JQ(ctx, ".store.book[].author", in)
+	assert.NoError(t, err)
+	assert.Len(t, out, 4)
+	assert.Contains(t, out, "Nigel Rees")
+	assert.Contains(t, out, "Evelyn Waugh")
+	assert.Contains(t, out, "Herman Melville")
+	assert.Contains(t, out, "J. R. R. Tolkien")
+
+	out, err = JQ(ctx, ".store.book[]|select(.price < 10.0 )", in)
+	assert.NoError(t, err)
+	expected := []interface{}{
+		map[string]interface{}{
+			"category": "reference",
+			"author":   "Nigel Rees",
+			"title":    "Sayings of the Century",
+			"price":    8.95,
+		},
+		map[string]interface{}{
+			"category": "fiction",
+			"author":   "Herman Melville",
+			"title":    "Moby Dick",
+			"isbn":     "0-553-21311-3",
+			"price":    8.99,
+		},
+	}
+	assert.EqualValues(t, expected, out)
+
+	in = map[string]interface{}{
+		"a": map[string]interface{}{
+			"aa": map[string]interface{}{
+				"foo": map[string]interface{}{
+					"aaa": map[string]interface{}{
+						"aaaa": map[string]interface{}{
+							"bar": 1234,
+						},
+					},
+				},
+			},
+			"ab": map[string]interface{}{
+				"aba": map[string]interface{}{
+					"foo": map[string]interface{}{
+						"abaa": true,
+						"abab": "baz",
+					},
+				},
+			},
+		},
+	}
+	out, err = JQ(ctx, `tostream|select((.[0]|index("foo")) and (.[0][-1]!="foo") and (.[1])) as $s|($s[0]|index("foo")+1) as $ind|($ind|truncate_stream($s)) as $newstream|$newstream|reduce . as [$p,$v] ({};setpath($p;$v))|add`, in)
+	assert.NoError(t, err)
+	assert.Len(t, out, 3)
+	assert.Contains(t, out, map[string]interface{}{"aaaa": map[string]interface{}{"bar": 1234}})
+	assert.Contains(t, out, true)
+	assert.Contains(t, out, "baz")
+}
+
+func TestJQ_typeConversions(t *testing.T) {
+	ctx := context.Background()
+
+	type bicycleType struct {
+		Color string
+	}
+	type storeType struct {
+		Bicycle *bicycleType
+		safe    interface{}
+	}
+
+	structIn := &storeType{
+		Bicycle: &bicycleType{
+			Color: "red",
+		},
+		safe: "hidden",
+	}
+
+	out, err := JQ(ctx, ".Bicycle.Color", structIn)
+	assert.NoError(t, err)
+	assert.Equal(t, "red", out)
+
+	out, err = JQ(ctx, ".safe", structIn)
+	assert.NoError(t, err)
+	assert.Nil(t, out)
+
+	_, err = JQ(ctx, ".*", structIn)
+	assert.Error(t, err)
+
+	// a type with an underlying type of map[string]interface{}, just like
+	// gomplate.tmplctx
+	type mapType map[string]interface{}
+
+	out, err = JQ(ctx, ".foo", mapType{"foo": "bar"})
+	assert.NoError(t, err)
+	assert.Equal(t, "bar", out)
+
+	// sometimes it'll be a pointer...
+	out, err = JQ(ctx, ".foo", &mapType{"foo": "bar"})
+	assert.NoError(t, err)
+	assert.Equal(t, "bar", out)
+
+	// underlying slice type
+	type sliceType []interface{}
+
+	out, err = JQ(ctx, ".[1]", sliceType{"foo", "bar"})
+	assert.NoError(t, err)
+	assert.Equal(t, "bar", out)
+
+	out, err = JQ(ctx, ".[2]", &sliceType{"foo", "bar", "baz"})
+	assert.NoError(t, err)
+	assert.Equal(t, "baz", out)
+
+	// other basic types
+	out, err = JQ(ctx, ".", []byte("hello"))
+	assert.NoError(t, err)
+	assert.EqualValues(t, "hello", out)
+
+	out, err = JQ(ctx, ".", "hello")
+	assert.NoError(t, err)
+	assert.EqualValues(t, "hello", out)
+
+	out, err = JQ(ctx, ".", 1234)
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1234, out)
+
+	out, err = JQ(ctx, ".", true)
+	assert.NoError(t, err)
+	assert.EqualValues(t, true, out)
+
+	out, err = JQ(ctx, ".", nil)
+	assert.NoError(t, err)
+	assert.Nil(t, out)
+
+	// underlying basic types
+	type intType int
+	out, err = JQ(ctx, ".", intType(1234))
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1234, out)
+
+	type byteArrayType []byte
+	out, err = JQ(ctx, ".", byteArrayType("hello"))
+	assert.NoError(t, err)
+	assert.EqualValues(t, "hello", out)
+}
+
+func TestJQConvertType_passthroughTypes(t *testing.T) {
+	// non-marshalable values, like recursive structs, can't be used
+	type recursive struct{ Self *recursive }
+	v := &recursive{}
+	v.Self = v
+	_, err := jqConvertType(v)
+	assert.Error(t, err)
+
+	testdata := []interface{}{
+		map[string]interface{}{"foo": 1234},
+		[]interface{}{"foo", "bar", "baz", 1, 2, 3},
+		"foo",
+		[]byte("foo"),
+		json.RawMessage(`{"foo": "bar"}`),
+		true,
+		nil,
+		int(1234), int8(123), int16(123), int32(123), int64(123),
+		uint(123), uint8(123), uint16(123), uint32(123), uint64(123),
+		float32(123.45), float64(123.45),
+	}
+
+	for _, d := range testdata {
+		out, err := jqConvertType(d)
+		assert.NoError(t, err)
+		assert.Equal(t, d, out)
+	}
+}
diff --git a/docs-src/content/functions/coll.yml b/docs-src/content/functions/coll.yml
@@ -105,6 +105,35 @@ funcs:
       - |
         $ gomplate -i '{{ .books | jsonpath `$..works[?( @.edition_count > 400 )].title` }}' -c books=https://openlibrary.org/subjects/fantasy.json
         [Alice's Adventures in Wonderland Gulliver's Travels]
+  - name: coll.JQ
+    alias: jq
+    description: |
+      Filters an input object or list using the [jq](https://stedolan.github.io/jq/) language, as implemented by [gojq](https://github.com/itchyny/gojq).
+
+      Any JSON datatype may be used as input (NOTE: strings are not JSON-parsed but passed in as is).
+      If the expression results in multiple items (no matter if streamed or as an array) they are wrapped in an array.
+      Otherwise a single item is returned (even if resulting in an array with a single contained element).
+
+      JQ filter expressions can be tested at https://jqplay.org/
+
+      See also:
+
+      - [jq manual](https://stedolan.github.io/jq/manual/)
+      - [gojq differences to jq](https://github.com/itchyny/gojq#difference-to-jq)
+    pipeline: true
+    arguments:
+      - name: expression
+        required: true
+        description: The JQ expression
+      - name: in
+        required: true
+        description: The object or list to query
+    examples:
+      - |
+        $ gomplate \
+           -i '{{ .books | jq `[.works[]|{"title":.title,"authors":[.authors[].name],"published":.first_publish_year}][0]` }}' \
+           -c books=https://openlibrary.org/subjects/fantasy.json
+        map[authors:[Lewis Carroll] published:1865 title:Alice's Adventures in Wonderland]
   - name: coll.Keys
     alias: keys
     description: |