Add output formats, interactive table browser, and multi-chunk results (#4646)

## Why

The `aitools tools query` command outputs results as JSON to stderr via
`cmdio.LogString`. This means output can't be piped, there's no tabular
display for interactive use, and large multi-chunk results are truncated
to the first chunk.

## Changes

Before: always JSON to stderr, single chunk only.
Now: format-aware output routing with interactive table browser.

**Output routing:**
- Interactive terminal (stdout is TTY): table display
  - Small results (<=30 rows): static formatted table, prints and exits
- Large results (>30 rows): interactive browser with scrolling and
search
- Non-interactive (piped stdout) or `--output json`: JSON array of
objects (backwards compatible)
- All output goes to stdout, not stderr

**Interactive table browser (`libs/tableview/`):**
Reusable shared component using `bubbles/viewport`. Features:
- Horizontal + vertical scrolling (arrow keys, pgup/pgdn)
- Cursor row highlighting (purple background)
- `/` search with yellow match highlighting, `n`/`N` to navigate matches
- `g`/`G` for top/bottom, `q` to quit
- Footer with row count, scroll position, keybinding hints

**Multi-chunk fetching:**
Queries returning more data than fits in a single response chunk now
fetch all chunks via `GetStatementResultChunkN` and combine them.



https://github.com/user-attachments/assets/60b299aa-01ef-4b75-89e0-1f15dac05e42



https://github.com/user-attachments/assets/86747be1-6c94-4399-b0f3-dd2fc25f3df8



## Test plan

- [x] Unit tests for `libs/tableview` (table rendering, search,
highlighting)
- [x] Unit tests for renderers (JSON format, static table,
extractColumns)
- [x] Unit tests for `fetchAllRows` (single chunk, multi-chunk, nil
result)
- [x] All existing query tests pass (polling, cancellation, input
resolution)
- [x] `make lintfull` clean
- [x] `make checks` clean
- [x] Manual: interactive table with scrolling and search
- [x] Manual: `| cat` produces JSON
- [x] Manual: `--output json` forces JSON

Author: simonfaltum

experimental/aitools/cmd/query.go

@@ -2,7 +2,6 @@ package mcp
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -17,6 +16,7 @@ import (
 	"github.com/databricks/cli/experimental/aitools/lib/session"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -34,8 +34,38 @@ const (
 
 	// cancelTimeout is how long to wait for server-side cancellation.
 	cancelTimeout = 10 * time.Second
+
+	// staticTableThreshold is the maximum number of rows rendered as a static table.
+	// Beyond this, an interactive scrollable table is used.
+	staticTableThreshold = 30
+)
+
+type queryOutputMode int
+
+const (
+	queryOutputModeJSON queryOutputMode = iota
+	queryOutputModeStaticTable
+	queryOutputModeInteractiveTable
 )
 
+func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSupported bool, rowCount int) queryOutputMode {
+	if outputType == flags.OutputJSON {
+		return queryOutputModeJSON
+	}
+	if !stdoutInteractive {
+		return queryOutputModeJSON
+	}
+	// Interactive table browsing requires keyboard input from stdin.
+	// If prompts are not supported, prefer static table output instead.
+	if !promptSupported {
+		return queryOutputModeStaticTable
+	}
+	if rowCount <= staticTableThreshold {
+		return queryOutputModeStaticTable
+	}
+	return queryOutputModeInteractiveTable
+}
+
 func newQueryCmd() *cobra.Command {
 	var warehouseID string
 	var filePath string
@@ -52,7 +82,8 @@ exists, it is read as a SQL file automatically.
 The command auto-detects an available warehouse unless --warehouse is set
 or the DATABRICKS_WAREHOUSE_ID environment variable is configured.
 
-Output includes the query results as JSON and row count.`,
+Output is JSON in non-interactive contexts. In interactive terminals it renders
+tables, and large results open an interactive table browser.`,
 		Example: `  databricks experimental aitools tools query "SELECT * FROM samples.nyctaxi.trips LIMIT 5"
   databricks experimental aitools tools query --warehouse abc123 "SELECT 1"
   databricks experimental aitools tools query --file report.sql
@@ -79,13 +110,30 @@ Output includes the query results as JSON and row count.`,
 				return err
 			}
 
-			output, err := formatQueryResult(resp)
+			columns := extractColumns(resp.Manifest)
+			rows, err := fetchAllRows(ctx, w.StatementExecution, resp)
 			if err != nil {
 				return err
 			}
 
-			cmdio.LogString(ctx, output)
-			return nil
+			if len(columns) == 0 && len(rows) == 0 {
+				fmt.Fprintln(cmd.OutOrStdout(), "Query executed successfully (no results)")
+				return nil
+			}
+
+			// Output format depends on stdout capabilities.
+			// Interactive table browsing also requires prompt-capable stdin.
+			stdoutInteractive := cmdio.SupportsColor(ctx, cmd.OutOrStdout())
+			promptSupported := cmdio.IsPromptSupported(ctx)
+
+			switch selectQueryOutputMode(root.OutputType(cmd), stdoutInteractive, promptSupported, len(rows)) {
+			case queryOutputModeJSON:
+				return renderJSON(cmd.OutOrStdout(), columns, rows)
+			case queryOutputModeStaticTable:
+				return renderStaticTable(cmd.OutOrStdout(), columns, rows)
+			default:
+				return renderInteractiveTable(cmd.OutOrStdout(), columns, rows)
+			}
 		},
 	}
 
