From db1002e73d3a61b14bd32a8818372b1f3633ffb1 Mon Sep 17 00:00:00 2001 From: Jonathan Rockway <2367+jrockway@users.noreply.github.com> Date: Wed, 29 Jun 2022 02:00:02 -0400 Subject: [PATCH] jq: allow loading modules; implement highlight as go instead of jq code --- cmd/internal/jlog/jlog.go | 15 ++++---- pkg/parse/filter.go | 31 +++++++++++++--- pkg/parse/filter_test.go | 78 +++++++++++++++++++++++++++++++-------- pkg/parse/parse_test.go | 4 +- 4 files changed, 99 insertions(+), 29 deletions(-) diff --git a/cmd/internal/jlog/jlog.go b/cmd/internal/jlog/jlog.go index 5c18ed7..d5fa436 100644 --- 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 index 58b2e66..bb19a86 100644 --- 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 index 0084114..ecd79d2 100644 --- 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 index 94f2b60..f2260d4 100644 --- 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 {