jq.go

  1package shell
  2
  3import (
  4	"encoding/json"
  5	"fmt"
  6	"io"
  7	"os"
  8	"strings"
  9
 10	"github.com/itchyny/gojq"
 11	"mvdan.cc/sh/v3/interp"
 12)
 13
 14const jqUsage = `jq - Go implementation of jq (gojq 0.12.19 builtin)
 15
 16Synopsis:
 17  %% echo '{"foo": 128}' | jq '.foo'
 18
 19Usage:
 20  jq [OPTIONS] [FILTER] [FILE...]
 21
 22Options:
 23  -r, --raw-output              output raw strings
 24  -j, --join-output             implies -r with no newline delimiter
 25  -c, --compact-output          output without pretty-printing
 26  -s, --slurp                   read all inputs into an array
 27  -n, --null-input              use null as input value
 28  -e, --exit-status             exit 1 when the last value is false or null
 29  -R, --raw-input               read input as raw strings
 30      --arg name value          set a string value to a variable
 31      --argjson name value      set a JSON value to a variable
 32  -h, --help                    display this help
 33`
 34
 35// handleJQ implements the jq builtin using gojq. It supports a subset of jq
 36// flags: -r (raw output), -c (compact output), -s (slurp), -n (null input),
 37// -e (exit status), -R (raw input), and --arg name value.
 38//
 39// Note that this is somewhat of a reimplmentation of the CLI of the glorious
 40// github.com/itchyny/gojq, and we'd ideally get the CLI exposed upstream to
 41// avoid this falling out of sync.
 42func handleJQ(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
 43	var (
 44		rawOutput  bool
 45		compact    bool
 46		slurp      bool
 47		nullInput  bool
 48		exitStatus bool
 49		rawInput   bool
 50		joinOutput bool
 51		argNames   []string
 52		argValues  []any
 53	)
 54
 55	// Parse flags and extract the query.
 56	var queryStr string
 57	var fileArgs []string
 58	i := 1 // skip "jq"
 59	for i < len(args) {
 60		arg := args[i]
 61		switch {
 62		case arg == "-h" || arg == "--help":
 63			fmt.Fprint(stdout, jqUsage)
 64			return nil
 65		case arg == "-r" || arg == "--raw-output":
 66			rawOutput = true
 67		case arg == "-j" || arg == "--join-output":
 68			joinOutput = true
 69			rawOutput = true
 70		case arg == "-c" || arg == "--compact-output":
 71			compact = true
 72		case arg == "-s" || arg == "--slurp":
 73			slurp = true
 74		case arg == "-n" || arg == "--null-input":
 75			nullInput = true
 76		case arg == "-e" || arg == "--exit-status":
 77			exitStatus = true
 78		case arg == "-R" || arg == "--raw-input":
 79			rawInput = true
 80		case arg == "--arg":
 81			if i+2 >= len(args) {
 82				fmt.Fprintf(stderr, "jq: --arg requires name and value\n")
 83				return interp.ExitStatus(2)
 84			}
 85			argNames = append(argNames, "$"+args[i+1])
 86			argValues = append(argValues, args[i+2])
 87			i += 2
 88		case arg == "--argjson":
 89			if i+2 >= len(args) {
 90				fmt.Fprintf(stderr, "jq: --argjson requires name and value\n")
 91				return interp.ExitStatus(2)
 92			}
 93			var val any
 94			if err := json.Unmarshal([]byte(args[i+2]), &val); err != nil {
 95				fmt.Fprintf(stderr, "jq: invalid JSON for --argjson %s: %s\n", args[i+1], err)
 96				return interp.ExitStatus(2)
 97			}
 98			argNames = append(argNames, "$"+args[i+1])
 99			argValues = append(argValues, val)
100			i += 2
101		case arg == "--":
102			i++
103			// Remaining args are file arguments.
104			for i < len(args) {
105				fileArgs = append(fileArgs, args[i])
106				i++
107			}
108			continue
109		case strings.HasPrefix(arg, "-") && queryStr != "":
110			fmt.Fprintf(stderr, "jq: unknown option: %s\n", arg)
111			return interp.ExitStatus(2)
112		default:
113			if queryStr == "" {
114				queryStr = arg
115			} else {
116				fileArgs = append(fileArgs, arg)
117			}
118		}
119		i++
120	}
121
122	if queryStr == "" {
123		queryStr = "."
124	}
125
126	query, err := gojq.Parse(queryStr)
127	if err != nil {
128		fmt.Fprintf(stderr, "jq: %s\n", err)
129		return interp.ExitStatus(3)
130	}
131
132	opts := []gojq.CompilerOption{
133		gojq.WithEnvironLoader(os.Environ),
134	}
135	if len(argNames) > 0 {
136		opts = append(opts, gojq.WithVariables(argNames))
137	}
138
139	code, err := gojq.Compile(query, opts...)
140	if err != nil {
141		fmt.Fprintf(stderr, "jq: %s\n", err)
142		return interp.ExitStatus(3)
143	}
144
145	// Build input values.
146	inputs, err := readInputs(stdin, fileArgs, nullInput, rawInput, slurp)
147	if err != nil {
148		fmt.Fprintf(stderr, "jq: %s\n", err)
149		return interp.ExitStatus(2)
150	}
151
152	var lastFalsy bool
153	for _, input := range inputs {
154		iter := code.Run(input, argValues...)
155		for {
156			v, ok := iter.Next()
157			if !ok {
158				break
159			}
160			if err, ok := v.(error); ok {
161				fmt.Fprintf(stderr, "jq: %s\n", err)
162				return interp.ExitStatus(5)
163			}
164			if exitStatus {
165				lastFalsy = v == nil || v == false
166			}
167			if err := writeValue(stdout, v, rawOutput, compact, joinOutput); err != nil {
168				return err
169			}
170		}
171	}
172
173	if exitStatus && lastFalsy {
174		return interp.ExitStatus(1)
175	}
176	return nil
177}
178
179// readInputs reads JSON (or raw) input values from stdin or files.
180func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool) ([]any, error) {
181	if nullInput {
182		return []any{nil}, nil
183	}
184
185	var readers []io.Reader
186	if len(files) > 0 {
187		for _, f := range files {
188			file, err := os.Open(f)
189			if err != nil {
190				return nil, err
191			}
192			defer file.Close()
193			readers = append(readers, file)
194		}
195	} else {
196		readers = []io.Reader{stdin}
197	}
198
199	var vals []any
200	for _, r := range readers {
201		data, err := io.ReadAll(r)
202		if err != nil {
203			return nil, err
204		}
205
206		if rawInput {
207			lines := strings.Split(string(data), "\n")
208			if slurp {
209				vals = append(vals, strings.Join(lines, "\n"))
210			} else {
211				for _, line := range lines {
212					if line != "" || !slurp {
213						vals = append(vals, line)
214					}
215				}
216			}
217			continue
218		}
219
220		// Decode potentially multiple JSON values from the stream.
221		dec := json.NewDecoder(strings.NewReader(string(data)))
222		var streamVals []any
223		for {
224			var v any
225			if err := dec.Decode(&v); err != nil {
226				if err == io.EOF {
227					break
228				}
229				return nil, fmt.Errorf("parse error: %w", err)
230			}
231			streamVals = append(streamVals, v)
232		}
233
234		if slurp {
235			vals = append(vals, streamVals)
236		} else {
237			vals = append(vals, streamVals...)
238		}
239	}
240
241	if len(vals) == 0 {
242		return []any{nil}, nil
243	}
244	return vals, nil
245}
246
247// writeValue writes a single jq output value.
248func writeValue(w io.Writer, v any, raw, compact, join bool) error {
249	if raw {
250		if s, ok := v.(string); ok {
251			if _, err := fmt.Fprint(w, s); err != nil {
252				return err
253			}
254			if !join {
255				_, err := fmt.Fprint(w, "\n")
256				return err
257			}
258			return nil
259		}
260	}
261
262	var bs []byte
263	var err error
264	if compact {
265		bs, err = gojq.Marshal(v)
266	} else {
267		bs, err = json.MarshalIndent(v, "", "  ")
268	}
269	if err != nil {
270		return err
271	}
272	if _, writeErr := w.Write(bs); writeErr != nil {
273		return writeErr
274	}
275	_, err = fmt.Fprint(w, "\n")
276	return err
277}