jq: allow loading modules; implement highlight as go instead of jq code

jrockway · Jun 29, 2022 · db1002e · db1002e
1 parent dfb3d9b
commit db1002e
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 29 deletions.
diff --git a/cmd/internal/jlog/jlog.go b/cmd/internal/jlog/jlog.go
@@ -28,12 +28,13 @@ type Output struct {
 }
 
 type General struct {
-	MatchRegex   string `short:"g" long:"regex" description:"A regular expression that removes lines from the output that don't match, like grep."`
-	NoMatchRegex string `short:"G" long:"no-regex" description:"A regular expression that removes lines from the output that DO match, like 'grep -v'."`
-	JQ           string `short:"e" long:"jq" description:"A jq program to run on each record in the processed input; use this to ignore certain lines, add fields, etc.  Hint: 'select(condition)' will remove lines that don't match 'condition'."`
-	NoColor      bool   `short:"M" long:"no-color" description:"Disable the use of color." env:"JLOG_FORCE_MONOCHROME"`
-	NoMonochrome bool   `short:"c" long:"no-monochrome" description:"Force the use of color." ENV:"JLOG_FORCE_COLOR"`
-	Profile      string `long:"profile" description:"If set, collect a CPU profile and write it to this file."`
+	MatchRegex   string   `short:"g" long:"regex" description:"A regular expression that removes lines from the output that don't match, like grep."`
+	NoMatchRegex string   `short:"G" long:"no-regex" description:"A regular expression that removes lines from the output that DO match, like 'grep -v'."`
+	JQ           string   `short:"e" long:"jq" description:"A jq program to run on each record in the processed input; use this to ignore certain lines, add fields, etc.  Hint: 'select(condition)' will remove lines that don't match 'condition'."`
+	JQSearchPath []string `long:"jq-search-path" env:"JLOG_JQ_SEARCH_PATH" description:"A list of directories in which to search for JQ modules.  A path entry named (not merely ending in) .jq is automatically loaded.  When set through the environment, use ':' as the delimiter (like $PATH)." default:"~/.jlog/jq/.jq" default:"~/.jlog/jq" env-delim:":"` //nolint
+	NoColor      bool     `short:"M" long:"no-color" description:"Disable the use of color." env:"JLOG_FORCE_MONOCHROME"`
+	NoMonochrome bool     `short:"c" long:"no-monochrome" description:"Force the use of color." env:"JLOG_FORCE_COLOR"`
+	Profile      string   `long:"profile" description:"If set, collect a CPU profile and write it to this file."`
 
 	Version bool `short:"v" long:"version" description:"Print version information and exit."`
 }
