From 80d3d293e5dd81a4d7cf1092b1dbf6288bdafe66 Mon Sep 17 00:00:00 2001
From: Robert Lin <robert@bobheadxi.dev>
Date: Fri, 27 May 2022 09:59:00 -0700
Subject: [PATCH] output: use io.Reader implementation more internally

---
 jq.go     | 31 +++++++++++++++++++++----------
 map.go    |  2 +-
 output.go | 17 +++++++----------
 3 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/jq.go b/jq.go
index 00b10e0..20495fb 100644
--- a/jq.go
+++ b/jq.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 
 	"github.com/itchyny/gojq"
 )
@@ -22,18 +23,28 @@ func buildJQ(query string) (*gojq.Code, error) {
 	return jqCode, nil
 }
 
-// execJQ executes the compiled jq query against content.
-func execJQ(ctx context.Context, jqCode *gojq.Code, content []byte) ([]byte, error) {
+// execJQBytes can be used to execute a compiled jq query against small content bytes,
+// e.g. lines. Errors are annotated with the provided content for ease of debugging.
+func execJQBytes(ctx context.Context, jqCoode *gojq.Code, content []byte) ([]byte, error) {
 	if len(content) == 0 {
 		return nil, nil
 	}
+	result, err := execJQ(ctx, jqCoode, bytes.NewReader(content))
+	if err != nil {
+		// Embed the consumed content
+		return nil, fmt.Errorf("%w: %s", err, string(content))
+	}
+	return result, nil
+}
 
+// execJQ executes the compiled jq query against content from reader.
+func execJQ(ctx context.Context, jqCode *gojq.Code, reader io.Reader) ([]byte, error) {
 	var input interface{}
-	if err := json.NewDecoder(bytes.NewReader(content)).Decode(&input); err != nil {
-		return nil, fmt.Errorf("json: %w: %s", err, string(content))
+	if err := json.NewDecoder(reader).Decode(&input); err != nil {
+		return nil, fmt.Errorf("json: %w", err)
 	}
 
-	var newLine bytes.Buffer
+	var result bytes.Buffer
 	iter := jqCode.RunWithContext(ctx, input)
 	for {
 		v, ok := iter.Next()
@@ -42,14 +53,14 @@ func execJQ(ctx context.Context, jqCode *gojq.Code, content []byte) ([]byte, err
 		}
 
 		if err, ok := v.(error); ok {
-			return nil, fmt.Errorf("jq: %w: %s", err, string(content))
+			return nil, fmt.Errorf("jq: %w", err)
 		}
 
-		result, err := gojq.Marshal(v)
+		encoded, err := gojq.Marshal(v)
 		if err != nil {
-			return nil, fmt.Errorf("jq: %w: %s", err, string(content))
+			return nil, fmt.Errorf("jq: %w", err)
 		}
-		newLine.Write(result)
+		result.Write(encoded)
 	}
-	return newLine.Bytes(), nil
+	return result.Bytes(), nil
 }
diff --git a/map.go b/map.go
index aff0e41..6dcac2b 100644
--- a/map.go
+++ b/map.go
@@ -27,7 +27,7 @@ func MapJQ(query string) (LineMap, error) {
 	}
 
 	return func(ctx context.Context, line []byte, dst io.Writer) (int, error) {
-		b, err := execJQ(ctx, jqCode, line)
+		b, err := execJQBytes(ctx, jqCode, line)
 		if err != nil {
 			return 0, err
 		}
diff --git a/output.go b/output.go
index cceffe9..924975e 100644
--- a/output.go
+++ b/output.go
@@ -180,22 +180,19 @@ func (o *commandOutput) JQ(query string) ([]byte, error) {
 		return nil, err
 	}
 
-	var buffer bytes.Buffer
-	if err := o.Stream(&buffer); err != nil {
-		return nil, err
-	}
-
-	b, err := execJQ(o.ctx, jqCode, buffer.Bytes())
+	result, err := execJQ(o.ctx, jqCode, o)
 	if err != nil {
 		return nil, err
 	}
-	return b, nil
+	return result, nil
 }
 
 func (o *commandOutput) String() (string, error) {
-	var sb strings.Builder
-	err := o.Stream(&sb)
-	return strings.TrimSuffix(sb.String(), "\n"), err
+	b, err := io.ReadAll(o)
+	if err != nil {
+		return "", err
+	}
+	return strings.TrimSuffix(string(b), "\n"), nil
 }
 
 func (o *commandOutput) Read(p []byte) (int, error) {