Pretty-print and color SQL results in terminals

README.md

@@ -239,6 +239,10 @@ SQL queries without an explicit `LIMIT` return at most 500 rows. Add a SQL
 does not fetch additional pages. Interactive mode (`-i`) fetches 500-row
 windows up to the SQL limit.
 
+Non-interactive SQL results are pretty-printed JSON arrays by default, with
+theme colors in a terminal. Piped output and `-o json` stay uncolored for
+scripts. Use `-i` for the interactive table view.
+
 OTel fields with dots like `service.name` and `http.status_code` work directly in queries without manual quoting.
 
 ## PromQL Workflows

cmd/query.go

@@ -28,6 +28,7 @@ import (
 	"time"
 
 	"github.com/parseablehq/pb/pkg/model"
+	"github.com/parseablehq/pb/pkg/ui"
 
 	internalHTTP "github.com/parseablehq/pb/pkg/http"
 
@@ -55,7 +56,7 @@ var query = &cobra.Command{
 	Use:          "run [query] [flags]",
 	Example:      "  pb sql run \"select * from frontend\" --from=10m --to=now\n  pb sql run \"select * from frontend\" -i",
 	Short:        "Run SQL query on a dataset",
-	Long:         "\nRun SQL query on a dataset. Default output format is text.\nUse --output json for JSON output, or -i for interactive table view.\nQueries without a SQL LIMIT return at most 500 rows.\nExplicit limits are not capped at 500 by pb; choose a reasonable value.",
+	Long:         "\nRun SQL query on a dataset. Results are pretty-printed JSON by default, with theme colors in a terminal.\nUse -o json for uncolored JSON or -i for the interactive table view.\nQueries without a SQL LIMIT return at most 500 rows.\nExplicit limits are not capped at 500 by pb; choose a reasonable value.",
 	Args:         cobra.MaximumNArgs(1),
 	SilenceUsage: true,
 	PreRunE:      PreRunDefaultProfile,
@@ -155,7 +156,7 @@ var query = &cobra.Command{
 func init() {
 	query.Flags().StringP(startFlag, startFlagShort, defaultStart, "Start time for query.")
 	query.Flags().StringP(endFlag, endFlagShort, defaultEnd, "End time for query.")
-	query.Flags().StringVarP(&outputFormat, "output", "o", "", "Output format (text|json)")
+	query.Flags().StringVarP(&outputFormat, "output", "o", "", "Output format (text|json); text colors JSON in terminals")
 	query.Flags().BoolP("interactive", "i", false, "Open interactive table view")
 	query.Flags().StringVar(&saveAsName, "save-as", "", "Save this query with a name for later use")
 }
@@ -746,7 +747,9 @@ func fetchData(client *internalHTTP.HTTPClient, query string, startTime, endTime
 	if outputFormat == "json" {
 		return streamSQLJSONResponse(os.Stdout, reader, resp.Status)
 	}
-	return streamSQLTextResponse(os.Stdout, reader, resp.Status)
+	color, restore := setupSQLTextColor()
+	defer restore()
+	return streamSQLTextResponse(os.Stdout, reader, resp.Status, color)
 }
 
 func readLimitedErrorPreview(body io.Reader) (string, error) {
@@ -766,49 +769,23 @@ func readLimitedErrorPreview(body io.Reader) (string, error) {
 	return preview, nil
 }
 
-func streamSQLTextResponse(out io.Writer, reader *bufio.Reader, status string) error {
-	prefix, empty, err := readResponsePrefix(reader)
-	if err != nil {
-		return fmt.Errorf("failed to read response body: %w", err)
-	}
-	if empty {
-		fmt.Fprintf(out, "No response body returned (status: %s).\n", status)
-		return nil
-	}
-
-	body := io.Reader(reader)
-	if firstNonSpace(prefix) == '[' {
-		isEmpty, consumed, err := consumeEmptyJSONArray(reader)
-		if err != nil {
-			return fmt.Errorf("failed to read response body: %w", err)
-		}
-		if isEmpty {
-			fmt.Fprintf(out, "Query succeeded: no rows returned (status: %s).\n", status)
-			return nil
-		}
-		body = io.MultiReader(bytes.NewReader(consumed), reader)
-	}
-
-	tracker := &trailingNewlineWriter{w: out}
-	if _, err := io.Copy(tracker, io.MultiReader(bytes.NewReader(prefix), body)); err != nil {
-		return fmt.Errorf("failed to stream response body: %w", err)
-	}
-	if tracker.wrote && tracker.last != '\n' {
-		fmt.Fprintln(out)
-	}
-	return nil
+func streamSQLTextResponse(out io.Writer, reader *bufio.Reader, status string, color bool) error {
+	return streamSQLResponse(out, reader, status, color)
 }
 
 func streamSQLJSONResponse(out io.Writer, reader *bufio.Reader, status string) error {
+	return streamSQLResponse(out, reader, status, false)
+}
+
+func streamSQLResponse(out io.Writer, reader *bufio.Reader, status string, color bool) error {
 	prefix, empty, err := readResponsePrefix(reader)
 	if err != nil {
 		return fmt.Errorf("failed to read response body: %w", err)
 	}
 	if empty {
-		fmt.Fprintf(out, "No response body returned (status: %s).\n", status)
-		return nil
+		return fmt.Errorf("empty SQL response body (status: %s)", status)
 	}
-	if err := writePrettyJSONArray(out, io.MultiReader(bytes.NewReader(prefix), reader)); err != nil {
+	if err := writePrettyJSONArray(out, io.MultiReader(bytes.NewReader(prefix), reader), color); err != nil {
 		return fmt.Errorf("error decoding JSON response: %w", err)
 	}
 	return nil
@@ -831,42 +808,7 @@ func readResponsePrefix(reader *bufio.Reader) ([]byte, bool, error) {
 	}
 }
 
-func consumeEmptyJSONArray(reader *bufio.Reader) (bool, []byte, error) {
-	var consumed []byte
-	for {
-		b, err := reader.ReadByte()
-		if err == io.EOF {
-			return false, consumed, nil
-		}
-		if err != nil {
-			return false, consumed, err
-		}
-		consumed = append(consumed, b)
-		if isSQLSpace(b) {
-			continue
-		}
-		if b != ']' {
-			return false, consumed, nil
-		}
-		break
-	}
-
-	for {
-		b, err := reader.ReadByte()
-		if err == io.EOF {
-			return true, consumed, nil
-		}
-		if err != nil {
-			return false, consumed, err
-		}
-		consumed = append(consumed, b)
-		if !isSQLSpace(b) {
-			return false, consumed, nil
-		}
-	}
-}
-
-func writePrettyJSONArray(out io.Writer, body io.Reader) error {
+func writePrettyJSONArray(out io.Writer, body io.Reader, color bool) error {
 	decoder := json.NewDecoder(body)
 	token, err := decoder.Token()
 	if err != nil {
@@ -880,12 +822,12 @@ func writePrettyJSONArray(out io.Writer, body io.Reader) error {
 	first := true
 	fmt.Fprint(out, "[")
 	for decoder.More() {
-		var item interface{}
+		var item json.RawMessage
 		if err := decoder.Decode(&item); err != nil {
 			return err
 		}
-		encoded, err := json.MarshalIndent(item, "", "  ")
-		if err != nil {
+		var encoded bytes.Buffer
+		if err := json.Indent(&encoded, item, "", "  "); err != nil {
 			return err
 		}
 		if first {
@@ -895,7 +837,11 @@ func writePrettyJSONArray(out io.Writer, body io.Reader) error {
 			fmt.Fprintln(out, ",")
 		}
 		fmt.Fprint(out, "  ")
-		fmt.Fprint(out, string(bytes.ReplaceAll(encoded, []byte("\n"), []byte("\n  "))))
+		formatted := string(bytes.ReplaceAll(encoded.Bytes(), []byte("\n"), []byte("\n  ")))
+		if color {
+			formatted = ui.HighlightJSON(formatted)
+		}
+		fmt.Fprint(out, formatted)
 	}
 	if _, err := decoder.Token(); err != nil {
 		return err
@@ -909,30 +855,6 @@ func writePrettyJSONArray(out io.Writer, body io.Reader) error {
 	return nil
 }
 
-func firstNonSpace(data []byte) byte {
-	for _, b := range data {
-		if !isSQLSpace(b) {
-			return b
-		}
-	}
-	return 0
-}
-
-type trailingNewlineWriter struct {
-	w     io.Writer
-	wrote bool
-	last  byte
-}
-
-func (w *trailingNewlineWriter) Write(p []byte) (int, error) {
-	n, err := w.w.Write(p)
-	if n > 0 {
-		w.wrote = true
-		w.last = p[n-1]
-	}
-	return n, err
-}
-
 func extractStreamName(query string) string {
 	re := regexp.MustCompile(`(?i)\bfrom\s+(?:"([^"]+)"|([a-zA-Z_][a-zA-Z0-9_-]*))`)
 	m := re.FindStringSubmatch(query)

cmd/query_test.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"bufio"
 	"bytes"
+	"regexp"
 	"strings"
 	"testing"
 )
@@ -149,32 +150,60 @@ func TestEnsureDefaultLimitKeepsRealLimitWithComments(t *testing.T) {
 	}
 }
 
-func TestStreamSQLTextResponseStreamsBodyAndAddsTrailingNewline(t *testing.T) {
+func TestStreamSQLTextResponsePrettyPrintsJSON(t *testing.T) {
 	var output bytes.Buffer
 	reader := bufio.NewReader(strings.NewReader(`[{"a":1}]`))
 
-	if err := streamSQLTextResponse(&output, reader, "200 OK"); err != nil {
+	if err := streamSQLTextResponse(&output, reader, "200 OK", false); err != nil {
 		t.Fatalf("streamSQLTextResponse failed: %v", err)
 	}
 
-	if got, want := output.String(), "[{\"a\":1}]\n"; got != want {
+	if got, want := output.String(), "[\n  {\n    \"a\": 1\n  }\n]\n"; got != want {
 		t.Fatalf("got %q, want %q", got, want)
 	}
 }
 
-func TestStreamSQLTextResponsePrintsNoRowsForEmptyArray(t *testing.T) {
+func TestStreamSQLTextResponsePrintsEmptyArray(t *testing.T) {
 	var output bytes.Buffer
 	reader := bufio.NewReader(strings.NewReader("  []\n"))
 
-	if err := streamSQLTextResponse(&output, reader, "200 OK"); err != nil {
+	if err := streamSQLTextResponse(&output, reader, "200 OK", false); err != nil {
 		t.Fatalf("streamSQLTextResponse failed: %v", err)
 	}
 
-	if got, want := output.String(), "Query succeeded: no rows returned (status: 200 OK).\n"; got != want {
+	if got, want := output.String(), "[]\n"; got != want {
 		t.Fatalf("got %q, want %q", got, want)
 	}
 }
 
+func TestStreamSQLTextResponseColorsTerminalJSON(t *testing.T) {
+	var output bytes.Buffer
+	reader := bufio.NewReader(strings.NewReader(`[{"c":2400}]`))
+
+	if err := streamSQLTextResponse(&output, reader, "200 OK", true); err != nil {
+		t.Fatalf("streamSQLTextResponse failed: %v", err)
+	}
+	if !strings.Contains(output.String(), "\x1b[") {
+		t.Fatalf("expected themed ANSI colors, got %q", output.String())
+	}
+	plain := regexp.MustCompile(`\x1b\[[0-9;]*m`).ReplaceAllString(output.String(), "")
+	if want := "[\n  {\n    \"c\": 2400\n  }\n]\n"; plain != want {
+		t.Fatalf("colored output changed JSON layout: got %q, want %q", plain, want)
+	}
+}
+
+func TestStreamSQLJSONResponsePreservesLargeNumbers(t *testing.T) {
+	var output bytes.Buffer
+	reader := bufio.NewReader(strings.NewReader(`[{"time_unix_nano":1790784848161039008}]`))
+
+	if err := streamSQLJSONResponse(&output, reader, "200 OK"); err != nil {
+		t.Fatalf("streamSQLJSONResponse failed: %v", err)
+	}
+	if !strings.Contains(output.String(), "1790784848161039008") {
+		t.Fatalf("large integer changed: %s", output.String())
+	}
+}
+
 func TestStreamSQLJSONResponsePrettyPrintsArray(t *testing.T) {
 	var output bytes.Buffer
 	reader := bufio.NewReader(strings.NewReader(`[{"b":2},{"a":1}]`))

cmd/utils.go

@@ -0,0 +1,25 @@
+package cmd
+
+import (
+	"os"
+
+	"github.com/muesli/termenv"
+	"golang.org/x/term"
+)
+
+// setupSQLTextColor enables ANSI colors only for terminal output. It returns
+// whether to color the result and a cleanup function that restores the console
+// mode if it was changed. Call the cleanup function after writing the result.
+func setupSQLTextColor() (bool, func()) {
+	noop := func() {}
+	_, noColor := os.LookupEnv("NO_COLOR")
+	if !term.IsTerminal(int(os.Stdout.Fd())) || noColor || os.Getenv("TERM") == "dumb" {
+		return false, noop
+	}
+
+	restore, err := termenv.EnableVirtualTerminalProcessing(termenv.NewOutput(os.Stdout, termenv.WithProfile(termenv.TrueColor)))
+	if err != nil {
+		return false, noop
+	}
+	return true, func() { _ = restore() }
+}

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/gofrs/flock v0.12.1
 	github.com/guptarohit/asciigraph v0.9.0
 	github.com/manifoldco/promptui v0.9.0
+	github.com/muesli/termenv v0.15.2
 	github.com/oklog/ulid/v2 v2.1.0
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/pkg/errors v0.9.1
@@ -173,7 +174,6 @@ require (
 require (
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/evertras/bubble-table v0.19.2
-	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.9
 	github.com/sahilm/fuzzy v0.1.1 // indirect
 )