@@ -178,7 +179,7 @@ func NewFilterScheme(gen General) (*parse.FilterScheme, error) { //nolint
 	if err := fsch.AddNoMatchRegex(gen.NoMatchRegex); err != nil {
 		return nil, fmt.Errorf("adding NoMatchRegex: %v", err)
 	}
-	if err := fsch.AddJQ(gen.JQ); err != nil {
+	if err := fsch.AddJQ(gen.JQ, &parse.JQOptions{SearchPath: gen.JQSearchPath}); err != nil {
 		return nil, fmt.Errorf("adding JQ: %v", err)
 	}
 	return fsch, nil

diff --git a/pkg/parse/filter.go b/pkg/parse/filter.go
@@ -3,6 +3,7 @@ package parse
 import (
 	"errors"
 	"fmt"
+	"os"
 	"regexp"
 
 	"github.com/itchyny/gojq"
@@ -34,28 +35,48 @@ func prepareVariables(l *line) []interface{} {
 // highlightKey is a special key that controls highlighting.
 const highlightKey = "__highlight"
 
-func compileJQ(p string) (*gojq.Code, error) {
+func compileJQ(p string, searchPath []string) (*gojq.Code, error) {
 	if p == "" {
 		return nil, nil
 	}
-	p = "def highlight($cond): . + {__highlight: $cond};\n" + p
 	q, err := gojq.Parse(p)
 	if err != nil {
 		return nil, fmt.Errorf("parsing jq program %q: %v", p, err)
 	}
-	jq, err := gojq.Compile(q, gojq.WithVariables(DefaultVariables))
+	jq, err := gojq.Compile(q,
+		gojq.WithFunction("highlight", 1, 1, func(dot interface{}, args []interface{}) interface{} {
+			hl, ok := args[0].(bool)
+			if !ok {
+				return fmt.Errorf("argument to highlight should be a boolean; not %#v", args[0])
+			}
+			if val, ok := dot.(map[string]interface{}); ok {
+				val[highlightKey] = hl
+			}
+			return dot
+		}),
+		gojq.WithEnvironLoader(os.Environ),
+		gojq.WithVariables(DefaultVariables),
+		gojq.WithModuleLoader(gojq.NewModuleLoader(searchPath)))
 	if err != nil {
 		return nil, fmt.Errorf("compiling jq program %q: %v", p, err)
 	}
 	return jq, nil
 }
 
+type JQOptions struct {
+	SearchPath []string
+}
+
 // AddJQ compiles the provided jq program and adds it to the filter.
-func (f *FilterScheme) AddJQ(p string) error {
+func (f *FilterScheme) AddJQ(p string, opts *JQOptions) error {
 	if f.JQ != nil {
 		return errors.New("jq program already added")
 	}
-	jq, err := compileJQ(p)
+	var searchPath []string
+	if opts != nil {
+		searchPath = opts.SearchPath
+	}
+	jq, err := compileJQ(p, searchPath)
 	if err != nil {
 		return err // already has decent annotation
 	}

diff --git a/pkg/parse/filter_test.go b/pkg/parse/filter_test.go
@@ -1,6 +1,8 @@
 package parse
 
 import (
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -9,9 +11,26 @@ import (
 
 func TestJQ(t *testing.T) {
 	referenceLine := func() *line { return &line{msg: "foo", fields: map[string]interface{}{"foo": 42, "bar": "hi"}} }
+	tmpdir, err := os.MkdirTemp("", "jlog-test-jq-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := os.RemoveAll(tmpdir); err != nil {
+			t.Fatalf("cleanup: %v", err)
+		}
+	})
+	if err := os.WriteFile(filepath.Join(tmpdir, ".jq"), []byte(`def initFunction: {"init": true};`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(tmpdir, "foo.jq"), []byte(`def no: select($MSG|test("foo")|not);`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+
 	testData := []struct {
 		jq           string
 		l            *line
+		searchPath   []string
 		wantLine     *line
 		wantFiltered bool
 		wantErr      error
@@ -79,22 +98,51 @@ func TestJQ(t *testing.T) {
 			wantFiltered: true,
 			wantErr:      nil,
 		},
+		{
+			jq:           ".",
+			searchPath:   []string{filepath.Join(tmpdir, ".jq"), tmpdir},
+			l:            referenceLine(),
+			wantLine:     referenceLine(),
+			wantFiltered: false,
+			wantErr:      nil,
+		},
+		{
+			jq:         "initFunction",
+			searchPath: []string{filepath.Join(tmpdir, ".jq"), tmpdir},
+			l:          referenceLine(),
+			wantLine: &line{
+				msg:    "foo",
+				fields: map[string]interface{}{"init": true},
+			},
+			wantFiltered: false,
+			wantErr:      nil,
+		},
+		{
+			jq:           `import "foo" as foo; foo::no`,
+			searchPath:   []string{filepath.Join(tmpdir, ".jq"), tmpdir},
+			l:            referenceLine(),
+			wantLine:     referenceLine(),
+			wantFiltered: true,
+			wantErr:      nil,
+		},
 	}
 	for _, test := range testData {
-		fs := new(FilterScheme)
-		if err := fs.AddJQ(test.jq); err != nil {
-			t.Fatal(err)
-		}
-		gotFiltered, gotErr := fs.runJQ(test.l)
-		if diff := cmp.Diff(test.l, test.wantLine, cmp.AllowUnexported(line{}), cmpopts.EquateEmpty()); diff != "" {
-			t.Errorf("line: %s", diff)
-		}
-		if got, want := gotFiltered, test.wantFiltered; got != want {
-			t.Errorf("filtered:\n  got: %v\n want: %v", got, want)
-		}
-		if got, want := gotErr, test.wantErr; !comperror(got, want) {
-			t.Errorf("error:\n  got: %v\n want: %v", got, want)
-		}
+		t.Run(test.jq, func(t *testing.T) {
+			fs := new(FilterScheme)
+			if err := fs.AddJQ(test.jq, &JQOptions{SearchPath: test.searchPath}); err != nil {
+				t.Fatal(err)
+			}
+			gotFiltered, gotErr := fs.runJQ(test.l)
+			if diff := cmp.Diff(test.l, test.wantLine, cmp.AllowUnexported(line{}), cmpopts.EquateEmpty()); diff != "" {
+				t.Errorf("line: %s", diff)
+			}
+			if got, want := gotFiltered, test.wantFiltered; got != want {
+				t.Errorf("filtered:\n  got: %v\n want: %v", got, want)
+			}
+			if got, want := gotErr, test.wantErr; !comperror(got, want) {
+				t.Errorf("error:\n  got: %v\n want: %v", got, want)
+			}
+		})
 	}
 }
 
@@ -189,7 +237,7 @@ func TestAdds(t *testing.T) {
 			f := new(FilterScheme)
 			var errs []error
 			for _, jq := range test.jq {
-				if err := f.AddJQ(jq); err != nil {
+				if err := f.AddJQ(jq, nil); err != nil {
 					errs = append(errs, err)
 				}
 			}

diff --git a/pkg/parse/parse_test.go b/pkg/parse/parse_test.go
@@ -954,7 +954,7 @@ func TestReadLog(t *testing.T) {
 
 		t.Run(test.name, func(t *testing.T) {
 			fs := new(FilterScheme)
-			if err := fs.AddJQ(test.jq); err != nil {
+			if err := fs.AddJQ(test.jq, nil); err != nil {
 				t.Fatalf("add jq: %v", err)
 			}
 			if err := fs.AddMatchRegex(test.matchrx); err != nil {
@@ -1277,7 +1277,7 @@ func TestFullLog(t *testing.T) {
 	for _, test := range testData {
 		t.Run(test.name, func(t *testing.T) {
 			fs := new(FilterScheme)
-			if err := fs.AddJQ(test.jq); err != nil {
+			if err := fs.AddJQ(test.jq, nil); err != nil {
 				t.Fatal(err)
 			}
 			if err := fs.AddMatchRegex(test.matchregex); err != nil {