Propagate lazy query errors through jq JSON conversion

jqresult/compile.go

@@ -4,11 +4,27 @@ import (
 	"github.com/wader/gojq"
 )
 
+// tojsonPrelude is parsed once and inserted before the user's function
+// definitions, so a later user def tojson or def format still wins.
+// The native format helper is bound first, so format("text") and format("html")
+// keep the builtin. format("json"), tojson, and @json use the helper that
+// returns materialization errors instead of panicking.
+var tojsonPrelude = syncOncePrelude()
+
+func syncOncePrelude() *gojq.Query {
+	q, err := gojq.Parse(`def _execspansql_native_format($f): format($f); def tojson: _execspansql_tojson; def format($f): if $f == "json" then _execspansql_tojson else _execspansql_native_format($f) end;`)
+	if err != nil {
+		panic(err)
+	}
+	return q
+}
+
 // Compile parses filter and returns executable jq code.
 func Compile(filter string) (*gojq.Code, error) {
 	q, err := gojq.Parse(filter)
 	if err != nil {
 		return nil, err
 	}
-	return gojq.Compile(q)
+	q.FuncDefs = append(append([]*gojq.FuncDef{}, tojsonPrelude.FuncDefs...), q.FuncDefs...)
+	return gojq.Compile(q, gojq.WithFunction("_execspansql_tojson", 0, 0, execspansqlToJSON))
 }

jqresult/print_panic_test.go

@@ -0,0 +1,28 @@
+package jqresult
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/wader/gojq"
+)
+
+type boomIter struct{}
+
+func (boomIter) Next() (any, bool) { panic("unrelated boom") }
+
+func TestPrintRethrowsUnrelatedPanic(t *testing.T) {
+	defer func() {
+		p := recover()
+		if p == nil || !strings.Contains(p.(string), "unrelated boom") {
+			t.Fatalf("panic = %#v, want unrelated boom", p)
+		}
+	}()
+	_ = Print(nopEncoder{}, boomIter{})
+}
+
+type nopEncoder struct{}
+
+func (nopEncoder) Encode(any) error { return nil }
+
+var _ gojq.Iter = boomIter{}

jqresult/tojson.go

@@ -0,0 +1,53 @@
+package jqresult
+
+import "github.com/wader/gojq"
+
+// execspansqlToJSON is the body of the prelude's tojson. Materialization
+// errors are returned as the original error so try/catch and the caller both
+// see them. Plain values are encoded with gojq.Marshal, which keeps jq's
+// number and string formatting.
+func execspansqlToJSON(v any, _ []any) any {
+	normalized, err := normalizeJQValue(v)
+	if err != nil {
+		return err
+	}
+	b, err := gojq.Marshal(normalized)
+	if err != nil {
+		return err
+	}
+	return string(b)
+}
+
+func normalizeJQValue(v any) (any, error) {
+	if jqv, ok := v.(gojq.JQValue); ok {
+		converted := jqv.JQValueToGoJQ()
+		if err, ok := converted.(error); ok {
+			return nil, err
+		}
+		return normalizeJQValue(converted)
+	}
+	switch x := v.(type) {
+	case []any:
+		out := make([]any, len(x))
+		for i, elem := range x {
+			n, err := normalizeJQValue(elem)
+			if err != nil {
+				return nil, err
+			}
+			out[i] = n
+		}
+		return out, nil
+	case map[string]any:
+		out := make(map[string]any, len(x))
+		for k, elem := range x {
+			n, err := normalizeJQValue(elem)
+			if err != nil {
+				return nil, err
+			}
+			out[k] = n
+		}
+		return out, nil
+	default:
+		return v, nil
+	}
+}

jqresult/tojson_test.go

@@ -0,0 +1,55 @@
+package jqresult
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestFormatAliasesKeepNonJSONBehavior(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		filter string
+		want   string
+		isErr  bool
+	}{
+		{filter: `format("text")`, want: "hi"},
+		{filter: `format("html")`, want: "hi"},
+		{filter: `format(3)`, isErr: true},
+		{filter: `def tojson: "user-defined"; tojson`, want: "user-defined"},
+		{filter: `def format($f): "user-format"; format("json")`, want: "user-format"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.filter, func(t *testing.T) {
+			code, err := Compile(tc.filter)
+			if err != nil {
+				t.Fatal(err)
+			}
+			v, ok := code.Run("hi").Next()
+			if !ok {
+				t.Fatal("no output")
+			}
+			if tc.isErr {
+				if _, ok := v.(error); !ok {
+					t.Fatalf("got %#v, want a format type error", v)
+				}
+				return
+			}
+			if err, ok := v.(error); ok {
+				t.Fatal(err)
+			}
+			if !strings.Contains(anyString(v), tc.want) {
+				t.Fatalf("got %#v, want %q", v, tc.want)
+			}
+		})
+	}
+}
+
+func anyString(v any) string {
+	switch x := v.(type) {
+	case string:
+		return x
+	default:
+		return ""
+	}
+}

lazy_json_error_test.go

