From 6c3b4c6e67f2be3128565b837bb379febe10c111 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 00:34:34 +0100 Subject: [PATCH 1/4] Add CSV output format for SQL query results Add a --format csv flag to the query command for exporting results as CSV. Uses Go's encoding/csv for proper escaping and quoting. Column headers are included as the first row. Co-authored-by: Isaac --- experimental/aitools/cmd/query.go | 16 +++++++- experimental/aitools/cmd/render.go | 22 +++++++++++ experimental/aitools/cmd/render_test.go | 49 +++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go index dd5b97c761..e5b1071b03 100644 --- a/experimental/aitools/cmd/query.go +++ b/experimental/aitools/cmd/query.go @@ -44,6 +44,7 @@ type queryOutputMode int const ( queryOutputModeJSON queryOutputMode = iota + queryOutputModeCSV queryOutputModeStaticTable queryOutputModeInteractiveTable ) @@ -69,6 +70,7 @@ func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSup func newQueryCmd() *cobra.Command { var warehouseID string var filePath string + var format string cmd := &cobra.Command{ Use: "query [SQL | file.sql]", @@ -83,15 +85,21 @@ The command auto-detects an available warehouse unless --warehouse is set or the DATABRICKS_WAREHOUSE_ID environment variable is configured. Output is JSON in non-interactive contexts. In interactive terminals it renders -tables, and large results open an interactive table browser.`, +tables, and large results open an interactive table browser. Use --format csv +to export results as CSV.`, 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 databricks experimental aitools tools query report.sql + databricks experimental aitools tools query --format csv "SELECT * FROM samples.nyctaxi.trips LIMIT 5" echo "SELECT 1" | databricks experimental aitools tools query`, Args: cobra.MaximumNArgs(1), PreRunE: root.MustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { + if format != "" && format != "csv" { + return fmt.Errorf("unsupported format %q, supported values: csv", format) + } + ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) @@ -121,6 +129,11 @@ tables, and large results open an interactive table browser.`, return nil } + // CSV format bypasses the normal output mode selection. + if format == "csv" { + return renderCSV(cmd.OutOrStdout(), columns, rows) + } + // Output format depends on stdout capabilities. // Interactive table browsing also requires prompt-capable stdin. stdoutInteractive := cmdio.SupportsColor(ctx, cmd.OutOrStdout()) @@ -139,6 +152,7 @@ tables, and large results open an interactive table browser.`, cmd.Flags().StringVarP(&warehouseID, "warehouse", "w", "", "SQL warehouse ID to use for execution") cmd.Flags().StringVarP(&filePath, "file", "f", "", "Path to a SQL file to execute") + cmd.Flags().StringVar(&format, "format", "", "Output format: csv (default uses --output flag behavior)") return cmd } diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go index b7eadb401c..4b15045123 100644 --- a/experimental/aitools/cmd/render.go +++ b/experimental/aitools/cmd/render.go @@ -1,6 +1,7 @@ package mcp import ( + "encoding/csv" "encoding/json" "fmt" "io" @@ -51,6 +52,27 @@ func renderJSON(w io.Writer, columns []string, rows [][]string) error { return nil } +// renderCSV writes query results as CSV with column headers as the first row. +func renderCSV(w io.Writer, columns []string, rows [][]string) error { + cw := csv.NewWriter(w) + if err := cw.Write(columns); err != nil { + return fmt.Errorf("write CSV header: %w", err) + } + for _, row := range rows { + record := make([]string, len(columns)) + for i := range columns { + if i < len(row) { + record[i] = row[i] + } + } + if err := cw.Write(record); err != nil { + return fmt.Errorf("write CSV row: %w", err) + } + } + cw.Flush() + return cw.Error() +} + // renderStaticTable writes query results as a formatted text table. func renderStaticTable(w io.Writer, columns []string, rows [][]string) error { tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0) diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go index f559609749..6234b5b733 100644 --- a/experimental/aitools/cmd/render_test.go +++ b/experimental/aitools/cmd/render_test.go @@ -93,3 +93,52 @@ func TestRenderStaticTableEmpty(t *testing.T) { assert.Contains(t, output, "id") assert.Contains(t, output, "0 rows") } + +func TestRenderCSVBasic(t *testing.T) { + var buf bytes.Buffer + columns := []string{"id", "name", "city"} + rows := [][]string{ + {"1", "Alice", "New York"}, + {"2", "Bob", "London"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "id,name,city\n1,Alice,New York\n2,Bob,London\n", buf.String()) +} + +func TestRenderCSVSpecialCharacters(t *testing.T) { + var buf bytes.Buffer + columns := []string{"name", "description"} + rows := [][]string{ + {"Alice", "has a comma, here"}, + {"Bob", `has "quotes" here`}, + {"Carol", "has a\nnewline"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "name,description\nAlice,\"has a comma, here\"\nBob,\"has \"\"quotes\"\" here\"\nCarol,\"has a\nnewline\"\n", buf.String()) +} + +func TestRenderCSVEmptyResultSet(t *testing.T) { + var buf bytes.Buffer + columns := []string{"id", "name"} + var rows [][]string + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "id,name\n", buf.String()) +} + +func TestRenderCSVShortRows(t *testing.T) { + var buf bytes.Buffer + columns := []string{"a", "b", "c"} + rows := [][]string{ + {"1"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "a,b,c\n1,,\n", buf.String()) +} From 4142f53b63b287398a94e7762adf74b453764352 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 03:58:47 +0100 Subject: [PATCH 2/4] Fix review findings: flag conflict, no-results ordering, dead code, tests Co-authored-by: Isaac --- experimental/aitools/cmd/query.go | 20 ++++++++++------ experimental/aitools/cmd/query_test.go | 33 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go index e5b1071b03..70187b01f4 100644 --- a/experimental/aitools/cmd/query.go +++ b/experimental/aitools/cmd/query.go @@ -44,7 +44,6 @@ type queryOutputMode int const ( queryOutputModeJSON queryOutputMode = iota - queryOutputModeCSV queryOutputModeStaticTable queryOutputModeInteractiveTable ) @@ -100,6 +99,10 @@ to export results as CSV.`, return fmt.Errorf("unsupported format %q, supported values: csv", format) } + if format != "" && cmd.Flag("output").Changed { + return fmt.Errorf("cannot use --format and --output together; use --format csv or --output json/text") + } + ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) @@ -124,16 +127,19 @@ to export results as CSV.`, return err } - if len(columns) == 0 && len(rows) == 0 { - fmt.Fprintln(cmd.OutOrStdout(), "Query executed successfully (no results)") - return nil - } - // CSV format bypasses the normal output mode selection. if format == "csv" { + if len(columns) == 0 && len(rows) == 0 { + return nil + } return renderCSV(cmd.OutOrStdout(), columns, rows) } + 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()) @@ -152,7 +158,7 @@ to export results as CSV.`, cmd.Flags().StringVarP(&warehouseID, "warehouse", "w", "", "SQL warehouse ID to use for execution") cmd.Flags().StringVarP(&filePath, "file", "f", "", "Path to a SQL file to execute") - cmd.Flags().StringVar(&format, "format", "", "Output format: csv (default uses --output flag behavior)") + cmd.Flags().StringVar(&format, "format", "", "Output format (supported: csv). When omitted, uses --output flag behavior.") return cmd } diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go index 22b4fbaf13..b009ce000e 100644 --- a/experimental/aitools/cmd/query_test.go +++ b/experimental/aitools/cmd/query_test.go @@ -447,3 +447,36 @@ func TestResolveSQLMissingFileReturnsError(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "read SQL file") } + +func TestQueryCommandUnsupportedFormatReturnsError(t *testing.T) { + cmd := newQueryCmd() + cmd.PreRunE = nil + cmd.SetArgs([]string{"--format", "xml", "SELECT 1"}) + err := cmd.Execute() + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported format") +} + +func TestQueryCommandFormatAndOutputConflictReturnsError(t *testing.T) { + cmd := newQueryCmd() + cmd.PreRunE = nil + cmd.PersistentFlags().String("output", "text", "output type") + cmd.SetArgs([]string{"--format", "csv", "--output", "json", "SELECT 1"}) + err := cmd.Execute() + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot use --format and --output together") +} + +func TestRenderCSVOutput(t *testing.T) { + var buf strings.Builder + err := renderCSV(&buf, []string{"id", "name"}, [][]string{{"1", "alice"}, {"2", "bob"}}) + require.NoError(t, err) + assert.Equal(t, "id,name\n1,alice\n2,bob\n", buf.String()) +} + +func TestRenderCSVHeadersOnlyWhenNoRows(t *testing.T) { + var buf strings.Builder + err := renderCSV(&buf, []string{"id", "name"}, nil) + require.NoError(t, err) + assert.Equal(t, "id,name\n", buf.String()) +} From 49e07265ec7fd5d70582f3c562ce8df295656b87 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 07:23:37 +0100 Subject: [PATCH 3/4] Use CRLF line endings for RFC 4180 compliance --- experimental/aitools/cmd/query_test.go | 4 ++-- experimental/aitools/cmd/render.go | 1 + experimental/aitools/cmd/render_test.go | 8 ++++---- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go index b009ce000e..2fe18473ba 100644 --- a/experimental/aitools/cmd/query_test.go +++ b/experimental/aitools/cmd/query_test.go @@ -471,12 +471,12 @@ func TestRenderCSVOutput(t *testing.T) { var buf strings.Builder err := renderCSV(&buf, []string{"id", "name"}, [][]string{{"1", "alice"}, {"2", "bob"}}) require.NoError(t, err) - assert.Equal(t, "id,name\n1,alice\n2,bob\n", buf.String()) + assert.Equal(t, "id,name\r\n1,alice\r\n2,bob\r\n", buf.String()) } func TestRenderCSVHeadersOnlyWhenNoRows(t *testing.T) { var buf strings.Builder err := renderCSV(&buf, []string{"id", "name"}, nil) require.NoError(t, err) - assert.Equal(t, "id,name\n", buf.String()) + assert.Equal(t, "id,name\r\n", buf.String()) } diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go index 4b15045123..09cffb3049 100644 --- a/experimental/aitools/cmd/render.go +++ b/experimental/aitools/cmd/render.go @@ -55,6 +55,7 @@ func renderJSON(w io.Writer, columns []string, rows [][]string) error { // renderCSV writes query results as CSV with column headers as the first row. func renderCSV(w io.Writer, columns []string, rows [][]string) error { cw := csv.NewWriter(w) + cw.UseCRLF = true if err := cw.Write(columns); err != nil { return fmt.Errorf("write CSV header: %w", err) } diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go index 6234b5b733..25b8b2386c 100644 --- a/experimental/aitools/cmd/render_test.go +++ b/experimental/aitools/cmd/render_test.go @@ -104,7 +104,7 @@ func TestRenderCSVBasic(t *testing.T) { err := renderCSV(&buf, columns, rows) require.NoError(t, err) - assert.Equal(t, "id,name,city\n1,Alice,New York\n2,Bob,London\n", buf.String()) + assert.Equal(t, "id,name,city\r\n1,Alice,New York\r\n2,Bob,London\r\n", buf.String()) } func TestRenderCSVSpecialCharacters(t *testing.T) { @@ -118,7 +118,7 @@ func TestRenderCSVSpecialCharacters(t *testing.T) { err := renderCSV(&buf, columns, rows) require.NoError(t, err) - assert.Equal(t, "name,description\nAlice,\"has a comma, here\"\nBob,\"has \"\"quotes\"\" here\"\nCarol,\"has a\nnewline\"\n", buf.String()) + assert.Equal(t, "name,description\r\nAlice,\"has a comma, here\"\r\nBob,\"has \"\"quotes\"\" here\"\r\nCarol,\"has a\r\nnewline\"\r\n", buf.String()) } func TestRenderCSVEmptyResultSet(t *testing.T) { @@ -128,7 +128,7 @@ func TestRenderCSVEmptyResultSet(t *testing.T) { err := renderCSV(&buf, columns, rows) require.NoError(t, err) - assert.Equal(t, "id,name\n", buf.String()) + assert.Equal(t, "id,name\r\n", buf.String()) } func TestRenderCSVShortRows(t *testing.T) { @@ -140,5 +140,5 @@ func TestRenderCSVShortRows(t *testing.T) { err := renderCSV(&buf, columns, rows) require.NoError(t, err) - assert.Equal(t, "a,b,c\n1,,\n", buf.String()) + assert.Equal(t, "a,b,c\r\n1,,\r\n", buf.String()) } From 199578928528e8ea6c5ca6ce7ac453227b43c13a Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 07:58:48 +0100 Subject: [PATCH 4/4] Fix perfsprint lint: fmt.Errorf to errors.New --- experimental/aitools/cmd/query.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go index 70187b01f4..2c7cb51c21 100644 --- a/experimental/aitools/cmd/query.go +++ b/experimental/aitools/cmd/query.go @@ -100,7 +100,7 @@ to export results as CSV.`, } if format != "" && cmd.Flag("output").Changed { - return fmt.Errorf("cannot use --format and --output together; use --format csv or --output json/text") + return errors.New("cannot use --format and --output together; use --format csv or --output json/text") } ctx := cmd.Context()