@@ -274,6 +322,31 @@ func executeAndPoll(ctx context.Context, api sql.StatementExecutionInterface, wa
 	}
 }
 
+// fetchAllRows collects all result rows, fetching additional chunks if needed.
+func fetchAllRows(ctx context.Context, api sql.StatementExecutionInterface, resp *sql.StatementResponse) ([][]string, error) {
+	if resp.Result == nil {
+		return nil, nil
+	}
+
+	rows := append([][]string{}, resp.Result.DataArray...)
+
+	totalChunks := 0
+	if resp.Manifest != nil {
+		totalChunks = resp.Manifest.TotalChunkCount
+	}
+
+	for chunk := 1; chunk < totalChunks; chunk++ {
+		log.Debugf(ctx, "Fetching result chunk %d/%d for statement %s", chunk+1, totalChunks, resp.StatementId)
+		chunkResp, err := api.GetStatementResultChunkNByStatementIdAndChunkIndex(ctx, resp.StatementId, chunk)
+		if err != nil {
+			return nil, fmt.Errorf("fetch result chunk %d: %w", chunk, err)
+		}
+		rows = append(rows, chunkResp.DataArray...)
+	}
+
+	return rows, nil
+}
+
 // isTerminalState returns true if the statement has reached a final state.
 func isTerminalState(status *sql.StatementStatus) bool {
 	if status == nil {
@@ -334,43 +407,3 @@ func cleanSQL(s string) string {
 
 	return strings.Join(lines, "\n")
 }
-
-func formatQueryResult(resp *sql.StatementResponse) (string, error) {
-	var sb strings.Builder
-
-	if resp.Manifest == nil || resp.Result == nil {
-		sb.WriteString("Query executed successfully (no results)\n")
-		return sb.String(), nil
-	}
-
-	var columns []string
-	if resp.Manifest.Schema != nil {
-		for _, col := range resp.Manifest.Schema.Columns {
-			columns = append(columns, col.Name)
-		}
-	}
-
-	var rows []map[string]any
-	if resp.Result.DataArray != nil {
-		for _, row := range resp.Result.DataArray {
-			rowMap := make(map[string]any)
-			for i, val := range row {
-				if i < len(columns) {
-					rowMap[columns[i]] = val
-				}
-			}
-			rows = append(rows, rowMap)
-		}
-	}
-
-	output, err := json.MarshalIndent(rows, "", "  ")
-	if err != nil {
-		return "", fmt.Errorf("marshal results: %w", err)
-	}
-
-	sb.Write(output)
-	sb.WriteString("\n\n")
-	sb.WriteString(fmt.Sprintf("Row count: %d\n", len(rows)))
-
-	return sb.String(), nil
-}

experimental/aitools/cmd/query_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	mocksql "github.com/databricks/databricks-sdk-go/experimental/mocks/service/sql"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -159,32 +160,109 @@ func TestResolveWarehouseIDWithFlag(t *testing.T) {
 	assert.Equal(t, "explicit-id", id)
 }
 
-func TestFormatQueryResultNoResults(t *testing.T) {
+func TestSelectQueryOutputMode(t *testing.T) {
+	tests := []struct {
+		name              string
+		outputType        flags.Output
+		stdoutInteractive bool
+		promptSupported   bool
+		rowCount          int
+		want              queryOutputMode
+	}{
+		{
+			name:              "json flag always returns json",
+			outputType:        flags.OutputJSON,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          999,
+			want:              queryOutputModeJSON,
+		},
+		{
+			name:              "non interactive stdout returns json",
+			outputType:        flags.OutputText,
+			stdoutInteractive: false,
+			promptSupported:   true,
+			rowCount:          5,
+			want:              queryOutputModeJSON,
+		},
+		{
+			name:              "missing stdin interactivity falls back to static table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   false,
+			rowCount:          staticTableThreshold + 10,
+			want:              queryOutputModeStaticTable,
+		},
+		{
+			name:              "small results use static table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          staticTableThreshold,
+			want:              queryOutputModeStaticTable,
+		},
+		{
+			name:              "large results use interactive table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          staticTableThreshold + 1,
+			want:              queryOutputModeInteractiveTable,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := selectQueryOutputMode(tc.outputType, tc.stdoutInteractive, tc.promptSupported, tc.rowCount)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestFetchAllRowsSingleChunk(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
 	resp := &sql.StatementResponse{
-		Status: &sql.StatementStatus{State: sql.StatementStateSucceeded},
+		StatementId: "stmt-1",
+		Manifest:    &sql.ResultManifest{TotalChunkCount: 1},
+		Result:      &sql.ResultData{DataArray: [][]string{{"1", "alice"}, {"2", "bob"}}},
 	}
-	output, err := formatQueryResult(resp)
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
 	require.NoError(t, err)
-	assert.Contains(t, output, "no results")
+	assert.Equal(t, [][]string{{"1", "alice"}, {"2", "bob"}}, rows)
 }
 
-func TestFormatQueryResultWithData(t *testing.T) {
+func TestFetchAllRowsMultiChunk(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
 	resp := &sql.StatementResponse{
-		Status: &sql.StatementStatus{State: sql.StatementStateSucceeded},
-		Manifest: &sql.ResultManifest{
-			Schema: &sql.ResultSchema{
-				Columns: []sql.ColumnInfo{{Name: "id"}, {Name: "name"}},
-			},
-		},
-		Result: &sql.ResultData{
-			DataArray: [][]string{{"1", "alice"}, {"2", "bob"}},
-		},
+		StatementId: "stmt-1",
+		Manifest:    &sql.ResultManifest{TotalChunkCount: 3},
+		Result:      &sql.ResultData{DataArray: [][]string{{"1", "a"}}},
 	}
-	output, err := formatQueryResult(resp)
+
+	mockAPI.EXPECT().GetStatementResultChunkNByStatementIdAndChunkIndex(mock.Anything, "stmt-1", 1).
+		Return(&sql.ResultData{DataArray: [][]string{{"2", "b"}}}, nil).Once()
+	mockAPI.EXPECT().GetStatementResultChunkNByStatementIdAndChunkIndex(mock.Anything, "stmt-1", 2).
+		Return(&sql.ResultData{DataArray: [][]string{{"3", "c"}}}, nil).Once()
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
+	require.NoError(t, err)
+	assert.Equal(t, [][]string{{"1", "a"}, {"2", "b"}, {"3", "c"}}, rows)
+}
+
+func TestFetchAllRowsNilResult(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
+	resp := &sql.StatementResponse{StatementId: "stmt-1"}
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
 	require.NoError(t, err)
-	assert.Contains(t, output, "alice")
-	assert.Contains(t, output, "bob")
-	assert.Contains(t, output, "Row count: 2")
+	assert.Nil(t, rows)
 }
 
 func TestIsTerminalState(t *testing.T) {