@@ -0,0 +1,145 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type lazyJSONErrorServer struct {
+	queryStatsModeServer
+}
+
+func (s *lazyJSONErrorServer) ExecuteStreamingSql(*sppb.ExecuteSqlRequest, sppb.Spanner_ExecuteStreamingSqlServer) error {
+	return status.Error(codes.InvalidArgument, "invalid SQL from review fixture")
+}
+
+func TestLazyToJSONReturnsOriginalQueryError(t *testing.T) {
+	startQueryStatsModeServer(t, &lazyJSONErrorServer{})
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=THIS IS INVALID SQL",
+		"--jq-input-mode=lazy", "--filter=.rows | tojson",
+		"--output=" + filepath.Join(t.TempDir(), "out"), "--timeout=5s",
+	})
+	if err == nil {
+		t.Fatal("expected query error")
+	}
+	if status.Code(err) != codes.InvalidArgument {
+		t.Fatalf("code = %s, error = %v, want InvalidArgument", status.Code(err), err)
+	}
+	if !strings.Contains(err.Error(), "invalid SQL from review fixture") {
+		t.Fatalf("error = %v, want original message", err)
+	}
+}
+
+func TestLazyToJSONTryCatchSeesQueryError(t *testing.T) {
+	startQueryStatsModeServer(t, &lazyJSONErrorServer{})
+	path := filepath.Join(t.TempDir(), "out")
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=THIS IS INVALID SQL",
+		"--jq-input-mode=lazy", `--filter=try (.rows | tojson) catch .`,
+		"--output=" + path, "--timeout=5s",
+	})
+	if err != nil {
+		t.Fatalf("try/catch returned %v", err)
+	}
+	got, readErr := os.ReadFile(path)
+	if readErr != nil {
+		t.Fatal(readErr)
+	}
+	if !strings.Contains(string(got), "invalid SQL from review fixture") {
+		t.Fatalf("catch output = %q", got)
+	}
+}
+
+func TestLazyAtJSONReturnsOriginalQueryError(t *testing.T) {
+	startQueryStatsModeServer(t, &lazyJSONErrorServer{})
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=THIS IS INVALID SQL",
+		"--jq-input-mode=lazy", "--filter=@json",
+		"--output=" + filepath.Join(t.TempDir(), "out"), "--timeout=5s",
+	})
+	if err == nil {
+		t.Fatal("expected query error")
+	}
+	if status.Code(err) != codes.InvalidArgument || !strings.Contains(err.Error(), "invalid SQL from review fixture") {
+		t.Fatalf("error = %v, want InvalidArgument and the original message", err)
+	}
+}
+
+func TestLazyUserDefinedToJSONWins(t *testing.T) {
+	startQueryStatsModeServer(t, &queryStatsModeServer{})
+	path := filepath.Join(t.TempDir(), "out")
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=SELECT value FROM T",
+		"--jq-input-mode=lazy", `--filter=def tojson: "user-defined"; tojson`,
+		"--output=" + path, "--timeout=5s",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, readErr := os.ReadFile(path)
+	if readErr != nil {
+		t.Fatal(readErr)
+	}
+	if !strings.Contains(string(got), "user-defined") {
+		t.Fatalf("output = %q, want the user definition", got)
+	}
+}
+
+func TestLazyFormatJSONTryCatchSeesQueryError(t *testing.T) {
+	startQueryStatsModeServer(t, &lazyJSONErrorServer{})
+	path := filepath.Join(t.TempDir(), "out")
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=THIS IS INVALID SQL",
+		"--jq-input-mode=lazy", `--filter=try format("json") catch .`,
+		"--output=" + path, "--timeout=5s",
+	})
+	if err != nil {
+		t.Fatalf("try/catch returned %v", err)
+	}
+	got, readErr := os.ReadFile(path)
+	if readErr != nil {
+		t.Fatal(readErr)
+	}
+	if !strings.Contains(string(got), "invalid SQL from review fixture") {
+		t.Fatalf("catch output = %q", got)
+	}
+}
+
+func TestLazyFormatJSONReturnsOriginalQueryError(t *testing.T) {
+	startQueryStatsModeServer(t, &lazyJSONErrorServer{})
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=THIS IS INVALID SQL",
+		"--jq-input-mode=lazy", `--filter=format("json")`,
+		"--output=" + filepath.Join(t.TempDir(), "out"), "--timeout=5s",
+	})
+	if status.Code(err) != codes.InvalidArgument || !strings.Contains(err.Error(), "invalid SQL from review fixture") {
+		t.Fatalf("error = %v, want InvalidArgument and the original message", err)
+	}
+}
+
+func TestLazyToJSONHealthyRows(t *testing.T) {
+	startQueryStatsModeServer(t, &queryStatsModeServer{})
+	path := filepath.Join(t.TempDir(), "out")
+	err := runCLI(t.Context(), []string{
+		"db", "--project=p", "--instance=i", "--sql=SELECT value FROM T",
+		"--jq-input-mode=lazy", "--filter=.rows | tojson",
+		"--output=" + path, "--timeout=5s",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, readErr := os.ReadFile(path)
+	if readErr != nil {
+		t.Fatal(readErr)
+	}
+	if !strings.Contains(string(got), "value") {
+		t.Fatalf("tojson = %q, want row contents", got)
+	}
+}